agents: poll tool tasks in the Python Claude Agent SDK adapter

agents/python/claude-agent-sdk/README.md

@@ -286,10 +286,14 @@ flow = claude_agent_session_flow(
   `CLAUDE_CODE_BUILT_IN_TOOL_NAMES`, or another source's tool. The same names
   listed at run time fail that `RunModel` attempt. An `allowed_tools` entry for
   a source tool fails the attempt, because the call must defer.
-- The runner polls no tool tasks. A source call that starts one records
-  `failed` for a `read` call and `outcome_unknown` otherwise. A source whose
-  listing fails or takes longer than `tool_source_listing_timeout` offers no
-  tools in that attempt.
+- A source call can start a tool task by returning `ToolTaskStarted`. The
+  session polls it on a Timer through the source's `poll_tool_task`, and the
+  resumed CLI receives the task's result. `McpToolSource` hands its Streamable
+  HTTP tasks to the session this way. `cancel_turn` records a running task's
+  call `outcome_unknown` and cancels the task through the source's
+  `cancel_tool_task`; a source without that method leaves the task running.
+- A source whose listing fails or takes longer than
+  `tool_source_listing_timeout` offers no tools in that attempt.
 - A source tool that the turn already called but the current listing lacks is
   still served, replay-only, so the resumed CLI receives its recorded result.
   A new call of that tool is denied in the turn ("its tool source does not

agents/python/claude-agent-sdk/src/dex_agents_claude_agent_sdk/runner.py

@@ -53,6 +53,9 @@
     ToolCallStepOptions,
     ToolExecutionError,
     ToolExecutionResult,
+    ToolTaskCancellation,
+    ToolTaskPoll,
+    ToolTaskPollOutcome,
     TurnCompleted,
     collect_turn_images,
     load_images,
@@ -208,6 +211,16 @@ class ClaudeAgentRunner(AgentFrameworkRunner[DeferredToolOutcome]):
     deferrable call is refused, and the model sees the earlier result before
     it calls again.
 
+    A source call can start a tool task: its ``execute_deferred_tool_call``
+    returns ``ToolTaskStarted``, and the session polls the task on a Timer
+    through ``poll_tool_task`` until the source returns its result, which the
+    resumed Claude session receives. ``McpToolSource`` hands its Streamable
+    HTTP tasks to the session this way. ``cancel_turn`` abandons a running
+    task and cancels it through ``cancel_tool_task``.
+
+    .. versionchanged:: 0.6.6
+        The runner polls and cancels tool tasks of its sources.
+
     With a ``media_store``, the runner takes images. A user message's images
     reach the model as Claude ``image`` blocks after its text, and a tool
     result's images as ``image`` blocks inside its ``tool_result``. The
@@ -294,12 +307,16 @@ def __init__(
                 when the source's policy asks, and runs in its own tool Step
                 through the source's ``execute_deferred_tool_call``, with the
                 Step options of its ``tool_call_step_options``. Policies and
-                records use the tool's own name. The runner polls no tool
-                tasks, so a call that starts one records ``FAILED`` for a
-                ``READ`` call and ``OUTCOME_UNKNOWN`` otherwise. Defaults to
-                none.
+                records use the tool's own name. A call that returns
+                ``ToolTaskStarted`` waits for its task: the session polls it
+                through the source's ``poll_tool_task`` and cancels it through
+                the source's ``cancel_tool_task`` when ``cancel_turn``
+                abandons it. Defaults to none.
 
                 .. versionadded:: 0.4.0
+
+                .. versionchanged:: 0.6.6
+                    The runner polls and cancels the sources' tool tasks.
             tool_source_listing_timeout: How long one ``RunModel`` attempt
                 waits for each source's ``list_tools``. A source that fails or
                 takes longer offers no tools in that attempt. Defaults to 10
@@ -585,6 +602,78 @@ async def execute_deferred_tool_call(self, call: DeferredToolExecution) -> Defer
             raise ToolExecutionError(text)
         return ToolExecutionResult(output=text, images=images)
 
+    async def poll_tool_task(self, poll: ToolTaskPoll) -> ToolTaskPollOutcome:
+        """Poll a tool source's task once, through the source that built the call.
+
+        A source call whose ``execute_deferred_tool_call`` returns
+        ``ToolTaskStarted`` records ``WAITING_FOR_TASK``. The session's
+        ``PollToolTask`` Step then calls this method after each poll interval,
+        on any Worker, until the source returns a result. The resumed Claude
+        session receives that result as the call's ``tool_result``.
+
+        .. versionadded:: 0.6.6
+
+        Args:
+            poll: The call's record and its idempotency key from ``PollToolTask``.
+
+        Returns:
+            The source's ``poll_tool_task`` outcome: ``ToolExecutionResult``
+            when the task finished, ``ToolTaskPending`` while it runs, or
+            ``ToolInputRequested`` when it needs the user's answer.
+
+        Raises:
+            LookupError: If no tool source built the call, or the source that
+                built it is not configured on this Worker. The task may still
+                run, so ``PollToolTask`` retries and then records the call by
+                its effect.
+            ToolExecutionError: The source reports that the task failed. The
+                session records ``FAILED``.
+            ToolOutcomeUnknownError: The source cannot tell whether the task
+                took effect. The session records ``OUTCOME_UNKNOWN``.
+            Exception: Any other exception the source raises fails the poll
+                attempt.
+        """
+        record = poll.tool_call
+        routed = self._tool_sources.route(record)
+        if routed is None:
+            raise LookupError(f"no tool source built the call of tool {record.tool_name}")
+        source, reference = routed
+        return await source.poll_tool_task(
+            replace(poll, tool_call=replace(record, tool_source_reference=reference))
+        )
+
+    async def cancel_tool_task(self, cancellation: ToolTaskCancellation) -> None:
+        """Cancel a tool source's task that ``cancel_turn`` abandoned.
+
+        The session's ``CancelToolTask`` Step calls this method best-effort
+        with the call's record, whose ``cancellation_reason`` is
+        ``TURN_CANCELLED``, and then records the call ``OUTCOME_UNKNOWN``.
+        The source that built the call gets it with its own reference. A
+        source without ``cancel_tool_task`` leaves its task running.
+
+        .. versionadded:: 0.6.6
+
+        Args:
+            cancellation: The abandoned call's record from ``CancelToolTask``.
+
+        Raises:
+            LookupError: If no tool source built the call, or the source that
+                built it is not configured on this Worker. ``CancelToolTask``
+                retries and then records the call ``OUTCOME_UNKNOWN``.
+            Exception: What the source raises.
+        """
+        record = cancellation.tool_call
+        routed = self._tool_sources.route(record)
+        if routed is None:
+            raise LookupError(f"no tool source built the call of tool {record.tool_name}")
+        source, reference = routed
+        cancel = getattr(source, "cancel_tool_task", None)
+        if cancel is None:
+            return
+        await cancel(
+            replace(cancellation, tool_call=replace(record, tool_source_reference=reference))
+        )
+
     def tool_call_step_options(self, tool_call: AgentToolCall) -> ToolCallStepOptions | None:
         """Return a tool source's options for one call's tool Step.
 

agents/python/claude-agent-sdk/tests/claude_test_agent.py

@@ -23,7 +23,10 @@
 needs approval above 50; ``post_note`` is a single-attempt Mutation and
 ``policy_note`` the same under a session policy that asks for approval and
 claims ``READ``. ``HidingToolSource`` lists ``hiding_note``, the same
-single-attempt Mutation, until a test hides it.
+single-attempt Mutation, until a test hides it. ``TaskToolSource`` lists
+``report``, a ``READ`` tool without approval that starts a tool task, and
+``UncancellableTaskToolSource`` lists the same tool as ``uncancellable_report``
+without ``cancel_tool_task``.
 
 An application PreToolUse hook rewrites ``publish`` input ``"rewrite me"`` to
 ``"rewritten"``, raises a ``payout_checked`` amount to 500 when its memo is
@@ -84,10 +87,14 @@
     RunModelLimit,
     ToolCallStepOptions,
     ToolEffect,
+    ToolExecutionError,
     ToolExecutionResult,
     ToolPolicy,
+    ToolTaskCancellation,
+    ToolTaskPending,
     ToolTaskPoll,
     ToolTaskPollOutcome,
+    ToolTaskStarted,
 )
 from dex_agents.builtin_loop import SourceToolDefinition
 from dex_agents.connectors import ConnectorToolSource, connector_tool
@@ -139,6 +146,8 @@
     "policy_note": ToolPolicy(ToolEffect.READ, requires_approval=True),
     "screenshot": ToolPolicy(ToolEffect.READ, requires_approval=True),
     "draw": ToolPolicy(ToolEffect.NON_IDEMPOTENT_WRITE, requires_approval=False),
+    "report": ToolPolicy(ToolEffect.READ, requires_approval=False),
+    "uncancellable_report": ToolPolicy(ToolEffect.READ, requires_approval=False),
 }
 TEST_IMAGES = {
     image["name"]: image
@@ -191,7 +200,12 @@ def build_options(context: ClaudeAgentOptionsContext) -> ClaudeAgentOptions:
         build_media_store(state_dir) if os.environ.get(MEDIA_STORE_ENVIRONMENT_VARIABLE) else None
     )
     custom_tools = build_custom_tools(state_dir)
-    tool_sources: list[Any] = [build_connector_source(state_dir), HidingToolSource(state_dir)]
+    tool_sources: list[Any] = [
+        build_connector_source(state_dir),
+        HidingToolSource(state_dir),
+        TaskToolSource(state_dir),
+        UncancellableTaskToolSource(state_dir),
+    ]
     if media_store is not None:
         custom_tools.append(build_draw_tool(state_dir))
         tool_sources.append(ScreenshotToolSource(state_dir, media_store))
@@ -358,6 +372,84 @@ async def poll_tool_task(self, poll: ToolTaskPoll) -> ToolTaskPollOutcome:
         return await self._inner.poll_tool_task(poll)
 
 
+class UncancellableTaskToolSource:
+    """Lists ``uncancellable_report``, whose call starts a task that ``poll_tool_task`` polls.
+
+    The call returns ``ToolTaskStarted`` with task ID ``task-<tool call ID>``.
+    Arguments: ``topic`` names the result ``report on <topic>``; ``polls`` is
+    the poll that finishes the task, counted from 1; ``fail`` makes that poll
+    raise ``ToolExecutionError``; ``hold`` keeps the task pending. The
+    invocation log records ``start <task>``, ``poll <task> <n>``, and
+    ``cancel-task <task> <status> <cancellation_reason>`` for a source with
+    ``cancel_tool_task``. This one has none, so ``cancel_turn`` leaves its task
+    running.
+    """
+
+    source_name = "uncancellable-reports"
+    tool_name = "uncancellable_report"
+
+    def __init__(self, state_dir: Path) -> None:
+        self._invocations = InvocationLog(state_dir / INVOCATION_LOG_NAME)
+
+    async def list_tools(self) -> Sequence[SourceToolDefinition]:
+        return [_ListedTool(self.tool_name, "Write a report as a task")]
+
+    def deferred_tool_call(
+        self,
+        turn: AgentTurnContext,
+        *,
+        tool_call_id: str,
+        tool_name: str,
+        arguments: str,
+        listed_tool: SourceToolDefinition,
+    ) -> DeferredToolCall:
+        return DeferredToolCall(
+            tool_call_id=tool_call_id,
+            tool_name=tool_name,
+            arguments=arguments,
+            policy=turn.resolve_tool_policy(tool_name, ToolEffect.READ),
+            tool_source_reference=json.dumps({"kind": "report"}),
+        )
+
+    async def execute_deferred_tool_call(self, call: DeferredToolExecution) -> DeferredToolOutcome:
+        if call.tool_call.tool_source_reference != json.dumps({"kind": "report"}):
+            raise RuntimeError("report got another source's reference")
+        if not call.supports_tool_tasks:
+            raise RuntimeError("the runner does not poll tool tasks")
+        task_id = f"task-{call.tool_call.tool_call_id}"
+        self._invocations.record(f"start {task_id}")
+        return ToolTaskStarted(task_id=task_id, poll_interval_seconds=1)
+
+    async def poll_tool_task(self, poll: ToolTaskPoll) -> ToolTaskPollOutcome:
+        record = poll.tool_call
+        if record.tool_source_reference != json.dumps({"kind": "report"}):
+            raise RuntimeError("report got another source's reference")
+        task_id = record.task_id or ""
+        earlier_polls = [
+            name for name in self._invocations.names() if name.startswith(f"poll {task_id} ")
+        ]
+        poll_number = len(earlier_polls) + 1
+        self._invocations.record(f"poll {task_id} {poll_number}")
+        arguments = json.loads(record.arguments)
+        if arguments.get("hold") or poll_number < int(arguments.get("polls", 1)):
+            return ToolTaskPending(poll_interval_seconds=1)
+        if arguments.get("fail"):
+            raise ToolExecutionError(f"the report on {arguments['topic']} failed")
+        return ToolExecutionResult(output=f"report on {arguments['topic']}")
+
+
+class TaskToolSource(UncancellableTaskToolSource):
+    """Lists ``report``, the same task tool, and cancels its tasks."""
+
+    source_name = "reports"
+    tool_name = "report"
+
+    async def cancel_tool_task(self, cancellation: ToolTaskCancellation) -> None:
+        record = cancellation.tool_call
+        reason = record.cancellation_reason.value if record.cancellation_reason else "none"
+        self._invocations.record(f"cancel-task {record.task_id} {record.status.value} {reason}")
+
+
 def _run_model_limit_from_environment() -> RunModelLimit | None:
     limit = os.environ.get(RUN_MODEL_LIMIT_ENVIRONMENT_VARIABLE)
     if limit is None: