From cdcab6c4f236d0787305e1b3be1b5fd3d8f60b20 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Mon, 4 May 2026 14:30:16 +0800 Subject: [PATCH 1/2] feat(session): publish subagent runtime messages Product-side session management needs a narrow subagent_message publishing contract so the existing pipeline prompt wiring can consume messages without expanding agentkit or HTTP scope. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../src/coding_agent/ui/session_manager.py | 44 ++++++ .../ui/test_session_manager_public_api.py | 141 ++++++++++++++++++ 2 files changed, 185 insertions(+) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 14183235f..1e0f2f2b1 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -77,6 +77,10 @@ def _approval_decision_message_id(session_id: str, request_id: str) -> str: return f"approval_decision:{session_id}:{request_id}" +def _subagent_message_id(session_id: str) -> str: + return f"subagent_message:{session_id}:{uuid.uuid4().hex}" + + def _approval_response_projection(response: ApprovalResponse) -> dict[str, Any]: return { "request_id": response.request_id, @@ -1622,6 +1626,46 @@ async def submit_approval( ) return response is not None + async def publish_subagent_message( + self, + session_id: str, + text: str, + *, + message_id: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> bool: + if not isinstance(text, str) or not text.strip(): + raise ValueError("text must be a non-empty string") + if message_id is not None and not message_id: + raise ValueError("message_id must be None or a non-empty string") + + await self._assert_owner(session_id) + session = await self.get_session_async(session_id) + effective_message_id = message_id or _subagent_message_id(session_id) + payload: dict[str, Any] = {"text": text} + if metadata is not None: + payload["metadata"] = dict(metadata) + + try: + await session.runtime_message_bus.publish( + RuntimeMessage( + message_id=effective_message_id, + kind=RuntimeMessageKind.SUBAGENT_MESSAGE, + payload=payload, + ) + ) + except DuplicateRuntimeMessageError as exc: + if exc.message_id != effective_message_id: + raise + logger.info( + "subagent_message already published for session %s message %s", + session_id, + effective_message_id, + ) + session.last_activity = datetime.now() + await self._persist_session_async(session) + return True + async def wait_for_http_approval( self, session_id: str, diff --git a/coding-agent/tests/ui/test_session_manager_public_api.py b/coding-agent/tests/ui/test_session_manager_public_api.py index 36ef476ca..0889ca7a4 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -4,12 +4,18 @@ import types from datetime import datetime from pathlib import Path +from typing import Any, cast from unittest.mock import patch import pytest from agentkit.checkpoint import CheckpointService from agentkit.checkpoint.models import CheckpointMeta +from agentkit.plugin.registry import PluginRegistry +from agentkit.providers.models import DoneEvent, TextEvent +from agentkit.runtime import RuntimeMessageCursor, RuntimeMessageKind +from agentkit.runtime.hook_runtime import HookRuntime +from agentkit.runtime.pipeline import Pipeline, PipelineContext from agentkit.tape.tape import Tape from coding_agent.approval.store import ApprovalStore from coding_agent.approval import ApprovalPolicy @@ -31,6 +37,46 @@ ) +class _PromptCaptureProvider: + def __init__(self) -> None: + self.messages: list[list[dict[str, object]]] = [] + + async def stream(self, messages, tools=None, **kwargs): + del tools, kwargs + self.messages.append(messages) + yield TextEvent(text="ok") + yield DoneEvent() + + +class _PromptCapturePlugin: + state_key = "prompt_capture" + + def __init__(self) -> None: + self.provider = _PromptCaptureProvider() + + def hooks(self): + return {"provide_llm": self.provide_llm} + + def provide_llm(self, **kwargs): + del kwargs + return self.provider + + +def _pipeline_with_prompt_capture() -> tuple[Pipeline, _PromptCaptureProvider]: + registry = PluginRegistry() + plugin = _PromptCapturePlugin() + registry.register(plugin) + return Pipeline(runtime=HookRuntime(registry), registry=registry), plugin.provider + + +def _system_prompt_content(messages: list[dict[str, object]]) -> str: + return "\n".join( + str(message.get("content", "")) + for message in messages + if message.get("role") == "system" + ) + + def test_register_session_uses_public_api() -> None: manager = SessionManager() approval_store = ApprovalStore() @@ -1198,6 +1244,101 @@ async def test_submit_approval_rejects_stale_pending_projection_without_store_re assert session.approval_event.is_set() is False +@pytest.mark.asyncio +async def test_publish_subagent_message_reaches_next_pipeline_prompt() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session_id = await manager.create_session() + session = manager.get_session(session_id) + pipeline, provider = _pipeline_with_prompt_capture() + ctx = PipelineContext( + tape=Tape(), + session_id=session_id, + config={"system_prompt": "system"}, + ) + session.runtime_pipeline = pipeline + session.runtime_ctx = ctx + session.runtime_adapter = object() + + published = await manager.publish_subagent_message( + session_id, + "worker finished analysis", + message_id="subagent-msg-1", + ) + ctx.runtime_message_bus = session.runtime_message_bus + await pipeline.run_turn(ctx) + + assert published is True + system_content = _system_prompt_content(provider.messages[0]) + assert "subagent_message subagent-msg-1: worker finished analysis" in system_content + assert ctx.runtime_message_cursor.sequence == 1 + + +@pytest.mark.asyncio +async def test_publish_subagent_message_duplicate_id_is_idempotent() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session_id = await manager.create_session() + session = manager.get_session(session_id) + + first = await manager.publish_subagent_message( + session_id, + "first body wins", + message_id="subagent-msg-dup", + ) + second = await manager.publish_subagent_message( + session_id, + "retry body ignored", + message_id="subagent-msg-dup", + ) + batch = await session.runtime_message_bus.consume_after( + RuntimeMessageCursor(), + kinds={RuntimeMessageKind.SUBAGENT_MESSAGE}, + ) + + assert first is True + assert second is True + assert [item.message.payload for item in batch.messages] == [ + {"text": "first body wins"} + ] + + +@pytest.mark.asyncio +async def test_publish_subagent_message_does_not_advance_approval_cursor() -> None: + manager = SessionManager(store=InMemorySessionStore()) + session_id = await manager.create_session() + session = manager.get_session(session_id) + + published = await manager.publish_subagent_message( + session_id, + "child update", + message_id="subagent-msg-cursor", + ) + + assert published is True + assert session.approval_decision_cursor == RuntimeMessageCursor() + approval_batch = await session.runtime_message_bus.consume_after( + RuntimeMessageCursor(), + kinds={RuntimeMessageKind.APPROVAL_DECISION}, + ) + subagent_batch = await session.runtime_message_bus.consume_after( + RuntimeMessageCursor(), + kinds={RuntimeMessageKind.SUBAGENT_MESSAGE}, + ) + assert approval_batch.messages == () + assert [item.message.message_id for item in subagent_batch.messages] == [ + "subagent-msg-cursor" + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("text", ["", " ", 123]) +async def test_publish_subagent_message_rejects_invalid_text(text: object) -> None: + manager = SessionManager(store=InMemorySessionStore()) + session_id = await manager.create_session() + + with pytest.raises(ValueError, match="text must be a non-empty string"): + await manager.publish_subagent_message(session_id, cast(Any, text)) + + @pytest.mark.asyncio async def test_list_checkpoints_returns_metadata_for_session_tape() -> None: manager = SessionManager(store=InMemorySessionStore()) From 709deb10d0f3a67e166d798e1340984ee1733b08 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Mon, 4 May 2026 15:26:19 +0800 Subject: [PATCH 2/2] fix(session): prevent publish resurrection race Hold the session lock across subagent message publish and persistence so a concurrent remove cannot restore deleted state. Add a regression test that proves removal waits until publish finishes and leaves the session deleted. --- .../src/coding_agent/ui/session_manager.py | 47 ++++++++++--------- .../ui/test_session_manager_public_api.py | 42 +++++++++++++++++ 2 files changed, 66 insertions(+), 23 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 1e0f2f2b1..2008574c6 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -1639,31 +1639,32 @@ async def publish_subagent_message( if message_id is not None and not message_id: raise ValueError("message_id must be None or a non-empty string") - await self._assert_owner(session_id) - session = await self.get_session_async(session_id) - effective_message_id = message_id or _subagent_message_id(session_id) - payload: dict[str, Any] = {"text": text} - if metadata is not None: - payload["metadata"] = dict(metadata) + async with self._lock: + await self._assert_owner(session_id) + session = await self.get_session_async(session_id) + effective_message_id = message_id or _subagent_message_id(session_id) + payload: dict[str, Any] = {"text": text} + if metadata is not None: + payload["metadata"] = dict(metadata) - try: - await session.runtime_message_bus.publish( - RuntimeMessage( - message_id=effective_message_id, - kind=RuntimeMessageKind.SUBAGENT_MESSAGE, - payload=payload, + try: + await session.runtime_message_bus.publish( + RuntimeMessage( + message_id=effective_message_id, + kind=RuntimeMessageKind.SUBAGENT_MESSAGE, + payload=payload, + ) ) - ) - except DuplicateRuntimeMessageError as exc: - if exc.message_id != effective_message_id: - raise - logger.info( - "subagent_message already published for session %s message %s", - session_id, - effective_message_id, - ) - session.last_activity = datetime.now() - await self._persist_session_async(session) + except DuplicateRuntimeMessageError as exc: + if exc.message_id != effective_message_id: + raise + logger.info( + "subagent_message already published for session %s message %s", + session_id, + effective_message_id, + ) + session.last_activity = datetime.now() + await self._persist_session_async(session) return True async def wait_for_http_approval( diff --git a/coding-agent/tests/ui/test_session_manager_public_api.py b/coding-agent/tests/ui/test_session_manager_public_api.py index 0889ca7a4..3457f6daf 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -1339,6 +1339,48 @@ async def test_publish_subagent_message_rejects_invalid_text(text: object) -> No await manager.publish_subagent_message(session_id, cast(Any, text)) +@pytest.mark.asyncio +async def test_publish_subagent_message_does_not_resurrect_removed_session() -> None: + store = InMemorySessionStore() + manager = SessionManager(store=store) + session_id = await manager.create_session() + session = manager.get_session(session_id) + publish_started = asyncio.Event() + allow_publish = asyncio.Event() + + class BlockingBus: + async def publish(self, message: object) -> object: + del message + publish_started.set() + await allow_publish.wait() + return object() + + session.runtime_message_bus = cast(Any, BlockingBus()) + + publish_task = asyncio.create_task( + manager.publish_subagent_message( + session_id, + "worker finished analysis", + message_id="subagent-msg-race", + ) + ) + await publish_started.wait() + + remove_task = asyncio.create_task(manager.remove_session_async(session_id)) + await asyncio.sleep(0) + + assert remove_task.done() is False + assert manager.has_session(session_id) is True + assert store.get(session_id) is not None + + allow_publish.set() + assert await publish_task is True + await remove_task + + assert manager.has_session(session_id) is False + assert store.get(session_id) is None + + @pytest.mark.asyncio async def test_list_checkpoints_returns_metadata_for_session_tape() -> None: manager = SessionManager(store=InMemorySessionStore())