From 70823294abc4a17c2725eb19098fa7cc8a450aa6 Mon Sep 17 00:00:00 2001 From: Declan Brady Date: Thu, 16 Jul 2026 11:09:10 -0400 Subject: [PATCH 1/5] feat(interrupt): task/interrupt hook, protocol, resume-safe session capture (AGX1-391) Hand-editable SDK support for the platform's non-terminal task/interrupt verb (the generated client.tasks.interrupt + param/response types + RPC unions regenerate from the scale-agentex OpenAPI change via Stainless; see the TODO in protocol/acp.py). - protocol/acp.py: RPCMethod.TASK_INTERRUPT, InterruptTaskParams, PARAMS_MODEL_BY_METHOD - Temporal ACP layer: SignalName.INTERRUPT_TURN and an overridable BaseWorkflow.on_interrupt hook (no-op default so existing agents are unaffected), TemporalTaskService.interrupt() that signals interrupt_turn (never cancels/terminates), on_task_interrupt ACP-server decorator + handle_interrupt forwarding - resume fix: capture session_id from the early system/init (claude) and thread.started (codex) envelopes with a cancellation-safe finally, so a turn interrupted before completion stays resumable - unit tests Part of the Stop-mid-stream + queue-at-boundary effort (AGX1-391). Pairs with scaleapi/scale-agentex#365. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../lib/adk/_modules/_claude_code_sync.py | 51 ++++- .../lib/adk/_modules/_claude_code_turn.py | 26 ++- src/agentex/lib/adk/_modules/_codex_sync.py | 37 ++++ src/agentex/lib/adk/_modules/_codex_turn.py | 22 ++- .../services/temporal_task_service.py | 19 +- .../lib/core/temporal/types/workflow.py | 5 + .../lib/core/temporal/workflows/workflow.py | 31 ++- .../lib/sdk/fastacp/base/base_acp_server.py | 15 ++ .../lib/sdk/fastacp/impl/temporal_acp.py | 19 ++ src/agentex/lib/types/acp.py | 1 + src/agentex/protocol/acp.py | 33 ++++ tests/lib/adk/test_claude_code_turn.py | 66 +++++++ tests/lib/adk/test_codex_turn.py | 57 ++++++ tests/test_acp_interrupt.py | 178 ++++++++++++++++++ 14 files changed, 545 insertions(+), 15 deletions(-) create mode 100644 tests/test_acp_interrupt.py diff --git a/src/agentex/lib/adk/_modules/_claude_code_sync.py b/src/agentex/lib/adk/_modules/_claude_code_sync.py index 93a639118..efb05ef89 100644 --- a/src/agentex/lib/adk/_modules/_claude_code_sync.py +++ b/src/agentex/lib/adk/_modules/_claude_code_sync.py @@ -74,6 +74,33 @@ def _extract_summary(text: str, max_len: int = 300) -> str: async def convert_claude_code_to_agentex_events( lines: AsyncIterator[str | dict[str, Any]], on_result: Callable[[dict[str, Any]], Awaitable[None]] | None = None, + on_init: Callable[[dict[str, Any]], Awaitable[None]] | None = None, +) -> AsyncIterator[StreamTaskMessageStart | StreamTaskMessageDelta | StreamTaskMessageFull | StreamTaskMessageDone]: + """Public tap: convert a claude-code ``stream-json`` line stream to events. + + Thin wrapper over :func:`_convert_claude_code_impl` that owns the + cancellation backstop: a ``finally`` closes the underlying ``lines`` iterator + (when it exposes ``aclose``) whenever this generator is closed — including on + the ``GeneratorExit``/``CancelledError`` raised when a consuming task is + cancelled mid-turn by an interrupt. This terminates the CLI stdout handle / + subprocess instead of leaking it. Kept as a wrapper (rather than a + ``try/finally`` inside the impl) so the large parser body stays untouched. + """ + inner = _convert_claude_code_impl(lines, on_result=on_result, on_init=on_init) + try: + async for event in inner: + yield event + finally: + await inner.aclose() + aclose = getattr(lines, "aclose", None) + if aclose is not None: + await aclose() + + +async def _convert_claude_code_impl( + lines: AsyncIterator[str | dict[str, Any]], + on_result: Callable[[dict[str, Any]], Awaitable[None]] | None = None, + on_init: Callable[[dict[str, Any]], Awaitable[None]] | None = None, ) -> AsyncIterator[StreamTaskMessageStart | StreamTaskMessageDelta | StreamTaskMessageFull | StreamTaskMessageDone]: """Convert a claude-code ``stream-json`` line stream into Agentex ``StreamTaskMessage*`` events. @@ -85,7 +112,20 @@ async def convert_claude_code_to_agentex_events( caller can capture usage and cost. It is awaited before the generator continues. When ``None``, the result envelope is silently dropped. + ``on_init`` is called with the ``system``/``init`` envelope the moment it + arrives (the FIRST envelope of a claude-code stream), so the caller can + capture ``session_id`` EARLY — before the turn completes. This is what makes + an interrupted-before-completion turn resumable: the terminal ``result`` + envelope (which ``on_result`` reads) never arrives when a turn is cut short, + so relying on it alone loses the session id. Mirrors how the inline + ``claude_agents`` activity captures session_id from its SystemMessage init. + When ``None``, the init envelope's session metadata is not surfaced early. + Envelope → canonical mapping is documented in this module's docstring. + + A ``finally`` closes the underlying ``lines`` iterator (when it exposes + ``aclose``) as a backstop, so a cancellation mid-turn (e.g. an interrupt that + cancels the consuming task) does not leak the CLI stdout handle / subprocess. """ next_index = 0 tool_call_count = 0 @@ -355,9 +395,14 @@ async def convert_claude_code_to_agentex_events( # system / init — session metadata (ignored at this layer) # ----------------------------------------------------------------------- elif evt_type == "system": - # Session ID tracking and MCP status logging are provider concerns. - # This pure parser layer intentionally emits nothing for system events. - pass + # Session ID tracking and MCP status logging are provider concerns: + # this pure parser layer emits no StreamTaskMessage for system events. + # It DOES surface the init envelope's session metadata early via + # on_init so a caller can capture session_id before the turn's + # terminal ``result`` arrives — required for resuming a turn that was + # interrupted before completion (no ``result`` is ever emitted then). + if on_init is not None and evt.get("subtype") == "init": + await on_init(evt) # ----------------------------------------------------------------------- # result — carries usage + cost; fired to on_result, not emitted as msgs diff --git a/src/agentex/lib/adk/_modules/_claude_code_turn.py b/src/agentex/lib/adk/_modules/_claude_code_turn.py index 6c052976a..d41afafe7 100644 --- a/src/agentex/lib/adk/_modules/_claude_code_turn.py +++ b/src/agentex/lib/adk/_modules/_claude_code_turn.py @@ -118,17 +118,28 @@ class ClaudeCodeTurn: def __init__(self, lines: AsyncIterator[str | dict[str, Any]]) -> None: self._lines = lines self._result_envelope: dict[str, Any] | None = None + # session_id captured from the early system/init envelope. This is what + # keeps an interrupted-before-completion turn resumable: the terminal + # ``result`` envelope never arrives when a turn is cut short, so we must + # capture the session id up front (from init) rather than only at the end. + self._init_session_id: str | None = None self._events_stream: AsyncIterator[StreamTaskMessage] | None = None async def _on_result(self, envelope: dict[str, Any]) -> None: self._result_envelope = envelope + async def _on_init(self, envelope: dict[str, Any]) -> None: + sid = envelope.get("session_id") + if sid: + self._init_session_id = sid + @property def events(self) -> AsyncIterator[StreamTaskMessage]: if self._events_stream is None: self._events_stream = convert_claude_code_to_agentex_events( self._lines, on_result=self._on_result, + on_init=self._on_init, ) return self._events_stream @@ -136,13 +147,16 @@ def events(self) -> AsyncIterator[StreamTaskMessage]: def session_id(self) -> str | None: """The Claude Code session id, for resuming a multi-turn session. - Valid only after ``events`` has been fully consumed (populated by the - ``result`` envelope). Returns ``None`` if the stream was truncated or - Claude Code reported no session id. + Prefers the id from the terminal ``result`` envelope (fully-complete + turn), then falls back to the id captured from the early ``system/init`` + envelope. The init fallback is what makes a turn that was interrupted + before completion still resumable (no ``result`` is emitted then). + Returns ``None`` only if neither envelope was seen (e.g. the stream was + truncated before init) or Claude Code reported no session id. """ - if not self._result_envelope: - return None - return self._result_envelope.get("session_id") + if self._result_envelope and self._result_envelope.get("session_id"): + return self._result_envelope.get("session_id") + return self._init_session_id def usage(self) -> TurnUsage: """Return normalised usage for this turn. diff --git a/src/agentex/lib/adk/_modules/_codex_sync.py b/src/agentex/lib/adk/_modules/_codex_sync.py index 5a951d57e..321974194 100644 --- a/src/agentex/lib/adk/_modules/_codex_sync.py +++ b/src/agentex/lib/adk/_modules/_codex_sync.py @@ -527,6 +527,32 @@ def _handle_item(self, evt_type: str, item: dict[str, Any]) -> list[StreamTaskMe async def convert_codex_to_agentex_events( events: AsyncIterator[str | dict[str, Any]], on_result: Callable[[dict[str, Any]], None] | None = None, + on_init: Callable[[dict[str, Any]], None] | None = None, +) -> AsyncIterator[StreamTaskMessage]: + """Public tap: convert a ``codex exec --json`` event stream to events. + + Thin wrapper over :func:`_convert_codex_impl` that owns the cancellation + backstop: a ``finally`` closes the underlying ``events`` iterator (when it + exposes ``aclose``) whenever this generator is closed — including on the + ``GeneratorExit``/``CancelledError`` raised when a consuming task is + cancelled mid-turn by an interrupt. This terminates the CLI stdout handle / + subprocess instead of leaking it. + """ + inner = _convert_codex_impl(events, on_result=on_result, on_init=on_init) + try: + async for event in inner: + yield event + finally: + await inner.aclose() + aclose = getattr(events, "aclose", None) + if aclose is not None: + await aclose() + + +async def _convert_codex_impl( + events: AsyncIterator[str | dict[str, Any]], + on_result: Callable[[dict[str, Any]], None] | None = None, + on_init: Callable[[dict[str, Any]], None] | None = None, ) -> AsyncIterator[StreamTaskMessage]: """Convert a ``codex exec --json`` event stream into Agentex stream events. @@ -546,6 +572,12 @@ async def convert_codex_to_agentex_events( ``reasoning_count`` — int Use this to record turn-level metrics / usage in the caller's span without coupling this module to span/tracing APIs. + on_init: Optional callback invoked once when the ``thread.started`` event + is seen (the first event of a codex stream). Receives ``{"session_id": + }``. This surfaces the session id EARLY — before + the turn completes — so a turn interrupted before completion is still + resumable (``turn.completed`` / ``on_result`` never fires then). The + codex counterpart of the claude-code ``system/init`` early capture. Yields: Canonical ``StreamTaskMessage*`` events (Start/Delta/Full/Done) with @@ -590,6 +622,11 @@ async def convert_codex_to_agentex_events( for msg in messages: yield msg + # Surface session_id early (processor sets it while handling + # thread.started) so an interrupted turn is still resumable. + if on_init is not None and evt.get("type") == "thread.started": + on_init({"session_id": processor.session_id}) + if on_result is not None: on_result( { diff --git a/src/agentex/lib/adk/_modules/_codex_turn.py b/src/agentex/lib/adk/_modules/_codex_turn.py index e7fa1d929..05429bce5 100644 --- a/src/agentex/lib/adk/_modules/_codex_turn.py +++ b/src/agentex/lib/adk/_modules/_codex_turn.py @@ -160,6 +160,10 @@ def __init__( # Populated by the on_result callback once the stream is exhausted. self._result: dict[str, Any] | None = None + # Populated by the on_init callback when thread.started arrives (early), + # so session_id is available even if the turn is interrupted before + # completion (turn.completed / on_result never fires then). + self._init_session_id: str | None = None # The events generator is created at most once: ``_raw_events`` is a # single-consumption AsyncIterator, so re-wrapping it would yield an # already-exhausted stream that fires on_result with zeros and clobbers @@ -179,21 +183,31 @@ def events(self) -> AsyncIterator[StreamTaskMessage]: self._events_gen = convert_codex_to_agentex_events( self._raw_events, on_result=self._on_result, + on_init=self._on_init, ) return self._events_gen def _on_result(self, result: dict[str, Any]) -> None: self._result = result + def _on_init(self, init: dict[str, Any]) -> None: + sid = init.get("session_id") + if sid: + self._init_session_id = sid + @property def session_id(self) -> str | None: """The codex session id, for resuming a multi-turn session. - Valid only after ``events`` has been fully consumed (populated by the - ``on_result`` callback). Returns ``None`` if the stream is not yet - exhausted or codex reported no session id. + Prefers the id from the terminal ``turn.completed`` result (fully-complete + turn), then falls back to the id captured early from ``thread.started``. + The early fallback keeps a turn interrupted before completion resumable + (no ``turn.completed`` / ``on_result`` fires then). Returns ``None`` only + if neither was seen or codex reported no session id. """ - return self._result.get("session_id") if self._result else None + if self._result and self._result.get("session_id"): + return self._result.get("session_id") + return self._init_session_id def usage(self) -> TurnUsage: """Return normalized ``TurnUsage`` for this turn. diff --git a/src/agentex/lib/core/temporal/services/temporal_task_service.py b/src/agentex/lib/core/temporal/services/temporal_task_service.py index cb7a181d6..5f6c0c381 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -6,7 +6,7 @@ from agentex.types.task import Task from agentex.types.agent import Agent from agentex.types.event import Event -from agentex.protocol.acp import SendEventParams, CreateTaskParams +from agentex.protocol.acp import SendEventParams, CreateTaskParams, InterruptTaskParams from agentex.lib.environment_variables import EnvironmentVariables from agentex.lib.core.clients.temporal.types import WorkflowState from agentex.lib.core.temporal.types.workflow import SignalName @@ -74,6 +74,23 @@ async def send_event(self, agent: Agent, task: Task, event: Event, request: dict ).model_dump(), ) + async def interrupt(self, agent: Agent, task: Task, request: dict | None = None) -> None: + """Forward a task/interrupt to the running workflow as a dedicated signal. + + Non-terminal: unlike ``cancel``/``terminate`` this does NOT tear down the + workflow. It signals ``interrupt_turn`` so the workflow's ``on_interrupt`` + hook can stop the in-flight turn while leaving the task continuable. + """ + return await self._temporal_client.send_signal( + workflow_id=task.id, + signal=SignalName.INTERRUPT_TURN.value, + payload=InterruptTaskParams( + agent=agent, + task=task, + request=request, + ).model_dump(), + ) + async def cancel(self, task_id: str) -> None: return await self._temporal_client.cancel_workflow( workflow_id=task_id, diff --git a/src/agentex/lib/core/temporal/types/workflow.py b/src/agentex/lib/core/temporal/types/workflow.py index 973bb52c2..62624832e 100644 --- a/src/agentex/lib/core/temporal/types/workflow.py +++ b/src/agentex/lib/core/temporal/types/workflow.py @@ -3,3 +3,8 @@ class SignalName(str, Enum): RECEIVE_EVENT = "receive_event" + # Dedicated non-terminal "stop the current turn" signal (design doc section 7). + # Routed to the overridable BaseWorkflow.on_interrupt hook. Kept separate from + # RECEIVE_EVENT so the interrupt handler can interleave with (and cancel) the + # in-flight turn WITHOUT waiting on the turn lock the running turn holds. + INTERRUPT_TURN = "interrupt_turn" diff --git a/src/agentex/lib/core/temporal/workflows/workflow.py b/src/agentex/lib/core/temporal/workflows/workflow.py index a353bea6e..e47fd9a5c 100644 --- a/src/agentex/lib/core/temporal/workflows/workflow.py +++ b/src/agentex/lib/core/temporal/workflows/workflow.py @@ -6,7 +6,7 @@ from temporalio import workflow -from agentex.protocol.acp import SendEventParams, CreateTaskParams +from agentex.protocol.acp import SendEventParams, CreateTaskParams, InterruptTaskParams from agentex.lib.utils.logging import make_logger from agentex.lib.core.temporal.types.workflow import SignalName @@ -160,3 +160,32 @@ def is_continued_run(self) -> bool: gated here runs again on each history rollover. """ return workflow.info().continued_run_id is not None + + @workflow.signal(name=SignalName.INTERRUPT_TURN) + async def on_interrupt(self, params: InterruptTaskParams) -> None: + """Handle a task/interrupt: stop the in-flight turn, keep the task continuable. + + This is the durable transport for the platform's ``task/interrupt`` verb + (design doc section 7). The control plane forwards ``task/interrupt`` to the + agent pod over HTTP JSON-RPC; the ACP/Temporal layer routes it to the + ``interrupt_turn`` signal, which invokes this hook. + + Temporal runs signal handlers on the same deterministic event loop as the + workflow ``run`` method, interleaving at ``await`` points. This handler runs + CONCURRENTLY with the in-flight turn coroutine, so it MUST NOT acquire the + turn lock (the running turn already holds it) — doing so would deadlock: the + interrupt could never fire because it would be waiting on the very turn it is + meant to stop. + + The base implementation is a no-op so existing agents keep working (and the + signal is still accepted, avoiding "unhandled signal" warnings). Interruptible + agents (the golden agent is the reference) override this to actually stop the + in-flight model / CLI subprocess — e.g. SIGINT the leased sandbox CLI by pid, + or cancel the inline model ``asyncio.Task`` — while preserving partial output + and resume state so the next turn continues the conversation. + """ + logger.info( + "on_interrupt (no-op base impl) for task %s; override to make this agent " + "interruptible", + params.task.id, + ) diff --git a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py index b3dac487e..1ea8e82e6 100644 --- a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py +++ b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py @@ -22,6 +22,7 @@ CancelTaskParams, CreateTaskParams, SendMessageParams, + InterruptTaskParams, ) from agentex.lib.utils.logging import make_logger, ctx_var_request_id from agentex.protocol.json_rpc import JSONRPCError, JSONRPCRequest, JSONRPCResponse @@ -68,6 +69,7 @@ class BaseACPServer(FastAPI): Available methods: - event/send → Send a message to a task - task/cancel → Cancel a task + - task/interrupt → Interrupt the in-flight turn (non-terminal; task stays continuable) - task/approve → Approve a task """ @@ -324,6 +326,7 @@ async def _process_request( - on_task_create - on_task_event_send - on_task_cancel + - on_task_interrupt ACP Type: Sync Decorators: @@ -362,6 +365,18 @@ def on_task_cancel(self, fn: Callable[[CancelTaskParams], Awaitable[Any]]): self._handlers[RPCMethod.TASK_CANCEL] = wrapped return fn + # Type: Async + def on_task_interrupt(self, fn: Callable[[InterruptTaskParams], Awaitable[Any]]): + """Handle task/interrupt method. + + Non-terminal counterpart to ``on_task_cancel``: forwards the interrupt to + the agent so it can stop the in-flight turn while leaving the task + continuable. See the interrupt-and-queue design doc, section 7. + """ + wrapped = self._wrap_handler(fn) + self._handlers[RPCMethod.TASK_INTERRUPT] = wrapped + return fn + # Type: Sync def on_message_send( self, diff --git a/src/agentex/lib/sdk/fastacp/impl/temporal_acp.py b/src/agentex/lib/sdk/fastacp/impl/temporal_acp.py index 69d843720..1a9cce7a8 100644 --- a/src/agentex/lib/sdk/fastacp/impl/temporal_acp.py +++ b/src/agentex/lib/sdk/fastacp/impl/temporal_acp.py @@ -10,6 +10,7 @@ SendEventParams, CancelTaskParams, CreateTaskParams, + InterruptTaskParams, ) from agentex.lib.utils.logging import make_logger from agentex.lib.environment_variables import EnvironmentVariables @@ -132,3 +133,21 @@ async def handle_cancel(params: CancelTaskParams) -> None: except Exception as e: logger.error(f"Failed to cancel task: {e}") raise + + @self.on_task_interrupt + async def handle_interrupt(params: InterruptTaskParams) -> None: + """Forward task/interrupt to the running workflow via TaskService. + + Non-terminal: signals the workflow's ``interrupt_turn`` handler rather + than tearing the workflow down, so the task stays continuable. + """ + try: + if self._temporal_task_service is not None: + await self._temporal_task_service.interrupt( + agent=params.agent, + task=params.task, + request=params.request, + ) + except Exception as e: + logger.error(f"Failed to interrupt task: {e}") + raise diff --git a/src/agentex/lib/types/acp.py b/src/agentex/lib/types/acp.py index 74295ef88..e86ff7ddf 100644 --- a/src/agentex/lib/types/acp.py +++ b/src/agentex/lib/types/acp.py @@ -12,4 +12,5 @@ CancelTaskParams, CreateTaskParams, SendMessageParams, + InterruptTaskParams, ) diff --git a/src/agentex/protocol/acp.py b/src/agentex/protocol/acp.py index d719b4fd5..8ee469f21 100644 --- a/src/agentex/protocol/acp.py +++ b/src/agentex/protocol/acp.py @@ -18,6 +18,7 @@ class RPCMethod(str, Enum): MESSAGE_SEND = "message/send" TASK_CANCEL = "task/cancel" TASK_CREATE = "task/create" + TASK_INTERRUPT = "task/interrupt" class CreateTaskParams(BaseModel): @@ -104,6 +105,28 @@ class CancelTaskParams(BaseModel): ) +class InterruptTaskParams(BaseModel): + """Parameters for task/interrupt method. + + Non-terminal counterpart to :class:`CancelTaskParams`. The control plane + forwards ``task/interrupt`` to the agent so it can stop the in-flight turn + while leaving the task continuable (status ``INTERRUPTED``, not a terminal + status). See the interrupt-and-queue design doc, sections 5-7. + + Attributes: + agent: The agent that the task was sent to. + task: The task that was interrupted. + request: Additional request context including headers forwarded to this agent. + """ + + agent: Agent = Field(..., description="The agent that the task was sent to") + task: Task = Field(..., description="The task that was interrupted") + request: dict[str, Any] | None = Field( + default=None, + description="Additional request context including headers forwarded to this agent", + ) + + RPC_SYNC_METHODS = [ RPCMethod.MESSAGE_SEND, ] @@ -113,4 +136,14 @@ class CancelTaskParams(BaseModel): RPCMethod.TASK_CANCEL: CancelTaskParams, RPCMethod.MESSAGE_SEND: SendMessageParams, RPCMethod.TASK_CREATE: CreateTaskParams, + RPCMethod.TASK_INTERRUPT: InterruptTaskParams, } + +# TODO(interrupt): the client-facing REST method ``client.tasks.interrupt(task_id, +# reason=...)`` plus its param/response types (``task_interrupt_params.py``) and the +# ``AgentRPCMethod`` union additions are Stainless-GENERATED. They regenerate from the +# upstream OpenAPI change in scale-agentex (design doc sections 9.1 / 9.2) once +# ``TASK_INTERRUPT`` and the ``POST /tasks/{task_id}/interrupt`` route land there. Do NOT +# hand-write them under ``src/agentex/resources/**`` or ``src/agentex/types/**``. The +# hand-editable protocol shapes above (and the agent-side hook in ``agentex.lib.**``) +# survive regeneration and are what agents rely on today. diff --git a/tests/lib/adk/test_claude_code_turn.py b/tests/lib/adk/test_claude_code_turn.py index 4fbb2f913..fac92be0a 100644 --- a/tests/lib/adk/test_claude_code_turn.py +++ b/tests/lib/adk/test_claude_code_turn.py @@ -281,3 +281,69 @@ async def test_events_property_returns_same_iterator(self): it1 = turn.events it2 = turn.events assert it1 is it2 + + +# --------------------------------------------------------------------------- +# Early session_id capture (resume after interrupt) +# --------------------------------------------------------------------------- + + +class TestClaudeCodeTurnSessionIdCapture: + async def test_session_id_from_result_when_turn_completes(self): + envelopes = [ + {"type": "system", "subtype": "init", "session_id": "sess-init"}, + {"type": "assistant", "message": {"content": [{"type": "text", "text": "hi"}]}}, + {"type": "result", "session_id": "sess-final", "usage": {}}, + ] + turn = ClaudeCodeTurn(_aiter(envelopes)) + await _drain(turn) + # Terminal result wins for a fully-completed turn. + assert turn.session_id == "sess-final" + + async def test_session_id_from_init_when_interrupted_before_result(self): + """No `result` envelope (turn cut short) must still yield a session_id. + + This is the resume fix: capture session_id from the early system/init + envelope so an interrupted-before-completion turn stays resumable. + """ + envelopes = [ + {"type": "system", "subtype": "init", "session_id": "sess-init"}, + {"type": "assistant", "message": {"content": [{"type": "text", "text": "partial"}]}}, + # stream ends here (interrupted) — no result envelope + ] + turn = ClaudeCodeTurn(_aiter(envelopes)) + await _drain(turn) + assert turn.session_id == "sess-init" + + async def test_session_id_none_when_no_init_or_result(self): + envelopes = [ + {"type": "assistant", "message": {"content": [{"type": "text", "text": "x"}]}}, + ] + turn = ClaudeCodeTurn(_aiter(envelopes)) + await _drain(turn) + assert turn.session_id is None + + async def test_guarded_lines_closes_source_on_early_break(self): + """The converter's finally-backstop closes the source stdout iterator. + + Breaking out of the event loop early (as a cancellation/interrupt would) + must trigger aclose() on the underlying lines iterator so the CLI stdout + handle is not leaked. + """ + closed = {"value": False} + + async def _lines(): + try: + yield {"type": "system", "subtype": "init", "session_id": "sess-init"} + yield {"type": "assistant", "message": {"content": [{"type": "text", "text": "a"}]}} + yield {"type": "assistant", "message": {"content": [{"type": "text", "text": "b"}]}} + finally: + closed["value"] = True + + turn = ClaudeCodeTurn(_lines()) + events = turn.events + # Consume only the first event, then close the stream early. + async for _ in events: + break + await events.aclose() + assert closed["value"] is True diff --git a/tests/lib/adk/test_codex_turn.py b/tests/lib/adk/test_codex_turn.py index f6a046478..e3bd8c988 100644 --- a/tests/lib/adk/test_codex_turn.py +++ b/tests/lib/adk/test_codex_turn.py @@ -280,3 +280,60 @@ async def test_reasoning_tokens_propagated(self) -> None: turn = CodexTurn(_aiter(events), model="o4-mini") await _collect(turn) assert turn.usage().reasoning_tokens == 40 + + +# --------------------------------------------------------------------------- +# Early session_id capture (resume after interrupt) +# --------------------------------------------------------------------------- + + +class TestCodexTurnSessionIdCapture: + async def test_session_id_from_result_when_turn_completes(self) -> None: + events = [ + {"type": "thread.started", "thread_id": "thread-abc"}, + {"type": "turn.completed", "usage": None}, + ] + turn = CodexTurn(_aiter(events), model="o4-mini") + await _collect(turn) + # on_result carries session_id from processor.session_id at end of stream. + assert turn.session_id == "thread-abc" + + async def test_session_id_from_init_when_interrupted_before_completion(self) -> None: + """thread.started must yield session_id even with no turn.completed. + + Mirrors the claude-code early-capture fix: an interrupted-before-completion + codex turn never emits turn.completed, so the early thread.started capture + is what keeps it resumable. + """ + events = [ + {"type": "thread.started", "thread_id": "thread-abc"}, + # stream ends here (interrupted) — no turn.completed + ] + turn = CodexTurn(_aiter(events), model="o4-mini") + await _collect(turn) + assert turn.session_id == "thread-abc" + + async def test_session_id_none_without_thread_started(self) -> None: + turn = CodexTurn(_aiter([{"type": "turn.completed", "usage": None}]), model="o4-mini") + await _collect(turn) + assert turn.session_id is None + + async def test_guarded_events_closes_source_on_early_break(self) -> None: + closed = {"value": False} + + async def _events(): + try: + yield {"type": "thread.started", "thread_id": "thread-abc"} + yield { + "type": "item.completed", + "item": {"id": "m1", "type": "agent_message", "text": "hi"}, + } + finally: + closed["value"] = True + + turn = CodexTurn(_events(), model="o4-mini") + gen = turn.events + async for _ in gen: + break + await gen.aclose() + assert closed["value"] is True diff --git a/tests/test_acp_interrupt.py b/tests/test_acp_interrupt.py new file mode 100644 index 000000000..f53cfa6b4 --- /dev/null +++ b/tests/test_acp_interrupt.py @@ -0,0 +1,178 @@ +"""Unit tests for the hand-editable task/interrupt additions. + +Covers the regeneration-safe surfaces added for the interrupt-and-queue design +(design doc sections 6, 7, 9.2): + +1. Protocol (``agentex.protocol.acp``): ``RPCMethod.TASK_INTERRUPT``, the + ``InterruptTaskParams`` model (mirror of ``CancelTaskParams``), and the + ``PARAMS_MODEL_BY_METHOD`` entry, plus the back-compat shim re-export. +2. ACP server routing: ``BaseACPServer.on_task_interrupt`` registers a handler + under ``RPCMethod.TASK_INTERRUPT``. +3. Temporal transport: ``BaseWorkflow.on_interrupt`` is a ``@workflow.signal`` + named ``interrupt_turn`` (and is NOT abstract, so existing agents keep + working), and ``TemporalTaskService.interrupt`` forwards that signal without + tearing the workflow down. +""" + +from __future__ import annotations + +from unittest.mock import Mock, AsyncMock + +import pytest + +from agentex.types.task import Task +from agentex.types.agent import Agent +from agentex.protocol.acp import ( + PARAMS_MODEL_BY_METHOD, + RPCMethod, + CancelTaskParams, + InterruptTaskParams, +) + + +def _agent() -> Agent: + return Agent( + id="test-agent-456", + name="test-agent", + description="test-agent", + acp_type="async", + created_at="2023-01-01T00:00:00Z", + updated_at="2023-01-01T00:00:00Z", + ) + + +def _task() -> Task: + return Task(id="test-task-123", status="RUNNING") + + +# --------------------------------------------------------------------------- +# Protocol additions +# --------------------------------------------------------------------------- + + +class TestInterruptProtocol: + def test_rpc_method_value(self) -> None: + assert RPCMethod.TASK_INTERRUPT.value == "task/interrupt" + # Constructible from the wire string (the ACP server does RPCMethod(str)). + assert RPCMethod("task/interrupt") is RPCMethod.TASK_INTERRUPT + + def test_params_model_registered(self) -> None: + assert PARAMS_MODEL_BY_METHOD[RPCMethod.TASK_INTERRUPT] is InterruptTaskParams + + def test_params_mirror_cancel_shape(self) -> None: + """InterruptTaskParams mirrors CancelTaskParams field-for-field.""" + assert set(InterruptTaskParams.model_fields) == set(CancelTaskParams.model_fields) + assert set(InterruptTaskParams.model_fields) == {"agent", "task", "request"} + + def test_params_validate_round_trip(self) -> None: + params = InterruptTaskParams(agent=_agent(), task=_task()) + assert params.task.id == "test-task-123" + assert params.request is None + # Header forwarding path (BaseACPServer populates params.request). + with_headers = InterruptTaskParams.model_validate( + { + "agent": _agent().model_dump(), + "task": _task().model_dump(), + "request": {"headers": {"x-foo": "bar"}}, + } + ) + assert with_headers.request == {"headers": {"x-foo": "bar"}} + + def test_shim_reexports_interrupt_params(self) -> None: + """The back-compat shim must re-export the new model as the same object.""" + from agentex.protocol import acp as canon + from agentex.lib.types import acp as shim + + assert shim.InterruptTaskParams is canon.InterruptTaskParams + + +# --------------------------------------------------------------------------- +# ACP server routing +# --------------------------------------------------------------------------- + + +class TestACPServerInterruptRouting: + def test_on_task_interrupt_registers_handler(self) -> None: + from unittest.mock import patch + + from agentex.lib.sdk.fastacp.base.base_acp_server import BaseACPServer + + with patch.dict("os.environ", {"AGENTEX_BASE_URL": ""}): + server = BaseACPServer() + + assert RPCMethod.TASK_INTERRUPT not in server._handlers + + @server.on_task_interrupt + async def _handle(params: InterruptTaskParams) -> None: # noqa: ARG001 + return None + + assert RPCMethod.TASK_INTERRUPT in server._handlers + assert server._handlers[RPCMethod.TASK_INTERRUPT] is not None + + def test_temporal_acp_wires_interrupt_handler(self) -> None: + from unittest.mock import patch + + from agentex.lib.sdk.fastacp.impl.temporal_acp import TemporalACP + + with patch.dict("os.environ", {"AGENTEX_BASE_URL": ""}): + server = TemporalACP.create(temporal_address="localhost:7233") + + assert RPCMethod.TASK_INTERRUPT in server._handlers + + +# --------------------------------------------------------------------------- +# Temporal transport: signal + service forwarding +# --------------------------------------------------------------------------- + + +class TestWorkflowInterruptSignal: + def test_on_interrupt_is_signal_named_interrupt_turn(self) -> None: + from agentex.lib.core.temporal.types.workflow import SignalName + from agentex.lib.core.temporal.workflows.workflow import BaseWorkflow + + # getattr avoids the dunder name-mangling that would otherwise rewrite + # this to _TestWorkflowInterruptSignal__temporal_signal_definition. + sd = getattr(BaseWorkflow.on_interrupt, "__temporal_signal_definition") + assert sd is not None + # str-enum: equal by value to the wire string "interrupt_turn". + assert sd.name == SignalName.INTERRUPT_TURN + assert sd.name == "interrupt_turn" + + def test_on_interrupt_not_abstract(self) -> None: + """A default no-op keeps existing (non-interruptible) workflows valid.""" + from agentex.lib.core.temporal.workflows.workflow import BaseWorkflow + + assert "on_interrupt" not in BaseWorkflow.__abstractmethods__ + + +class TestTemporalTaskServiceInterrupt: + async def test_interrupt_sends_signal_not_cancel(self) -> None: + from agentex.lib.core.temporal.types.workflow import SignalName + from agentex.lib.core.temporal.services.temporal_task_service import ( + TemporalTaskService, + ) + + temporal_client = Mock() + temporal_client.send_signal = AsyncMock() + temporal_client.cancel_workflow = AsyncMock() + temporal_client.terminate_workflow = AsyncMock() + + service = TemporalTaskService(temporal_client=temporal_client, env_vars=Mock()) + + await service.interrupt(agent=_agent(), task=_task(), request={"headers": {"x-a": "b"}}) + + # Non-terminal: it signals, it does NOT cancel or terminate the workflow. + temporal_client.cancel_workflow.assert_not_called() + temporal_client.terminate_workflow.assert_not_called() + temporal_client.send_signal.assert_awaited_once() + + kwargs = temporal_client.send_signal.await_args.kwargs + assert kwargs["workflow_id"] == "test-task-123" + assert kwargs["signal"] == SignalName.INTERRUPT_TURN.value == "interrupt_turn" + # Payload is a serialized InterruptTaskParams (task/agent/request). + assert kwargs["payload"]["task"]["id"] == "test-task-123" + assert kwargs["payload"]["request"] == {"headers": {"x-a": "b"}} + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(pytest.main([__file__, "-v"])) From 96bfc571d08e944eac13f6cbe3be11123bc676d6 Mon Sep 17 00:00:00 2001 From: Declan Brady Date: Thu, 16 Jul 2026 14:11:09 -0400 Subject: [PATCH 2/5] feat(adk): add adk.tasks.interrupt for agents to record INTERRUPTED (AGX1-391) Agents call adk.tasks.interrupt(task_id=...) from their interrupt handler, after stopping the in-flight turn, to record the non-terminal INTERRUPTED status (the control plane forwards task/interrupt but no longer writes status). Mirrors adk.tasks.cancel across the facade / activity / service layers and registers the interrupt-task activity. The service uses a generic POST /tasks/{id}/interrupt for now; switches to client.tasks.interrupt once Stainless regenerates the typed method. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/agentex/lib/adk/_modules/tasks.py | 47 +++++++++++++++++++ src/agentex/lib/core/services/adk/tasks.py | 26 ++++++++++ .../lib/core/temporal/activities/__init__.py | 1 + .../activities/adk/tasks_activities.py | 10 ++++ 4 files changed, 84 insertions(+) diff --git a/src/agentex/lib/adk/_modules/tasks.py b/src/agentex/lib/adk/_modules/tasks.py index 7b304a656..d842f23ee 100644 --- a/src/agentex/lib/adk/_modules/tasks.py +++ b/src/agentex/lib/adk/_modules/tasks.py @@ -174,6 +174,53 @@ async def cancel( parent_span_id=parent_span_id, ) + async def interrupt( + self, + *, + task_id: str, + reason: str | None = None, + trace_id: str | None = None, + parent_span_id: str | None = None, + start_to_close_timeout: timedelta = timedelta(seconds=5), + heartbeat_timeout: timedelta = timedelta(seconds=5), + retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY, + ) -> Task: + """ + Mark a running task as interrupted (non-terminal). + + Interrupt is cooperative: an agent calls this from its interrupt handler, + after it has actually stopped its in-flight turn, to record INTERRUPTED. + The task stays continuable; the control plane resumes it to RUNNING on the + next turn. + Args: + task_id: The ID of the task to interrupt. + reason: Optional reason for the interrupt. + Returns: + The updated task entry. + """ + params = TaskStatusTransitionParams( + task_id=task_id, + reason=reason, + trace_id=trace_id, + parent_span_id=parent_span_id, + ) + if in_temporal_workflow(): + return await ActivityHelpers.execute_activity( + activity_name=TasksActivityName.INTERRUPT_TASK, + request=params, + response_type=Task, + start_to_close_timeout=start_to_close_timeout, + retry_policy=retry_policy, + heartbeat_timeout=heartbeat_timeout, + ) + else: + return await self._tasks_service.interrupt_task( + task_id=task_id, + reason=reason, + trace_id=trace_id, + parent_span_id=parent_span_id, + ) + async def complete( self, *, diff --git a/src/agentex/lib/core/services/adk/tasks.py b/src/agentex/lib/core/services/adk/tasks.py index 7748799e4..fecd399f4 100644 --- a/src/agentex/lib/core/services/adk/tasks.py +++ b/src/agentex/lib/core/services/adk/tasks.py @@ -98,6 +98,32 @@ async def cancel_task( span.output = task_model.model_dump() return task_model + async def interrupt_task( + self, + task_id: str, + reason: str | None = None, + trace_id: str | None = None, + parent_span_id: str | None = None, + ) -> Task: + trace = self._tracer.trace(trace_id) + async with trace.span( + parent_id=parent_span_id, + name="interrupt_task", + input={"task_id": task_id, "reason": reason}, + ) as span: + heartbeat_if_in_workflow("interrupt task") + # TODO(AGX1-391): switch to self._agentex_client.tasks.interrupt(...) once the + # POST /tasks/{id}/interrupt endpoint is added to the Stainless config and + # regenerated. Using the generic post keeps this working before regeneration. + task_model = await self._agentex_client.post( + f"/tasks/{task_id}/interrupt", + cast_to=Task, + body={"reason": reason}, + ) + if span: + span.output = task_model.model_dump() + return task_model + async def complete_task( self, task_id: str, diff --git a/src/agentex/lib/core/temporal/activities/__init__.py b/src/agentex/lib/core/temporal/activities/__init__.py index 4660afdde..93e6b69e6 100644 --- a/src/agentex/lib/core/temporal/activities/__init__.py +++ b/src/agentex/lib/core/temporal/activities/__init__.py @@ -181,6 +181,7 @@ def get_all_activities(sgp_client=None): tasks_activities.get_task, tasks_activities.delete_task, tasks_activities.cancel_task, + tasks_activities.interrupt_task, tasks_activities.complete_task, tasks_activities.fail_task, tasks_activities.terminate_task, diff --git a/src/agentex/lib/core/temporal/activities/adk/tasks_activities.py b/src/agentex/lib/core/temporal/activities/adk/tasks_activities.py index 38eecd447..bd1c430c4 100644 --- a/src/agentex/lib/core/temporal/activities/adk/tasks_activities.py +++ b/src/agentex/lib/core/temporal/activities/adk/tasks_activities.py @@ -18,6 +18,7 @@ class TasksActivityName(str, Enum): GET_TASK = "get-task" DELETE_TASK = "delete-task" CANCEL_TASK = "cancel-task" + INTERRUPT_TASK = "interrupt-task" COMPLETE_TASK = "complete-task" FAIL_TASK = "fail-task" TERMINATE_TASK = "terminate-task" @@ -83,6 +84,15 @@ async def cancel_task(self, params: TaskStatusTransitionParams) -> Task: parent_span_id=params.parent_span_id, ) + @activity.defn(name=TasksActivityName.INTERRUPT_TASK) + async def interrupt_task(self, params: TaskStatusTransitionParams) -> Task: + return await self._tasks_service.interrupt_task( + task_id=params.task_id, + reason=params.reason, + trace_id=params.trace_id, + parent_span_id=params.parent_span_id, + ) + @activity.defn(name=TasksActivityName.COMPLETE_TASK) async def complete_task(self, params: TaskStatusTransitionParams) -> Task: return await self._tasks_service.complete_task( From 226789c8284c4cdb92665a06e14e412482c356d3 Mon Sep 17 00:00:00 2001 From: Declan Brady Date: Thu, 16 Jul 2026 15:35:10 -0400 Subject: [PATCH 3/5] refactor(adk): use generated tasks.interrupt now that it's in the SDK (AGX1-391) scale-agentex#365 merged and Stainless regenerated client.tasks.interrupt, so drop the interim generic POST in the adk interrupt service and call the typed method (matches cancel_task). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/agentex/lib/core/services/adk/tasks.py | 9 +-------- src/agentex/protocol/acp.py | 9 --------- 2 files changed, 1 insertion(+), 17 deletions(-) diff --git a/src/agentex/lib/core/services/adk/tasks.py b/src/agentex/lib/core/services/adk/tasks.py index fecd399f4..f1dba08bd 100644 --- a/src/agentex/lib/core/services/adk/tasks.py +++ b/src/agentex/lib/core/services/adk/tasks.py @@ -112,14 +112,7 @@ async def interrupt_task( input={"task_id": task_id, "reason": reason}, ) as span: heartbeat_if_in_workflow("interrupt task") - # TODO(AGX1-391): switch to self._agentex_client.tasks.interrupt(...) once the - # POST /tasks/{id}/interrupt endpoint is added to the Stainless config and - # regenerated. Using the generic post keeps this working before regeneration. - task_model = await self._agentex_client.post( - f"/tasks/{task_id}/interrupt", - cast_to=Task, - body={"reason": reason}, - ) + task_model = await self._agentex_client.tasks.interrupt(task_id=task_id, reason=reason) if span: span.output = task_model.model_dump() return task_model diff --git a/src/agentex/protocol/acp.py b/src/agentex/protocol/acp.py index 8ee469f21..7e310cd89 100644 --- a/src/agentex/protocol/acp.py +++ b/src/agentex/protocol/acp.py @@ -138,12 +138,3 @@ class InterruptTaskParams(BaseModel): RPCMethod.TASK_CREATE: CreateTaskParams, RPCMethod.TASK_INTERRUPT: InterruptTaskParams, } - -# TODO(interrupt): the client-facing REST method ``client.tasks.interrupt(task_id, -# reason=...)`` plus its param/response types (``task_interrupt_params.py``) and the -# ``AgentRPCMethod`` union additions are Stainless-GENERATED. They regenerate from the -# upstream OpenAPI change in scale-agentex (design doc sections 9.1 / 9.2) once -# ``TASK_INTERRUPT`` and the ``POST /tasks/{task_id}/interrupt`` route land there. Do NOT -# hand-write them under ``src/agentex/resources/**`` or ``src/agentex/types/**``. The -# hand-editable protocol shapes above (and the agent-side hook in ``agentex.lib.**``) -# survive regeneration and are what agents rely on today. From bc019620b89710420cda239a3d07afe07557a698 Mon Sep 17 00:00:00 2001 From: Declan Brady Date: Thu, 16 Jul 2026 15:43:11 -0400 Subject: [PATCH 4/5] fix(adk): make aclose calls type-safe on AsyncIterator (AGX1-391) The resume backstop and its tests called .aclose() directly on streams typed AsyncIterator (only AsyncGenerator exposes aclose), failing the pyright lint. Guard via getattr like the existing source-stream close, keeping runtime behavior identical. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/agentex/lib/adk/_modules/_claude_code_sync.py | 4 +++- src/agentex/lib/adk/_modules/_codex_sync.py | 4 +++- tests/lib/adk/test_claude_code_turn.py | 4 +++- tests/lib/adk/test_codex_turn.py | 4 +++- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/agentex/lib/adk/_modules/_claude_code_sync.py b/src/agentex/lib/adk/_modules/_claude_code_sync.py index efb05ef89..e8daa44a3 100644 --- a/src/agentex/lib/adk/_modules/_claude_code_sync.py +++ b/src/agentex/lib/adk/_modules/_claude_code_sync.py @@ -91,7 +91,9 @@ async def convert_claude_code_to_agentex_events( async for event in inner: yield event finally: - await inner.aclose() + inner_aclose = getattr(inner, "aclose", None) + if inner_aclose is not None: + await inner_aclose() aclose = getattr(lines, "aclose", None) if aclose is not None: await aclose() diff --git a/src/agentex/lib/adk/_modules/_codex_sync.py b/src/agentex/lib/adk/_modules/_codex_sync.py index 321974194..951622b13 100644 --- a/src/agentex/lib/adk/_modules/_codex_sync.py +++ b/src/agentex/lib/adk/_modules/_codex_sync.py @@ -543,7 +543,9 @@ async def convert_codex_to_agentex_events( async for event in inner: yield event finally: - await inner.aclose() + inner_aclose = getattr(inner, "aclose", None) + if inner_aclose is not None: + await inner_aclose() aclose = getattr(events, "aclose", None) if aclose is not None: await aclose() diff --git a/tests/lib/adk/test_claude_code_turn.py b/tests/lib/adk/test_claude_code_turn.py index fac92be0a..be80dfe1f 100644 --- a/tests/lib/adk/test_claude_code_turn.py +++ b/tests/lib/adk/test_claude_code_turn.py @@ -345,5 +345,7 @@ async def _lines(): # Consume only the first event, then close the stream early. async for _ in events: break - await events.aclose() + events_aclose = getattr(events, "aclose", None) + assert events_aclose is not None + await events_aclose() assert closed["value"] is True diff --git a/tests/lib/adk/test_codex_turn.py b/tests/lib/adk/test_codex_turn.py index e3bd8c988..c2843e90e 100644 --- a/tests/lib/adk/test_codex_turn.py +++ b/tests/lib/adk/test_codex_turn.py @@ -335,5 +335,7 @@ async def _events(): gen = turn.events async for _ in gen: break - await gen.aclose() + gen_aclose = getattr(gen, "aclose", None) + assert gen_aclose is not None + await gen_aclose() assert closed["value"] is True From 5e306ac8134b5a5c945b3b0295e1d3ae67a94f32 Mon Sep 17 00:00:00 2001 From: Declan Brady Date: Thu, 16 Jul 2026 15:59:01 -0400 Subject: [PATCH 5/5] fix(adk): cancellation-safe teardown for openai/langgraph/pydantic-ai taps (AGX1-391) Mirror the claude-code/codex tap hardening across the other harness converters: wrap each convert_*_to_agentex_events in a thin public wrapper with a cancellation-safe finally that closes the source stream (defensive getattr aclose), so an interrupted turn tears the source down instead of leaking it. Parser bodies moved untouched into _convert_*_impl. The early-session_id capture is NOT ported: these frameworks carry conversation state via the input list, not a resumable session id. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../lib/adk/_modules/_langgraph_sync.py | 21 +++++++++++++++++++ src/agentex/lib/adk/_modules/_openai_sync.py | 19 +++++++++++++++++ .../lib/adk/_modules/_pydantic_ai_sync.py | 21 +++++++++++++++++++ 3 files changed, 61 insertions(+) diff --git a/src/agentex/lib/adk/_modules/_langgraph_sync.py b/src/agentex/lib/adk/_modules/_langgraph_sync.py index 9d7b73847..02f2ba416 100644 --- a/src/agentex/lib/adk/_modules/_langgraph_sync.py +++ b/src/agentex/lib/adk/_modules/_langgraph_sync.py @@ -32,6 +32,27 @@ async def convert_langgraph_to_agentex_events( stream: Any, on_final_ai_message: Optional[Callable[..., None]] = None, +) -> AsyncGenerator[Any, None]: + """Public LangGraph tap: convert events, closing the source stream on exit. + + Thin wrapper over ``_convert_langgraph_impl`` that adds a cancellation-safe + ``finally`` so an interrupted turn tears down the LangGraph ``astream`` source + instead of leaking it. + """ + inner = _convert_langgraph_impl(stream, on_final_ai_message=on_final_ai_message) + try: + async for event in inner: + yield event + finally: + for _src in (inner, stream): + _aclose = getattr(_src, "aclose", None) + if _aclose is not None: + await _aclose() + + +async def _convert_langgraph_impl( + stream: Any, + on_final_ai_message: Optional[Callable[..., None]] = None, ) -> AsyncGenerator[Any, None]: """Convert LangGraph streaming events to Agentex TaskMessageUpdate events. diff --git a/src/agentex/lib/adk/_modules/_openai_sync.py b/src/agentex/lib/adk/_modules/_openai_sync.py index ac404bef1..b857a5be6 100644 --- a/src/agentex/lib/adk/_modules/_openai_sync.py +++ b/src/agentex/lib/adk/_modules/_openai_sync.py @@ -145,6 +145,25 @@ def _extract_tool_response_info(tool_map: dict[str, Any], tool_output_item: Any) async def convert_openai_to_agentex_events(stream_response): + """Public OpenAI tap: parse the event stream, closing the source on exit. + + Thin wrapper over ``_convert_openai_impl`` that adds a cancellation-safe + ``finally`` so an interrupted turn tears down the source event stream instead + of leaking it. (Resume state is carried by the OpenAI Agents SDK input list, + not a session id, so there is no early-session_id capture like the CLI taps.) + """ + inner = _convert_openai_impl(stream_response) + try: + async for event in inner: + yield event + finally: + for _src in (inner, stream_response): + _aclose = getattr(_src, "aclose", None) + if _aclose is not None: + await _aclose() + + +async def _convert_openai_impl(stream_response): """Convert OpenAI streaming events to AgentEx TaskMessageUpdate events with reasoning support. This is an enhanced version of the base converter that includes support for: diff --git a/src/agentex/lib/adk/_modules/_pydantic_ai_sync.py b/src/agentex/lib/adk/_modules/_pydantic_ai_sync.py index 0f9aaeb55..75bfb64db 100644 --- a/src/agentex/lib/adk/_modules/_pydantic_ai_sync.py +++ b/src/agentex/lib/adk/_modules/_pydantic_ai_sync.py @@ -120,6 +120,27 @@ def _tool_return_content(result: ToolReturnPart | Any) -> Any: async def convert_pydantic_ai_to_agentex_events( stream_response: AsyncIterator[Any], on_result: Callable[[AgentRunResultEvent], Any] | None = None, +) -> AsyncIterator[StreamTaskMessageStart | StreamTaskMessageDelta | StreamTaskMessageFull | StreamTaskMessageDone]: + """Public Pydantic AI tap: convert events, closing the source stream on exit. + + Thin wrapper over ``_convert_pydantic_ai_impl`` that adds a cancellation-safe + ``finally`` so an interrupted turn tears down the source stream instead of + leaking it. + """ + inner = _convert_pydantic_ai_impl(stream_response, on_result=on_result) + try: + async for event in inner: + yield event + finally: + for _src in (inner, stream_response): + _aclose = getattr(_src, "aclose", None) + if _aclose is not None: + await _aclose() + + +async def _convert_pydantic_ai_impl( + stream_response: AsyncIterator[Any], + on_result: Callable[[AgentRunResultEvent], Any] | None = None, ) -> AsyncIterator[StreamTaskMessageStart | StreamTaskMessageDelta | StreamTaskMessageFull | StreamTaskMessageDone]: """Convert a Pydantic AI agent event stream into Agentex stream events.