From bf71e37d050201aebebc088654f3eac33cb26ec2 Mon Sep 17 00:00:00 2001
From: DongXM <2139383712@qq.com>
Date: Sun, 5 Jul 2026 12:19:29 +0800
Subject: [PATCH 1/3] fix(persistence): register asyncpg JSON/JSONB codecs on
 PG engine

asyncpg returns JSON / JSONB columns as raw strings by default, which
breaks every LangBot code path that expects a dict (mcp server loading,
bot adapter config, pipeline config, etc.) when running on PostgreSQL.

The asyncpg dialect's built-in on_connect hook should install the
codecs, but it does not fire reliably on the SQLAlchemy 2.0 async
engine path. Monkey-patch AsyncAdapt_asyncpg_dbapi.connect so that
each new connection registers json / jsonb codecs eagerly via the
wrapper's run_async hook.
---
 .../pkg/persistence/databases/postgresql.py   | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/src/langbot/pkg/persistence/databases/postgresql.py b/src/langbot/pkg/persistence/databases/postgresql.py
index f63d8f61f4..be2bca055e 100644
--- a/src/langbot/pkg/persistence/databases/postgresql.py
+++ b/src/langbot/pkg/persistence/databases/postgresql.py
@@ -1,10 +1,66 @@
 from __future__ import annotations
 
+import json
+
 import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
 
 from .. import database
 
 
+async def _register_json_codecs(pg_conn) -> None:
+    """Register asyncpg type codecs for JSON / JSONB so values come back as
+    Python ``dict`` / ``list`` instead of raw ``str``.
+    """
+    await pg_conn.set_type_codec(
+        'json',
+        encoder=json.dumps,
+        decoder=json.loads,
+        schema='pg_catalog',
+        format='text',
+    )
+    await pg_conn.set_type_codec(
+        'jsonb',
+        encoder=json.dumps,
+        decoder=json.loads,
+        schema='pg_catalog',
+        format='text',
+    )
+
+
+def _patch_asyncpg_dialect() -> None:
+    """Wrap ``AsyncAdapt_asyncpg_dbapi.connect`` so that every new
+    asyncpg-backed SQLAlchemy connection gets our JSON codecs installed
+    eagerly, via the wrapper's ``run_async`` hook.
+    """
+    from sqlalchemy.dialects.postgresql.asyncpg import AsyncAdapt_asyncpg_dbapi
+
+    _orig = AsyncAdapt_asyncpg_dbapi.connect
+
+    def _patched(self, *args, **kwargs):
+        wrapper = _orig(self, *args, **kwargs)
+
+        # `AsyncAdapt_asyncpg_connection.run_async(fn)` invokes `fn(conn)`
+        # where `conn` is the real asyncpg.Connection. The hook must accept
+        # that argument; some SQLAlchemy versions forward it positionally,
+        # others don't, so we tolerate both.
+        async def _init(conn=None):
+            pg = conn
+            if pg is None:
+                adapt = wrapper.connection
+                pg = adapt._connection if hasattr(adapt, '_connection') else adapt
+            await _register_json_codecs(pg)
+
+        wrapper.run_async(_init)
+        return wrapper
+
+    AsyncAdapt_asyncpg_dbapi.connect = _patched
+
+
+# Apply the patch at module import time so subsequent engine creation
+# inherits the codec registration.
+_patch_asyncpg_dialect()
+
+
 @database.manager_class('postgresql')
 class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
     """PostgreSQL database manager"""

From 8de58978e6df3a46fe188a4df15f21a64f9562e3 Mon Sep 17 00:00:00 2001
From: DongXM <2139383712@qq.com>
Date: Thu, 16 Jul 2026 16:25:47 +0800
Subject: [PATCH 2/3] test(postgresql): add unit tests for asyncpg JSON/JSONB
 codec registration

---
 .../persistence/test_postgresql_codec.py      | 280 ++++++++++++++++++
 1 file changed, 280 insertions(+)
 create mode 100644 tests/unit_tests/persistence/test_postgresql_codec.py

diff --git a/tests/unit_tests/persistence/test_postgresql_codec.py b/tests/unit_tests/persistence/test_postgresql_codec.py
new file mode 100644
index 0000000000..ff46214b88
--- /dev/null
+++ b/tests/unit_tests/persistence/test_postgresql_codec.py
@@ -0,0 +1,280 @@
+"""Unit tests for asyncpg JSON/JSONB codec registration in PostgreSQL manager.
+
+Tests cover:
+- _register_json_codecs registers both json and jsonb codecs
+- _patch_asyncpg_dialect wraps the original connect method
+- The patched connect invokes run_async to install codecs on each new connection
+
+Note: Uses import isolation to break circular import chains (same pattern as
+test_database_decorator.py).
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from contextlib import contextmanager
+from typing import Generator
+from unittest.mock import AsyncMock, MagicMock
+
+from sqlalchemy.dialects.postgresql.asyncpg import AsyncAdapt_asyncpg_dbapi
+
+
+@contextmanager
+def isolated_database_import() -> Generator[None, None, None]:
+    """Context manager to isolate circular imports for database testing."""
+    mock_app = MagicMock()
+    mock_importutil = MagicMock()
+    mock_importutil.import_modules_in_pkg = lambda pkg: None
+    mock_importutil.import_modules_in_pkgs = lambda pkgs: None
+    mock_mgr = MagicMock()
+
+    mocks = {
+        'langbot.pkg.core.app': mock_app,
+        'langbot.pkg.utils.importutil': mock_importutil,
+        'langbot.pkg.persistence.mgr': mock_mgr,
+    }
+
+    saved: dict[str, object] = {}
+    for name in mocks:
+        if name in sys.modules:
+            saved[name] = sys.modules[name]
+
+    database_name = 'langbot.pkg.persistence.database'
+    if database_name in sys.modules:
+        saved[database_name] = sys.modules[database_name]
+
+    for sub in ['sqlite', 'postgresql']:
+        full_name = f'langbot.pkg.persistence.databases.{sub}'
+        if full_name in sys.modules:
+            saved[full_name] = sys.modules[full_name]
+
+    try:
+        for name, module in mocks.items():
+            sys.modules[name] = module
+
+        sys.modules.pop(database_name, None)
+        for sub in ['sqlite', 'postgresql']:
+            sys.modules.pop(f'langbot.pkg.persistence.databases.{sub}', None)
+
+        yield
+    finally:
+        for name in mocks:
+            if name in saved:
+                sys.modules[name] = saved[name]
+            else:
+                sys.modules.pop(name, None)
+
+        if database_name in saved:
+            sys.modules[database_name] = saved[database_name]
+        else:
+            sys.modules.pop(database_name, None)
+
+        for sub in ['sqlite', 'postgresql']:
+            full_name = f'langbot.pkg.persistence.databases.{sub}'
+            if full_name in saved:
+                sys.modules[full_name] = saved[full_name]
+            else:
+                sys.modules.pop(full_name, None)
+
+
+def get_postgresql_module():
+    """Get the postgresql database module with import isolation.
+
+    Saves and restores the original ``AsyncAdapt_asyncpg_dbapi.connect``
+    because the module calls ``_patch_asyncpg_dialect()`` at import time.
+    """
+    orig_connect = AsyncAdapt_asyncpg_dbapi.connect
+    try:
+        with isolated_database_import():
+            from langbot.pkg.persistence.databases import postgresql
+
+            return postgresql
+    finally:
+        AsyncAdapt_asyncpg_dbapi.connect = orig_connect
+
+
+class TestRegisterJsonCodecs:
+    """Tests for _register_json_codecs."""
+
+    async def test_registers_json_and_jsonb_codecs(self):
+        """Test that both json and jsonb codecs are registered on the connection."""
+        pg_module = get_postgresql_module()
+
+        mock_conn = MagicMock()
+        mock_conn.set_type_codec = AsyncMock()
+
+        await pg_module._register_json_codecs(mock_conn)
+
+        assert mock_conn.set_type_codec.call_count == 2
+
+        first_call = mock_conn.set_type_codec.call_args_list[0]
+        assert first_call.args[0] == 'json'
+        assert first_call.kwargs['encoder'] == json.dumps
+        assert first_call.kwargs['decoder'] == json.loads
+        assert first_call.kwargs['schema'] == 'pg_catalog'
+        assert first_call.kwargs['format'] == 'text'
+
+        second_call = mock_conn.set_type_codec.call_args_list[1]
+        assert second_call.args[0] == 'jsonb'
+        assert second_call.kwargs['encoder'] == json.dumps
+        assert second_call.kwargs['decoder'] == json.loads
+        assert second_call.kwargs['schema'] == 'pg_catalog'
+        assert second_call.kwargs['format'] == 'text'
+
+    async def test_codec_decodes_json_string_to_dict(self):
+        """Test that the registered decoder correctly parses JSON strings."""
+        pg_module = get_postgresql_module()
+
+        mock_conn = MagicMock()
+        mock_conn.set_type_codec = AsyncMock()
+
+        await pg_module._register_json_codecs(mock_conn)
+
+        json_call = mock_conn.set_type_codec.call_args_list[0]
+        decoder = json_call.kwargs['decoder']
+        assert decoder('{"key": "value"}') == {'key': 'value'}
+
+        jsonb_call = mock_conn.set_type_codec.call_args_list[1]
+        decoder_b = jsonb_call.kwargs['decoder']
+        assert decoder_b('[1, 2, 3]') == [1, 2, 3]
+
+    async def test_codec_encodes_dict_to_json_string(self):
+        """Test that the registered encoder correctly serializes Python objects."""
+        pg_module = get_postgresql_module()
+
+        mock_conn = MagicMock()
+        mock_conn.set_type_codec = AsyncMock()
+
+        await pg_module._register_json_codecs(mock_conn)
+
+        json_call = mock_conn.set_type_codec.call_args_list[0]
+        encoder = json_call.kwargs['encoder']
+        assert encoder({'key': 'value'}) == '{"key": "value"}'
+
+
+class TestPatchAsyncpgDialect:
+    """Tests for _patch_asyncpg_dialect."""
+
+    def test_patch_replaces_connect_method(self):
+        """Test that _patch_asyncpg_dialect replaces the connect method."""
+        pg_module = get_postgresql_module()
+        orig_connect = AsyncAdapt_asyncpg_dbapi.connect
+
+        try:
+            pg_module._patch_asyncpg_dialect()
+            assert AsyncAdapt_asyncpg_dbapi.connect is not orig_connect
+        finally:
+            AsyncAdapt_asyncpg_dbapi.connect = orig_connect
+
+    def test_patched_connect_calls_original_and_run_async(self):
+        """Test that the patched connect calls original connect and run_async."""
+        pg_module = get_postgresql_module()
+
+        mock_wrapper = MagicMock()
+        mock_wrapper.run_async = MagicMock()
+        fake_orig_connect = MagicMock(return_value=mock_wrapper)
+
+        orig_connect = AsyncAdapt_asyncpg_dbapi.connect
+        AsyncAdapt_asyncpg_dbapi.connect = fake_orig_connect
+
+        try:
+            pg_module._patch_asyncpg_dialect()
+            patched = AsyncAdapt_asyncpg_dbapi.connect
+
+            mock_self = MagicMock()
+            result = patched(mock_self, 'arg1', kwarg='val')
+
+            fake_orig_connect.assert_called_once_with(mock_self, 'arg1', kwarg='val')
+            mock_wrapper.run_async.assert_called_once()
+            assert result is mock_wrapper
+        finally:
+            AsyncAdapt_asyncpg_dbapi.connect = orig_connect
+
+    async def test_patched_connect_init_registers_codecs(self):
+        """Test that the run_async callback registers JSON codecs on the connection."""
+        pg_module = get_postgresql_module()
+
+        mock_wrapper = MagicMock()
+        captured_init: list = []
+        mock_wrapper.run_async = lambda fn: captured_init.append(fn)
+
+        fake_orig_connect = MagicMock(return_value=mock_wrapper)
+        orig_connect = AsyncAdapt_asyncpg_dbapi.connect
+        AsyncAdapt_asyncpg_dbapi.connect = fake_orig_connect
+
+        try:
+            pg_module._patch_asyncpg_dialect()
+            patched = AsyncAdapt_asyncpg_dbapi.connect
+            patched(MagicMock())
+
+            assert len(captured_init) == 1
+
+            mock_pg_conn = MagicMock()
+            mock_pg_conn.set_type_codec = AsyncMock()
+            await captured_init[0](mock_pg_conn)
+
+            assert mock_pg_conn.set_type_codec.call_count == 2
+        finally:
+            AsyncAdapt_asyncpg_dbapi.connect = orig_connect
+
+    async def test_init_handles_conn_none_fallback(self):
+        """Test that _init falls back to wrapper.connection when conn is None."""
+        pg_module = get_postgresql_module()
+
+        mock_pg_conn = MagicMock()
+        mock_pg_conn.set_type_codec = AsyncMock()
+
+        mock_wrapper = MagicMock()
+        captured_init: list = []
+        mock_wrapper.run_async = lambda fn: captured_init.append(fn)
+
+        mock_adapt = MagicMock()
+        mock_adapt._connection = mock_pg_conn
+        mock_wrapper.connection = mock_adapt
+
+        fake_orig_connect = MagicMock(return_value=mock_wrapper)
+        orig_connect = AsyncAdapt_asyncpg_dbapi.connect
+        AsyncAdapt_asyncpg_dbapi.connect = fake_orig_connect
+
+        try:
+            pg_module._patch_asyncpg_dialect()
+            patched = AsyncAdapt_asyncpg_dbapi.connect
+            patched(MagicMock())
+
+            assert len(captured_init) == 1
+            await captured_init[0](None)
+
+            assert mock_pg_conn.set_type_codec.call_count == 2
+        finally:
+            AsyncAdapt_asyncpg_dbapi.connect = orig_connect
+
+    def test_double_patch_still_works(self):
+        """Test that calling _patch_asyncpg_dialect twice still yields a working patch.
+
+        Double-patching stacks wrappers, so run_async fires once per layer.
+        This verifies the stacked patch does not break.
+        """
+        pg_module = get_postgresql_module()
+
+        mock_wrapper = MagicMock()
+        mock_wrapper.run_async = MagicMock()
+
+        fake_orig_connect = MagicMock(return_value=mock_wrapper)
+        orig_connect = AsyncAdapt_asyncpg_dbapi.connect
+        AsyncAdapt_asyncpg_dbapi.connect = fake_orig_connect
+
+        try:
+            pg_module._patch_asyncpg_dialect()
+            first_patched = AsyncAdapt_asyncpg_dbapi.connect
+
+            pg_module._patch_asyncpg_dialect()
+            second_patched = AsyncAdapt_asyncpg_dbapi.connect
+
+            assert first_patched is not second_patched
+
+            mock_wrapper.run_async.reset_mock()
+            second_patched(MagicMock())
+            assert mock_wrapper.run_async.called
+        finally:
+            AsyncAdapt_asyncpg_dbapi.connect = orig_connect

From c274f5cf4fd770427fb194976856f2a88640eb12 Mon Sep 17 00:00:00 2001
From: DongXM <2139383712@qq.com>
Date: Thu, 16 Jul 2026 16:29:39 +0800
Subject: [PATCH 3/3] fix(monitoring): use scalars().first() to correctly fetch
 ORM object in tool context

The previous result.first() + row[0] pattern returned a raw string (primary key value) instead of the ORM entity on PostgreSQL/asyncpg. Use result.scalars().first() per SQLAlchemy 2.0 convention so the MonitoringMessage object is correctly retrieved.
---
 .../pkg/api/http/service/monitoring.py        |  17 +-
 .../api/test_monitoring_tool_context.py       | 202 ++++++++++++++++++
 2 files changed, 210 insertions(+), 9 deletions(-)
 create mode 100644 tests/unit_tests/api/test_monitoring_tool_context.py

diff --git a/src/langbot/pkg/api/http/service/monitoring.py b/src/langbot/pkg/api/http/service/monitoring.py
index 46a352ad1c..74b1287ec3 100644
--- a/src/langbot/pkg/api/http/service/monitoring.py
+++ b/src/langbot/pkg/api/http/service/monitoring.py
@@ -160,16 +160,16 @@ async def _get_message_for_tool_context(
         self,
         message_id: str | None = None,
         session_id: str | None = None,
-    ):
+    ) -> persistence_monitoring.MonitoringMessage | None:
         if message_id:
             result = await self.ap.persistence_mgr.execute_async(
                 sqlalchemy.select(persistence_monitoring.MonitoringMessage).where(
                     persistence_monitoring.MonitoringMessage.id == message_id
                 )
             )
-            row = result.first()
-            if row:
-                return row[0]
+            obj = result.scalars().first()
+            if obj:
+                return obj
 
         if not session_id:
             return None
@@ -186,9 +186,9 @@ async def _get_message_for_tool_context(
             .limit(1)
         )
         result = await self.ap.persistence_mgr.execute_async(user_query)
-        row = result.first()
-        if row:
-            return row[0]
+        obj = result.scalars().first()
+        if obj:
+            return obj
 
         any_query = (
             sqlalchemy.select(persistence_monitoring.MonitoringMessage)
@@ -197,8 +197,7 @@ async def _get_message_for_tool_context(
             .limit(1)
         )
         result = await self.ap.persistence_mgr.execute_async(any_query)
-        row = result.first()
-        return row[0] if row else None
+        return result.scalars().first()
 
     # ========== Recording Methods ==========
 
diff --git a/tests/unit_tests/api/test_monitoring_tool_context.py b/tests/unit_tests/api/test_monitoring_tool_context.py
new file mode 100644
index 0000000000..f9fc3de41a
--- /dev/null
+++ b/tests/unit_tests/api/test_monitoring_tool_context.py
@@ -0,0 +1,202 @@
+"""Unit tests for MonitoringService._get_message_for_tool_context.
+
+Tests cover:
+- Fetching a MonitoringMessage by message_id (first query path)
+- Falling back to session_id + role='user' query (second path)
+- Falling back to session_id + any role query (third path)
+- Returning None when no message is found (safe degradation)
+- Returning None when neither message_id nor session_id is provided
+
+The fix replaces the old ``result.first()`` + ``row[0]`` pattern with
+``result.scalars().first()`` (SQLAlchemy 2.0 convention) so the ORM entity
+is correctly retrieved on PostgreSQL/asyncpg.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from langbot.pkg.api.http.service.monitoring import MonitoringService
+from langbot.pkg.entity.persistence.monitoring import MonitoringMessage
+
+pytestmark = pytest.mark.asyncio
+
+
+def _make_mock_message(
+    msg_id: str = 'msg-1',
+    bot_id: str = 'bot-1',
+    bot_name: str = 'TestBot',
+    pipeline_id: str = 'pipe-1',
+    pipeline_name: str = 'TestPipeline',
+    session_id: str = 'sess-1',
+    role: str = 'user',
+) -> MonitoringMessage:
+    """Create a MonitoringMessage instance for testing."""
+    return MonitoringMessage(
+        id=msg_id,
+        bot_id=bot_id,
+        bot_name=bot_name,
+        pipeline_id=pipeline_id,
+        pipeline_name=pipeline_name,
+        message_content='hello',
+        session_id=session_id,
+        status='success',
+        level='info',
+        role=role,
+    )
+
+
+def _make_scalars_result(obj=None):
+    """Create a mock query result whose scalars().first() returns *obj*."""
+    result = MagicMock()
+    result.scalars.return_value.first.return_value = obj
+    return result
+
+
+def _make_ap(side_effect=None):
+    """Create a mock Application with persistence_mgr.execute_async."""
+    ap = SimpleNamespace()
+    ap.persistence_mgr = SimpleNamespace()
+    if side_effect is not None:
+        ap.persistence_mgr.execute_async = AsyncMock(side_effect=side_effect)
+    else:
+        ap.persistence_mgr.execute_async = AsyncMock()
+    return ap
+
+
+class TestGetMessageForToolContextByMessageId:
+    """Tests for the message_id lookup path."""
+
+    async def test_returns_message_when_found_by_message_id(self):
+        """Returns the MonitoringMessage when found by message_id."""
+        msg = _make_mock_message(msg_id='msg-123')
+        result = _make_scalars_result(msg)
+        ap = _make_ap(side_effect=[result])
+
+        service = MonitoringService(ap)
+        found = await service._get_message_for_tool_context(message_id='msg-123')
+
+        assert found is msg
+        assert found.id == 'msg-123'
+        ap.persistence_mgr.execute_async.assert_awaited_once()
+
+    async def test_falls_through_when_message_id_not_found(self):
+        """Falls through to session_id lookup when message_id yields nothing."""
+        empty_result = _make_scalars_result(None)
+        msg = _make_mock_message(session_id='sess-1', role='user')
+        user_result = _make_scalars_result(msg)
+        ap = _make_ap(side_effect=[empty_result, user_result])
+
+        service = MonitoringService(ap)
+        found = await service._get_message_for_tool_context(message_id='missing', session_id='sess-1')
+
+        assert found is msg
+        assert ap.persistence_mgr.execute_async.await_count == 2
+
+
+class TestGetMessageForToolContextBySessionId:
+    """Tests for the session_id lookup paths."""
+
+    async def test_returns_user_message_when_found(self):
+        """Returns the most recent user message for the session."""
+        msg = _make_mock_message(session_id='sess-1', role='user')
+        user_result = _make_scalars_result(msg)
+        ap = _make_ap(side_effect=[user_result])
+
+        service = MonitoringService(ap)
+        found = await service._get_message_for_tool_context(session_id='sess-1')
+
+        assert found is msg
+        assert found.role == 'user'
+        ap.persistence_mgr.execute_async.assert_awaited_once()
+
+    async def test_falls_back_to_any_role_when_no_user_message(self):
+        """Falls back to any-role query when no user message exists."""
+        empty_user_result = _make_scalars_result(None)
+        any_msg = _make_mock_message(session_id='sess-1', role='assistant')
+        any_result = _make_scalars_result(any_msg)
+        ap = _make_ap(side_effect=[empty_user_result, any_result])
+
+        service = MonitoringService(ap)
+        found = await service._get_message_for_tool_context(session_id='sess-1')
+
+        assert found is any_msg
+        assert found.role == 'assistant'
+        assert ap.persistence_mgr.execute_async.await_count == 2
+
+
+class TestGetMessageForToolContextNone:
+    """Tests for None / degradation paths."""
+
+    async def test_returns_none_when_no_message_id_and_no_session_id(self):
+        """Returns None when neither message_id nor session_id is given."""
+        ap = _make_ap()
+        service = MonitoringService(ap)
+
+        found = await service._get_message_for_tool_context()
+
+        assert found is None
+        ap.persistence_mgr.execute_async.assert_not_awaited()
+
+    async def test_returns_none_when_message_id_missing_and_no_session_id(self):
+        """Returns None when message_id yields nothing and session_id is absent."""
+        empty_result = _make_scalars_result(None)
+        ap = _make_ap(side_effect=[empty_result])
+
+        service = MonitoringService(ap)
+        found = await service._get_message_for_tool_context(message_id='missing')
+
+        assert found is None
+
+    async def test_returns_none_when_all_queries_empty(self):
+        """Returns None when all three query paths return nothing."""
+        empty1 = _make_scalars_result(None)
+        empty2 = _make_scalars_result(None)
+        empty3 = _make_scalars_result(None)
+        ap = _make_ap(side_effect=[empty1, empty2, empty3])
+
+        service = MonitoringService(ap)
+        found = await service._get_message_for_tool_context(message_id='x', session_id='sess-1')
+
+        assert found is None
+        assert ap.persistence_mgr.execute_async.await_count == 3
+
+    async def test_returns_none_when_session_only_and_all_empty(self):
+        """Returns None when session_id queries all return nothing."""
+        empty_user = _make_scalars_result(None)
+        empty_any = _make_scalars_result(None)
+        ap = _make_ap(side_effect=[empty_user, empty_any])
+
+        service = MonitoringService(ap)
+        found = await service._get_message_for_tool_context(session_id='empty-session')
+
+        assert found is None
+        assert ap.persistence_mgr.execute_async.await_count == 2
+
+
+class TestGetMessageForToolContextUsesScalars:
+    """Verify the method calls scalars().first(), not first() + row[0]."""
+
+    async def test_uses_scalars_first_not_first(self):
+        """Ensure result.scalars().first() is used, not result.first().
+
+        This is the core regression test: the old code used result.first()
+        and then row[0], which returned a raw string on PostgreSQL/asyncpg.
+        The fix uses result.scalars().first() per SQLAlchemy 2.0 convention.
+        """
+        msg = _make_mock_message()
+        result = _make_scalars_result(msg)
+        ap = _make_ap(side_effect=[result])
+
+        service = MonitoringService(ap)
+        await service._get_message_for_tool_context(message_id='msg-1')
+
+        # scalars() must have been called
+        result.scalars.assert_called_once()
+        # scalars().first() must have been called
+        result.scalars.return_value.first.assert_called_once()
+        # first() on the result directly should NOT have been called
+        result.first.assert_not_called()