From 62d2d2a0425aaca9fd01a57e499fc153bb8e7781 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Mon, 4 May 2026 01:40:31 +0800 Subject: [PATCH 1/3] fix(agentkit): preserve late runtime prompts Rebuild prompt context when runtime messages arrive after build_context but before model streaming, and avoid consuming prompt-affecting messages after run_model until the next turn. --- coding-agent/src/agentkit/runtime/pipeline.py | 68 +++++++-- .../tests/agentkit/runtime/test_pipeline.py | 134 +++++++++++++++++- 2 files changed, 181 insertions(+), 21 deletions(-) diff --git a/coding-agent/src/agentkit/runtime/pipeline.py b/coding-agent/src/agentkit/runtime/pipeline.py index d7368b085..6294eaaff 100644 --- a/coding-agent/src/agentkit/runtime/pipeline.py +++ b/coding-agent/src/agentkit/runtime/pipeline.py @@ -50,6 +50,13 @@ RuntimeMessageKind.SYSTEM_NOTICE, } ) +_PROMPT_RUNTIME_MESSAGE_KINDS = frozenset( + { + RuntimeMessageKind.USER_STEER, + RuntimeMessageKind.SUBAGENT_MESSAGE, + RuntimeMessageKind.SYSTEM_NOTICE, + } +) StructuredToolResultScopeFactory = Callable[[bool], AbstractContextManager[None]] @@ -248,16 +255,16 @@ async def _consume_runtime_messages( ctx: PipelineContext, *, stage: StageName, - ) -> None: + ) -> bool: if ctx.runtime_message_bus is None: - return + return False batch = await ctx.runtime_message_bus.consume_after( ctx.runtime_message_cursor, kinds=_PIPELINE_RUNTIME_MESSAGE_KINDS, ) if not batch.messages: - return + return False for item in batch.messages: message = item.message @@ -269,6 +276,30 @@ async def _consume_runtime_messages( ctx.runtime_messages.extend(batch.messages) ctx.runtime_message_cursor = batch.cursor + return any( + item.message.kind in _PROMPT_RUNTIME_MESSAGE_KINDS + for item in batch.messages + ) + + async def _raise_if_runtime_interrupted( + self, + ctx: PipelineContext, + *, + stage: StageName, + ) -> None: + if ctx.runtime_message_bus is None: + return + + batch = await ctx.runtime_message_bus.consume_after( + ctx.runtime_message_cursor, + kinds={RuntimeMessageKind.INTERRUPT}, + ) + for item in batch.messages: + message = item.message + reason = _string_payload_value(message.payload, "reason", "text") + if not reason: + reason = message.message_id + raise PipelineError(f"runtime interrupted: {reason}", stage=stage) def _runtime_context_grounding(self, ctx: PipelineContext) -> dict[str, Any] | None: lines: list[str] = [] @@ -342,9 +373,18 @@ async def run_turn(self, ctx: PipelineContext) -> PipelineContext: try: for stage in self.STAGES: try: - await self._consume_runtime_messages(ctx, stage=stage) + if stage in {"save_state", "render", "dispatch"}: + await self._raise_if_runtime_interrupted(ctx, stage=stage) + runtime_prompt_changed = False + else: + runtime_prompt_changed = await self._consume_runtime_messages( + ctx, + stage=stage, + ) handler = getattr(self, f"_stage_{stage}", None) if handler is not None: + if stage == "run_model" and runtime_prompt_changed: + await self._stage_build_context(ctx) if _tracer is not None: _tracer.info( "stage_start", stage=stage, entry_count=len(ctx.tape) @@ -519,7 +559,8 @@ async def _stage_run_model(self, ctx: PipelineContext) -> None: toolset = self._require_toolset(ctx, stage="run_model") for _round in range(max_tool_rounds): - await self._consume_runtime_messages(ctx, stage="run_model") + if await self._consume_runtime_messages(ctx, stage="run_model"): + await self._stage_build_context(ctx) tool_dicts = ( [s.to_openai_format() for s in ctx.tool_schemas] if ctx.tool_schemas @@ -643,7 +684,9 @@ async def _stage_run_model(self, ctx: PipelineContext) -> None: approval = await toolset.approve_tool_call(tool_call, ctx=ctx) if not approval.approved: - rejection_msg = f"Tool call rejected: {approval.reason or 'policy'}" + rejection_msg = ( + f"Tool call rejected: {approval.reason or 'policy'}" + ) ctx.tape.append( Entry( kind="tool_result", @@ -667,7 +710,8 @@ async def _stage_run_model(self, ctx: PipelineContext) -> None: executable_calls.append(tool_call) if executable_calls: - await self._consume_runtime_messages(ctx, stage="run_model") + if await self._consume_runtime_messages(ctx, stage="run_model"): + await self._stage_build_context(ctx) max_size = ctx.config.get("max_tool_result_size", 10000) structured_results_enabled = bool( ctx.config.get("structured_results", False) @@ -678,13 +722,9 @@ async def _stage_run_model(self, ctx: PipelineContext) -> None: executable_calls, ctx=ctx, options=ToolExecutionOptions( - timeout_seconds=ctx.config.get( - "tool_timeout_seconds" - ), - max_retries=int( - ctx.config.get("tool_max_retries", 0) - ), - ) + timeout_seconds=ctx.config.get("tool_timeout_seconds"), + max_retries=int(ctx.config.get("tool_max_retries", 0)), + ), ) for result in execution_results: diff --git a/coding-agent/tests/agentkit/runtime/test_pipeline.py b/coding-agent/tests/agentkit/runtime/test_pipeline.py index 7297ef362..e31af003e 100644 --- a/coding-agent/tests/agentkit/runtime/test_pipeline.py +++ b/coding-agent/tests/agentkit/runtime/test_pipeline.py @@ -374,12 +374,128 @@ async def mock_stream(messages, tools=None, **kwargs): assert "environment: local" in system_content assert "workspace_root: /repo" in system_content assert "elapsed_seconds:" in system_content - assert "context_budget: max_input=128000 reserved_output=4096 max_output=8192" in system_content + assert ( + "context_budget: max_input=128000 reserved_output=4096 max_output=8192" + in system_content + ) assert "active_approvals: req-1:bash_run" in system_content assert "user_steer msg-steer: Prefer a short answer" in system_content assert "system_notice msg-notice: Checkpoint restored" in system_content assert ctx.runtime_message_cursor.sequence == 2 + @pytest.mark.asyncio + async def test_runtime_message_consumed_after_build_context_reaches_prompt( + self, setup + ): + pipeline, plugin = setup + bus = InMemoryRuntimeMessageBus() + captured_messages: list[list[dict[str, object]]] = [] + + async def mock_stream(messages, tools=None, **kwargs): + del tools, kwargs + captured_messages.append(messages) + from agentkit.providers.models import DoneEvent, TextEvent + + yield TextEvent(text="Hello back!") + yield DoneEvent() + + original_build_context = pipeline._stage_build_context + published_late_message = False + + async def build_context_then_publish(ctx): + nonlocal published_late_message + await original_build_context(ctx) + if published_late_message: + return + published_late_message = True + await bus.publish( + RuntimeMessage( + message_id="msg-late-steer", + kind=RuntimeMessageKind.USER_STEER, + payload={"text": "Use the new instruction"}, + ) + ) + + plugin._mock_llm.stream = mock_stream + pipeline._stage_build_context = build_context_then_publish + ctx = PipelineContext( + tape=Tape(), + session_id="session-1", + runtime_message_bus=bus, + config={"system_prompt": "system"}, + ) + + await pipeline.run_turn(ctx) + + system_content = "\n".join( + str(message.get("content", "")) + for message in captured_messages[0] + if message.get("role") == "system" + ) + assert "user_steer msg-late-steer: Use the new instruction" in system_content + assert ctx.runtime_message_cursor.sequence == 1 + + @pytest.mark.asyncio + async def test_prompt_runtime_message_after_run_model_is_kept_for_next_turn( + self, setup + ): + pipeline, plugin = setup + bus = InMemoryRuntimeMessageBus() + captured_messages: list[list[dict[str, object]]] = [] + + async def mock_stream(messages, tools=None, **kwargs): + del tools, kwargs + captured_messages.append(messages) + from agentkit.providers.models import DoneEvent, TextEvent + + yield TextEvent(text="Hello back!") + yield DoneEvent() + + original_run_model = pipeline._stage_run_model + published_late_message = False + + async def run_model_then_publish(ctx): + nonlocal published_late_message + await original_run_model(ctx) + if published_late_message: + return + published_late_message = True + await bus.publish( + RuntimeMessage( + message_id="msg-after-model", + kind=RuntimeMessageKind.USER_STEER, + payload={"text": "Use this next turn"}, + ) + ) + + plugin._mock_llm.stream = mock_stream + pipeline._stage_run_model = run_model_then_publish + ctx = PipelineContext( + tape=Tape(), + session_id="session-1", + runtime_message_bus=bus, + config={"system_prompt": "system"}, + ) + + await pipeline.run_turn(ctx) + await pipeline.run_turn(ctx) + + first_system_content = "\n".join( + str(message.get("content", "")) + for message in captured_messages[0] + if message.get("role") == "system" + ) + second_system_content = "\n".join( + str(message.get("content", "")) + for message in captured_messages[1] + if message.get("role") == "system" + ) + assert ( + "user_steer msg-after-model: Use this next turn" not in first_system_content + ) + assert "user_steer msg-after-model: Use this next turn" in second_system_content + assert ctx.runtime_message_cursor.sequence == 1 + @pytest.mark.asyncio async def test_runtime_context_includes_elapsed_without_run_context(self, setup): pipeline, plugin = setup @@ -411,9 +527,7 @@ async def mock_stream(messages, tools=None, **kwargs): assert "elapsed_seconds:" in system_content @pytest.mark.asyncio - async def test_runtime_approval_decision_is_left_for_product_consumer( - self, setup - ): + async def test_runtime_approval_decision_is_left_for_product_consumer(self, setup): pipeline, plugin = setup bus = InMemoryRuntimeMessageBus() await bus.publish( @@ -499,7 +613,9 @@ async def mock_stream(messages, tools=None, **kwargs): runtime_message_bus=bus, ) - with pytest.raises(PipelineError, match="runtime interrupted: user stopped turn"): + with pytest.raises( + PipelineError, match="runtime interrupted: user stopped turn" + ): await pipeline.run_turn(ctx) assert stream_called is False @@ -542,7 +658,9 @@ async def mock_stream(messages, tools=None, **kwargs): runtime_message_bus=bus, ) - with pytest.raises(PipelineError, match="runtime interrupted: user stopped turn"): + with pytest.raises( + PipelineError, match="runtime interrupted: user stopped turn" + ): await pipeline.run_turn(ctx) retry_batch = await bus.consume_after(RuntimeMessageCursor()) @@ -721,7 +839,9 @@ async def mock_stream(messages, tools=None, **kwargs): await pipeline.run_turn(ctx) - tool_result_entries = [entry for entry in ctx.tape if entry.kind == "tool_result"] + tool_result_entries = [ + entry for entry in ctx.tape if entry.kind == "tool_result" + ] assert tool_result_entries[0].payload == { "tool_call_id": "tc-invalid", "content": "Tool call validation failed: unexpected argument: extra", From f33a3c2700c37b3a8c6b17838539bee5969813e6 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Mon, 4 May 2026 01:52:48 +0800 Subject: [PATCH 2/3] test(agentkit): cover stream-time runtime prompts Lock in that prompt-affecting runtime messages published during active model streaming stay queued until the next turn instead of being consumed by post-model stages. --- .../tests/agentkit/runtime/test_pipeline.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/coding-agent/tests/agentkit/runtime/test_pipeline.py b/coding-agent/tests/agentkit/runtime/test_pipeline.py index e31af003e..56dd521fe 100644 --- a/coding-agent/tests/agentkit/runtime/test_pipeline.py +++ b/coding-agent/tests/agentkit/runtime/test_pipeline.py @@ -496,6 +496,67 @@ async def run_model_then_publish(ctx): assert "user_steer msg-after-model: Use this next turn" in second_system_content assert ctx.runtime_message_cursor.sequence == 1 + @pytest.mark.asyncio + async def test_prompt_runtime_message_published_during_stream_is_kept_for_next_turn( + self, setup + ): + pipeline, plugin = setup + bus = InMemoryRuntimeMessageBus() + captured_messages: list[list[dict[str, object]]] = [] + published_during_stream = False + + async def mock_stream(messages, tools=None, **kwargs): + nonlocal published_during_stream + del tools, kwargs + captured_messages.append(messages) + from agentkit.providers.models import DoneEvent, TextEvent + + if not published_during_stream: + published_during_stream = True + yield TextEvent(text="Hello") + await bus.publish( + RuntimeMessage( + message_id="msg-during-stream", + kind=RuntimeMessageKind.USER_STEER, + payload={"text": "Use this after stream"}, + ) + ) + yield TextEvent(text=" back!") + else: + yield TextEvent(text="Hello back!") + yield DoneEvent() + + plugin._mock_llm.stream = mock_stream + ctx = PipelineContext( + tape=Tape(), + session_id="session-1", + runtime_message_bus=bus, + config={"system_prompt": "system"}, + ) + + await pipeline.run_turn(ctx) + await pipeline.run_turn(ctx) + + first_system_content = "\n".join( + str(message.get("content", "")) + for message in captured_messages[0] + if message.get("role") == "system" + ) + second_system_content = "\n".join( + str(message.get("content", "")) + for message in captured_messages[1] + if message.get("role") == "system" + ) + assert ( + "user_steer msg-during-stream: Use this after stream" + not in first_system_content + ) + assert ( + "user_steer msg-during-stream: Use this after stream" + in second_system_content + ) + assert ctx.runtime_message_cursor.sequence == 1 + @pytest.mark.asyncio async def test_runtime_context_includes_elapsed_without_run_context(self, setup): pipeline, plugin = setup From 672109f56728ff1b993780298de20f597f6705ba Mon Sep 17 00:00:00 2001 From: kina0630 Date: Mon, 4 May 2026 02:08:22 +0800 Subject: [PATCH 3/3] fix(agentkit): keep pre-tool checkpoint interrupt-only Leave prompt-affecting runtime messages queued when a tool-call turn has not yet reached the next model invocation, so failed tool execution or rebuild paths can replay them on the next turn. --- coding-agent/src/agentkit/runtime/pipeline.py | 6 +- .../tests/agentkit/runtime/test_pipeline.py | 77 +++++++++++++++++++ 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/coding-agent/src/agentkit/runtime/pipeline.py b/coding-agent/src/agentkit/runtime/pipeline.py index 6294eaaff..c7365d6fc 100644 --- a/coding-agent/src/agentkit/runtime/pipeline.py +++ b/coding-agent/src/agentkit/runtime/pipeline.py @@ -710,8 +710,10 @@ async def _stage_run_model(self, ctx: PipelineContext) -> None: executable_calls.append(tool_call) if executable_calls: - if await self._consume_runtime_messages(ctx, stage="run_model"): - await self._stage_build_context(ctx) + await self._raise_if_runtime_interrupted( + ctx, + stage="run_model", + ) max_size = ctx.config.get("max_tool_result_size", 10000) structured_results_enabled = bool( ctx.config.get("structured_results", False) diff --git a/coding-agent/tests/agentkit/runtime/test_pipeline.py b/coding-agent/tests/agentkit/runtime/test_pipeline.py index 56dd521fe..4d58a1ba6 100644 --- a/coding-agent/tests/agentkit/runtime/test_pipeline.py +++ b/coding-agent/tests/agentkit/runtime/test_pipeline.py @@ -817,6 +817,83 @@ async def mock_stream(messages, tools=None, **kwargs): assert any(e.kind == "tool_result" for e in entries) assert entries[-1].payload["role"] == "assistant" + @pytest.mark.asyncio + async def test_prompt_runtime_message_before_tool_execution_survives_failed_turn( + self, setup + ): + pipeline, plugin = setup + bus = InMemoryRuntimeMessageBus() + tape = Tape() + tape.append( + Entry(kind="message", payload={"role": "user", "content": "read file.txt"}) + ) + ctx = PipelineContext(tape=tape, session_id="s1", runtime_message_bus=bus) + await pipeline.mount(ctx) + + from agentkit.providers.models import DoneEvent, TextEvent, ToolCallEvent + + call_count = 0 + captured_messages: list[list[dict[str, object]]] = [] + published_late_message = False + original_build_context = pipeline._stage_build_context + fail_rebuild_once = False + + async def mock_stream(messages, tools=None, **kwargs): + nonlocal call_count, published_late_message + del tools, kwargs + captured_messages.append(messages) + call_count += 1 + if call_count == 1: + yield ToolCallEvent( + tool_call_id="tc1", name="file_read", arguments={"path": "file.txt"} + ) + if not published_late_message: + published_late_message = True + await bus.publish( + RuntimeMessage( + message_id="msg-before-tool-execution", + kind=RuntimeMessageKind.USER_STEER, + payload={"text": "Keep this steer after failure"}, + ) + ) + yield DoneEvent() + else: + yield TextEvent(text="File contents: test data") + yield DoneEvent() + + async def build_context_fail_after_tool_call(current_ctx): + nonlocal fail_rebuild_once + if fail_rebuild_once and any( + entry.kind == "tool_call" for entry in current_ctx.tape + ): + fail_rebuild_once = False + raise RuntimeError("rebuild exploded") + await original_build_context(current_ctx) + + mock_llm = MagicMock() + mock_llm.stream = mock_stream + plugin._mock_llm = mock_llm + pipeline._stage_build_context = build_context_fail_after_tool_call + fail_rebuild_once = True + + with pytest.raises(PipelineError, match="rebuild exploded"): + await pipeline.run_turn(ctx) + + assert ctx.runtime_message_cursor.sequence == 0 + + await pipeline.run_turn(ctx) + + second_system_content = "\n".join( + str(message.get("content", "")) + for message in captured_messages[-1] + if message.get("role") == "system" + ) + assert ( + "user_steer msg-before-tool-execution: Keep this steer after failure" + in second_system_content + ) + assert ctx.runtime_message_cursor.sequence == 1 + @pytest.mark.asyncio async def test_run_turn_allows_later_plugin_to_handle_unknown_tool(self): registry = PluginRegistry()