---
phase: 04-matrix-mvp-shared-agent-context-and-context-management-comma
plan: 01
type: execute
wave: 1
depends_on: []
files_modified:
- sdk/agent_api_wrapper.py
- sdk/agent_session.py
- sdk/real.py
- adapter/matrix/bot.py
- tests/platform/test_agent_session.py
- tests/platform/test_real.py
- tests/adapter/matrix/test_dispatcher.py
autonomous: true
requirements:
- Replace AgentSessionClient with AgentApi
- Wire AgentApi lifecycle into MatrixBot
must_haves:
truths:
- "RealPlatformClient uses AgentApiWrapper (wraps AgentApi), not AgentSessionClient"
- "AgentApiWrapper is connected before sync_forever and closed in finally block of main()"
- "build_thread_key and AgentSessionClient are gone from sdk/"
- "stream_message() yields MessageChunk objects including a final chunk with tokens_used from last_tokens_used"
- "AGENT_WS_URL is used unchanged (no thread_id query param)"
- "MATRIX_PLATFORM_BACKEND=real still works end-to-end without test crash"
- "All existing tests pass after the swap"
artifacts:
- path: "sdk/agent_api_wrapper.py"
provides: "AgentApiWrapper subclass of AgentApi with last_tokens_used tracking"
contains: "AgentApiWrapper"
- path: "sdk/real.py"
provides: "RealPlatformClient wrapping AgentApiWrapper"
contains: "AgentApiWrapper"
- path: "adapter/matrix/bot.py"
provides: "main() awaits agent_api.connect() and agent_api.close()"
contains: "agent_api.connect"
- path: "tests/platform/test_real.py"
provides: "Updated tests using FakeAgentApi instead of FakeAgentSessionClient"
key_links:
- from: "adapter/matrix/bot.py main()"
to: "RealPlatformClient._agent_api"
via: "runtime.platform.agent_api property"
pattern: "agent_api\\.connect"
- from: "sdk/real.py stream_message()"
to: "agent_api.last_tokens_used"
via: "attribute read after async-for loop"
pattern: "last_tokens_used"
---
Replace the custom per-request AgentSessionClient with a thin AgentApiWrapper that
subclasses AgentApi from lambda_agent_api and adds last_tokens_used tracking. Remove
build_thread_key and AgentSessionClient entirely. Wire AgentApiWrapper connect/close
into bot.py main(). Update all tests that referenced the old client.
Do NOT modify any file under external/. The external/ directory is managed by the
platform team. All customisation goes in sdk/agent_api_wrapper.py.
Purpose: The existing AgentSessionClient creates a new WebSocket per message and
injects thread_id into the URL — both incompatible with origin/main platform-agent.
AgentApi maintains a single persistent WS connection managed via connect()/close()
and exposes send_message() as an AsyncIterator. We capture tokens_used in a thin
subclass so sdk/real.py can include it in the final MessageChunk without touching
the upstream library.
Output: sdk/agent_api_wrapper.py (new), sdk/real.py (rewritten), sdk/agent_session.py
(stubbed), adapter/matrix/bot.py updated, tests green.
@$HOME/.claude/get-shit-done/workflows/execute-plan.md
@$HOME/.claude/get-shit-done/templates/summary.md
@.planning/PROJECT.md
@.planning/ROADMAP.md
@.planning/phases/04-matrix-mvp-shared-agent-context-and-context-management-comma/04-CONTEXT.md
@.planning/phases/04-matrix-mvp-shared-agent-context-and-context-management-comma/04-RESEARCH.md
From external/platform-agent_api/lambda_agent_api/agent_api.py (READ ONLY):
```python
class AgentApi:
def __init__(self, agent_id: str, url: str,
callback=None, on_disconnect=None): ...
async def connect(self) -> None: ... # opens WS, awaits MsgStatus, starts _listen task
async def close(self) -> None: ... # cancels _listen, closes WS+session
async def send_message(self, text: str) -> AsyncIterator[AgentEventUnion]:
# yields MsgEventTextChunk only; breaks on MsgEventEnd (does NOT yield it)
# MsgEventEnd.tokens_used is consumed internally at the break point
...
async def _listen(self) -> None:
# internal task: receives WS frames, puts AgentEventUnion into self._queue
# on MsgEventEnd: puts it in queue then breaks
...
# AgentEventUnion = Union[MsgEventTextChunk, MsgEventEnd] per server.py
```
From external/platform-agent_api/lambda_agent_api/server.py (READ ONLY):
```python
class MsgEventTextChunk(BaseModel):
type: Literal[EServerMessage.AGENT_EVENT_TEXT_CHUNK]
text: str
class MsgEventEnd(BaseModel):
type: Literal[EServerMessage.AGENT_EVENT_END]
tokens_used: int
```
New file to create — sdk/agent_api_wrapper.py:
```python
class AgentApiWrapper(AgentApi):
"""Thin subclass of AgentApi that captures tokens_used from MsgEventEnd.
AgentApi.send_message() yields only MsgEventTextChunk and breaks silently
on MsgEventEnd without storing tokens_used. This wrapper overrides _listen()
to intercept MsgEventEnd and store tokens_used before it is discarded.
"""
last_tokens_used: int = 0
async def _listen(self) -> None:
# Override: same as parent, but capture MsgEventEnd.tokens_used
...
```
From sdk/interface.py (unchanged):
```python
class MessageChunk(BaseModel):
message_id: str
delta: str
finished: bool
tokens_used: int = 0
class PlatformClient(Protocol):
async def send_message(self, user_id, chat_id, text, attachments=None) -> MessageResponse: ...
async def stream_message(self, user_id, chat_id, text, attachments=None) -> AsyncIterator[MessageChunk]: ...
```
Task 1: Create sdk/agent_api_wrapper.py; rewrite sdk/real.py; stub sdk/agent_session.py
- sdk/real.py (full file — being replaced)
- sdk/agent_session.py (full file — being stubbed)
- external/platform-agent_api/lambda_agent_api/agent_api.py (full file — READ ONLY, inspect _listen and send_message to understand override point)
- external/platform-agent_api/lambda_agent_api/server.py (full file — READ ONLY, MsgEventEnd.tokens_used)
- sdk/interface.py (MessageChunk, PlatformClient Protocol)
sdk/agent_api_wrapper.py, sdk/real.py, sdk/agent_session.py
- Create sdk/agent_api_wrapper.py with class AgentApiWrapper(AgentApi):
- __init__: calls super().__init__(...) and adds self.last_tokens_used: int = 0
- Override _listen(): copy the parent _listen() logic verbatim, then at the point where MsgEventEnd is received (before or as it is put into the queue), set self.last_tokens_used = chunk.tokens_used
- Do NOT modify agent_api.py in external/ — subclass only
- RealPlatformClient.__init__ accepts agent_api: AgentApiWrapper, prototype_state: PrototypeStateStore, platform: str = "matrix"
- RealPlatformClient exposes agent_api as property self.agent_api so bot.py main() can call connect/close
- stream_message() iterates agent_api.send_message(text) yielding MessageChunk per MsgEventTextChunk chunk; after loop yields final MessageChunk(finished=True, delta="", tokens_used=agent_api.last_tokens_used)
- send_message() collects all chunks from stream_message() and returns MessageResponse
- No thread_key, no build_thread_key references anywhere in sdk/real.py
- sdk/agent_session.py: replace file contents with single comment stub (keep file to avoid import errors until tests are updated in Task 2)
1. Read external/platform-agent_api/lambda_agent_api/agent_api.py fully to find the exact _listen() implementation and the line where MsgEventEnd is handled.
2. Create sdk/agent_api_wrapper.py:
```python
from __future__ import annotations
import sys
from pathlib import Path
# Ensure lambda_agent_api is importable (same sys.path trick as bot.py)
_api_root = Path(__file__).resolve().parents[1] / "external" / "platform-agent_api"
if str(_api_root) not in sys.path:
sys.path.insert(0, str(_api_root))
from lambda_agent_api.agent_api import AgentApi
from lambda_agent_api.server import MsgEventEnd
class AgentApiWrapper(AgentApi):
"""Thin subclass of AgentApi that captures tokens_used from MsgEventEnd.
AgentApi.send_message() yields MsgEventTextChunk events and breaks on
MsgEventEnd without storing tokens_used. This wrapper overrides _listen()
to intercept MsgEventEnd and set self.last_tokens_used before the event
is discarded, so RealPlatformClient can include it in the final MessageChunk.
Do NOT modify external/platform-agent_api — subclass only.
"""
def __init__(self, agent_id: str, url: str, **kwargs) -> None:
super().__init__(agent_id=agent_id, url=url, **kwargs)
self.last_tokens_used: int = 0
async def _listen(self) -> None:
# Copy parent _listen() logic.
# Read external/platform-agent_api/lambda_agent_api/agent_api.py _listen()
# and reproduce it here, adding:
# if isinstance(event, MsgEventEnd):
# self.last_tokens_used = event.tokens_used
# at the point where MsgEventEnd is processed.
#
# IMPORTANT: after reading agent_api.py, replace this entire method body
# with the exact parent implementation + the tokens_used capture line.
# Do not call super()._listen() — the parent creates a task; we need the
# override to run in the same task context.
raise NotImplementedError(
"Executor: replace this body with the copied _listen() from AgentApi "
"plus `self.last_tokens_used = event.tokens_used` at the MsgEventEnd branch."
)
```
IMPORTANT NOTE FOR EXECUTOR: The `_listen()` body above is a placeholder.
After reading agent_api.py, copy the actual _listen() implementation from AgentApi
into AgentApiWrapper._listen() and insert `self.last_tokens_used = event.tokens_used`
at the MsgEventEnd branch. The final file must NOT contain the NotImplementedError.
3. Rewrite sdk/real.py entirely:
```python
from __future__ import annotations
from typing import TYPE_CHECKING, AsyncIterator
from sdk.interface import Attachment, MessageChunk, MessageResponse, PlatformClient, User, UserSettings
from sdk.prototype_state import PrototypeStateStore
if TYPE_CHECKING:
from sdk.agent_api_wrapper import AgentApiWrapper
class RealPlatformClient(PlatformClient):
def __init__(
self,
agent_api: "AgentApiWrapper",
prototype_state: PrototypeStateStore,
platform: str = "matrix",
) -> None:
self._agent_api = agent_api
self._prototype_state = prototype_state
self._platform = platform
@property
def agent_api(self) -> "AgentApiWrapper":
return self._agent_api
async def get_or_create_user(
self,
external_id: str,
platform: str,
display_name: str | None = None,
) -> User:
return await self._prototype_state.get_or_create_user(
external_id=external_id,
platform=platform,
display_name=display_name,
)
async def send_message(
self,
user_id: str,
chat_id: str,
text: str,
attachments: list[Attachment] | None = None,
) -> MessageResponse:
parts: list[str] = []
tokens_used = 0
async for chunk in self.stream_message(user_id, chat_id, text, attachments):
if chunk.delta:
parts.append(chunk.delta)
if chunk.finished:
tokens_used = chunk.tokens_used
return MessageResponse(
message_id=user_id,
response="".join(parts),
tokens_used=tokens_used,
finished=True,
)
async def stream_message(
self,
user_id: str,
chat_id: str,
text: str,
attachments: list[Attachment] | None = None,
) -> AsyncIterator[MessageChunk]:
from lambda_agent_api.server import MsgEventTextChunk
async for event in self._agent_api.send_message(text):
if isinstance(event, MsgEventTextChunk):
yield MessageChunk(
message_id=user_id,
delta=event.text,
finished=False,
)
yield MessageChunk(
message_id=user_id,
delta="",
finished=True,
tokens_used=self._agent_api.last_tokens_used,
)
async def get_settings(self, user_id: str) -> UserSettings:
return await self._prototype_state.get_settings(user_id)
async def update_settings(self, user_id: str, action) -> None:
await self._prototype_state.update_settings(user_id, action)
```
4. Replace sdk/agent_session.py content with:
```python
# Deleted in Phase 4 — replaced by AgentApiWrapper from sdk/agent_api_wrapper.py
# File kept as stub to avoid import errors during migration; remove after test_agent_session.py is updated.
```
cd /Users/a/MAI/sem2/lambda/surfaces-bot && python -c "from sdk.real import RealPlatformClient; print('import ok')"
- sdk/agent_api_wrapper.py exists with AgentApiWrapper(AgentApi), __init__ sets self.last_tokens_used = 0, _listen() override captures MsgEventEnd.tokens_used
- sdk/real.py imports AgentApiWrapper (not AgentSessionClient or AgentApi directly), exposes self.agent_api property
- sdk/real.py stream_message yields final chunk with tokens_used from agent_api.last_tokens_used
- external/ directory has NO modifications
- sdk/agent_session.py contains only a comment stub (no class definitions)
- `python -c "from sdk.real import RealPlatformClient"` exits 0
- `grep "AgentApiWrapper" sdk/real.py` returns a match
- `grep "last_tokens_used" sdk/agent_api_wrapper.py` returns a match
Task 2: Wire AgentApiWrapper lifecycle into bot.py main(); update all broken tests
- adapter/matrix/bot.py (full file — _build_platform_from_env and main() need changes)
- tests/platform/test_agent_session.py (full file — delete or rewrite)
- tests/platform/test_real.py (full file — FakeAgentSessionClient → FakeAgentApi)
- tests/adapter/matrix/test_dispatcher.py (test_build_runtime_uses_real_platform — needs update)
adapter/matrix/bot.py, tests/platform/test_agent_session.py, tests/platform/test_real.py, tests/adapter/matrix/test_dispatcher.py
- _build_platform_from_env() returns a RealPlatformClient with an unconnected AgentApiWrapper (connect() NOT called here — called in main())
- main() calls await runtime.platform.agent_api.connect() after build_runtime() (only when backend is "real"; mock has no agent_api); wrap in `if hasattr(runtime.platform, "agent_api")` guard
- main() finally block: await agent_api.close() before await client.close()
- AGENT_WS_URL env var is passed unchanged to AgentApiWrapper(url=ws_url) — no query param manipulation
- test_agent_session.py: completely rewritten — remove all build_thread_key tests, remove AgentSessionClient tests; replace with 2 tests: (1) import check for lambda_agent_api module, (2) stub test that documents the deletion
- test_real.py: FakeAgentSessionClient replaced with FakeAgentApi that has send_message(text: str) -> AsyncIterator and last_tokens_used: int = 0; tests updated to construct RealPlatformClient(agent_api=FakeAgentApi(), prototype_state=PrototypeStateStore()); test_send_message no longer checks thread_key in message_id (now uses user_id); test_stream_message checks final chunk tokens_used comes from FakeAgentApi.last_tokens_used
- test_dispatcher.py: test_build_runtime_uses_real_platform_when_matrix_backend_is_real must NOT call agent_api.connect() (build_runtime only constructs, does not connect); update test to mock AgentApiWrapper so it does not attempt a real WS connection; assert isinstance(runtime.platform, RealPlatformClient) still passes
1. Edit adapter/matrix/bot.py:
a. Remove imports: `from sdk.agent_session import AgentSessionClient, AgentSessionConfig`
b. In _build_platform_from_env(), use AgentApiWrapper with lazy import:
```python
def _build_platform_from_env() -> PlatformClient:
backend = os.environ.get("MATRIX_PLATFORM_BACKEND", "mock").strip().lower()
if backend == "real":
import sys
_api_root = Path(__file__).resolve().parents[2] / "external" / "platform-agent_api"
if str(_api_root) not in sys.path:
sys.path.insert(0, str(_api_root))
from sdk.agent_api_wrapper import AgentApiWrapper
ws_url = os.environ["AGENT_WS_URL"]
agent_api = AgentApiWrapper(agent_id="matrix-bot", url=ws_url)
return RealPlatformClient(
agent_api=agent_api,
prototype_state=PrototypeStateStore(),
platform="matrix",
)
return MockPlatformClient()
```
c. In main(), after `runtime = build_runtime(store=SQLiteStore(db_path), client=client)`, add:
```python
if hasattr(runtime.platform, "agent_api"):
await runtime.platform.agent_api.connect()
```
d. In main() finally block, add before `await client.close()`:
```python
if hasattr(runtime.platform, "agent_api"):
await runtime.platform.agent_api.close()
```
2. Rewrite tests/platform/test_agent_session.py:
```python
"""
test_agent_session.py — stub after Phase 4 migration.
AgentSessionClient and build_thread_key were removed in Phase 4.
The platform client is now AgentApiWrapper wrapping AgentApi from lambda_agent_api.
See tests/platform/test_real.py for RealPlatformClient tests.
"""
import sys
from pathlib import Path
_api_root = Path(__file__).resolve().parents[2] / "external" / "platform-agent_api"
if str(_api_root) not in sys.path:
sys.path.insert(0, str(_api_root))
def test_lambda_agent_api_module_importable():
from lambda_agent_api.agent_api import AgentApi # noqa: F401
from lambda_agent_api.server import MsgEventTextChunk, MsgEventEnd # noqa: F401
assert True
def test_agent_session_module_is_stub():
"""Ensure old module no longer exposes AgentSessionClient or build_thread_key."""
import sdk.agent_session as mod
assert not hasattr(mod, "AgentSessionClient"), "AgentSessionClient should be removed"
assert not hasattr(mod, "build_thread_key"), "build_thread_key should be removed"
```
3. Rewrite tests/platform/test_real.py:
```python
from __future__ import annotations
import sys
from pathlib import Path
from typing import AsyncIterator
import pytest
from core.protocol import SettingsAction
from sdk.interface import MessageChunk, MessageResponse, UserSettings
from sdk.prototype_state import PrototypeStateStore
from sdk.real import RealPlatformClient
_api_root = Path(__file__).resolve().parents[2] / "external" / "platform-agent_api"
if str(_api_root) not in sys.path:
sys.path.insert(0, str(_api_root))
from lambda_agent_api.server import MsgEventTextChunk, EServerMessage # noqa: E402
class FakeAgentApi:
"""Minimal fake for AgentApiWrapper — no real WebSocket."""
def __init__(self) -> None:
self.last_tokens_used: int = 0
self.send_calls: list[str] = []
async def send_message(self, text: str) -> AsyncIterator[MsgEventTextChunk]:
self.send_calls.append(text)
self.last_tokens_used = 7
yield MsgEventTextChunk(type=EServerMessage.AGENT_EVENT_TEXT_CHUNK, text=text[:2])
yield MsgEventTextChunk(type=EServerMessage.AGENT_EVENT_TEXT_CHUNK, text=text[2:])
# send_message() in real AgentApi breaks on MsgEventEnd without yielding it;
# FakeAgentApi mirrors this by not yielding MsgEventEnd — last_tokens_used is set directly.
@pytest.mark.asyncio
async def test_real_platform_client_get_or_create_user_uses_local_state():
client = RealPlatformClient(
agent_api=FakeAgentApi(),
prototype_state=PrototypeStateStore(),
)
first = await client.get_or_create_user("u1", "matrix", "Alice")
second = await client.get_or_create_user("u1", "matrix")
assert first.user_id == "usr-matrix-u1"
assert first.is_new is True
assert second.user_id == first.user_id
assert second.is_new is False
assert second.display_name == "Alice"
@pytest.mark.asyncio
async def test_real_platform_client_send_message_calls_agent_with_text():
fake = FakeAgentApi()
client = RealPlatformClient(agent_api=fake, prototype_state=PrototypeStateStore())
result = await client.send_message("@alice:example.org", "C1", "hello")
assert result.response == "hello"
assert result.tokens_used == 7
assert fake.send_calls == ["hello"]
@pytest.mark.asyncio
async def test_real_platform_client_stream_message_yields_chunks_and_final_with_tokens():
fake = FakeAgentApi()
client = RealPlatformClient(agent_api=fake, prototype_state=PrototypeStateStore())
chunks = []
async for chunk in client.stream_message("@alice:example.org", "C1", "hello"):
chunks.append(chunk)
assert chunks[-1].finished is True
assert chunks[-1].tokens_used == 7
assert "".join(c.delta for c in chunks) == "hello"
@pytest.mark.asyncio
async def test_real_platform_client_settings_are_local():
client = RealPlatformClient(
agent_api=FakeAgentApi(),
prototype_state=PrototypeStateStore(),
)
await client.update_settings(
"usr-matrix-u1",
SettingsAction(action="toggle_skill", payload={"skill": "browser", "enabled": True}),
)
settings = await client.get_settings("usr-matrix-u1")
assert isinstance(settings, UserSettings)
assert settings.skills["browser"] is True
```
4. Edit tests/adapter/matrix/test_dispatcher.py — update `test_build_runtime_uses_real_platform_when_matrix_backend_is_real`:
- Add sys.path setup for lambda_agent_api (same pattern as above)
- Mock AgentApiWrapper so it does not open a real WS:
```python
async def test_build_runtime_uses_real_platform_when_matrix_backend_is_real(monkeypatch):
import sys
from pathlib import Path
_api_root = Path(__file__).resolve().parents[3] / "external" / "platform-agent_api"
if str(_api_root) not in sys.path:
sys.path.insert(0, str(_api_root))
monkeypatch.setenv("MATRIX_PLATFORM_BACKEND", "real")
monkeypatch.setenv("AGENT_WS_URL", "ws://agent.example/agent_ws/")
# Patch AgentApiWrapper to avoid real WS connection during build_runtime
import sdk.agent_api_wrapper as _mod
class _FakeAgentApiWrapper:
def __init__(self, agent_id, url, **kw):
self.last_tokens_used = 0
async def connect(self): pass
async def close(self): pass
async def send_message(self, text):
return; yield # empty async generator
monkeypatch.setattr(_mod, "AgentApiWrapper", _FakeAgentApiWrapper)
from adapter.matrix.bot import build_runtime
from sdk.real import RealPlatformClient
runtime = build_runtime()
assert isinstance(runtime.platform, RealPlatformClient)
```
cd /Users/a/MAI/sem2/lambda/surfaces-bot && python -m pytest tests/platform/test_agent_session.py tests/platform/test_real.py tests/adapter/matrix/test_dispatcher.py -v 2>&1 | tail -20
- All tests in test_agent_session.py, test_real.py, test_dispatcher.py pass
- main() in bot.py has agent_api.connect() call guarded by hasattr check
- main() finally block closes agent_api before matrix client
- grep confirms no "AgentSessionClient" or "build_thread_key" remain in sdk/real.py or adapter/matrix/bot.py
- grep confirms no modifications to any file under external/
## Trust Boundaries
| Boundary | Description |
|----------|-------------|
| bot → platform-agent WS | Outbound WS to agent service; input is user text |
| env vars → bot config | AGENT_WS_URL, MATRIX_PLATFORM_BACKEND read from environment |
## STRIDE Threat Register
| Threat ID | Category | Component | Disposition | Mitigation Plan |
|-----------|----------|-----------|-------------|-----------------|
| T-04-01-01 | Tampering | AgentApiWrapper.send_message() text | accept | Single-user prototype; text originates from authenticated Matrix user |
| T-04-01-02 | Denial of Service | AgentBusyException from concurrent sends | mitigate | AgentApi._request_lock already prevents concurrent sends; bot must surface error to user instead of crashing |
| T-04-01-03 | Information Disclosure | AGENT_WS_URL in env | accept | Internal service URL; not exposed to users |
Run full test suite after both tasks complete:
```bash
cd /Users/a/MAI/sem2/lambda/surfaces-bot && python -m pytest tests/ -v 2>&1 | tail -30
```
Grep checks:
```bash
# No old imports should remain
grep -r "AgentSessionClient\|build_thread_key" sdk/ adapter/ tests/ --include="*.py" | grep -v "stub\|Deleted\|removed"
# AgentApiWrapper wired in bot.py
grep "agent_api.connect\|agent_api.close" adapter/matrix/bot.py
# last_tokens_used set in wrapper
grep "last_tokens_used" sdk/agent_api_wrapper.py
# No external/ files modified
git diff --name-only external/
```
- `pytest tests/platform/ tests/adapter/matrix/test_dispatcher.py -v` exits 0 with no failures
- `grep -r "AgentSessionClient" sdk/ adapter/` returns empty (or only the stub comment)
- `grep -r "build_thread_key" sdk/ adapter/` returns empty
- `grep "agent_api.connect" adapter/matrix/bot.py` returns a match
- `grep "last_tokens_used" sdk/agent_api_wrapper.py` returns the assignment line
- `git diff --name-only external/` returns empty (external/ untouched)