langchain-ai/deepagents
1
0
Fork 0
mirror of https://github.com/langchain-ai/deepagents.git synced 2026-07-21 17:25:26 -04:00
Code Issues 169 Packages Projects Releases 16 Wiki Activity

feat(code): tool.use and tool.result hook events (#3954)

Browse Source 
Closes #3953

---

Adds two hook events — `tool.use` (fires before a tool executes) and
`tool.result` (fires after) — mirroring Claude Code's
PreToolUse/PostToolUse hooks. They cover audit logging, notifications,
guardrails, and latency tracking. Both events are wired into the two
execution surfaces — the interactive Textual TUI and the headless (`-p`)
runner — so hooks fire consistently regardless of how the agent is
invoked. The existing `tool.error` event is untouched and continues to
fire alongside `tool.result` on failure, so existing hooks are
unaffected.

**Payloads**

`tool.use` carries `tool_name`, `tool_id`, and the parsed `tool_args`.
`tool.result` adds `tool_status` (`"success"` or `"error"`) and
`tool_output` (the tool's returned content, truncated to a shared
`HOOK_TOOL_OUTPUT_LIMIT`). The full schema and a worked example live in
the `hooks` module docstring.

**Resolving streamed arguments**

Tool-call arguments arrive as JSON string fragments while the model
streams. Each in-progress call is tracked in a typed `_ToolCallBuffer`,
and `_parse_tool_call_args` reassembles the fragments; `tool.use`
dispatches only once the arguments parse into a complete value, so hooks
receive the real arguments rather than a partial or empty payload.
Arguments that look structurally complete but still fail to parse are
logged rather than silently dropped, so a missed hook is never
invisible.

**Dispatch and delivery**

Both events are dispatched fire-and-forget on both surfaces, so a slow
hook command never blocks the agent or stalls the UI. Because they run
as background tasks, each surface drains any in-flight hook tasks before
its event loop tears down — the headless runner before the process
exits, and the TUI on app shutdown — so the final `tool.result` is never
cancelled mid-flight and silently dropped. Both surfaces also emit
`tool.result` for every executed tool, including a tool call whose
result arrives without a rendered widget, so audit consumers see the
same event stream in either mode. Rejected HITL approvals that abort a
turn now emit terminal `tool.error`/`tool.result` events, so hook
consumers never observe an unterminated `tool.use` for a tool that did
not run.

## Release note

Deep Agents Code now emits `tool.use` and `tool.result` hook events (in
addition to the existing `tool.error`), fired before and after each tool
call in both the interactive TUI and headless (`-p`) modes. Rejected
HITL approvals also emit terminal hook events when the turn aborts, so
every `tool.use` is closed by a matching result or error event.
Configure them under `~/.deepagents/hooks.json`; see the `hooks` module
docstring for the payload schema.

LinkedIn: https://www.linkedin.com/in/ninaadrrao/

---------

Co-authored-by: Mason Daugherty <mason@langchain.dev>
Co-authored-by: Mason Daugherty <github@mdrxy.com>
This commit is contained in:
Ninaad R. Rao
2026-07-05 16:24:51 -07:00
committed by GitHub
parent 59a665f339
commit ba1979d64c
12 changed files with 6786 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/code/deepagents_code/_tool_stream.py
+656
View File
@@ -0,0 +1,656 @@
"""Shared streaming tool-call buffering and hook-payload construction.
Both execution surfaces reassemble the same streamed tool-call state and fire
the same `tool.use` / `tool.result` / `tool.error` hook payloads:
- the interactive Textual TUI (`deepagents_code.textual_adapter`), and
- the headless runner (`deepagents_code.non_interactive`).
This module holds the single implementation of the buffering, argument parsing,
and payload-shape logic, so the two surfaces cannot drift *in those layers*. The
dispatch, gating, result-correlation, and buffer-lifecycle layers are still
implemented separately in each surface (each calls
`dispatch_hook_fire_and_forget` from its own namespace — the dispatch seam tests
patch) and must be kept in sync by hand.
Why the split exists
--------------------
The two surfaces have fundamentally different object lifecycles, which prevents
a single shared dispatch function:
- **TUI**: tool calls are backed by `ToolCallMessage` widgets that persist
parsed args as instance state and render success/error visually. `tool.use`
fires at widget-mount time because that is when the parsed args are available
— they are what the widget is constructed with (or come from a validated HITL
interrupt). Result correlation reads args from `tool_msg.args` (a widget
property).
- **Headless**: there are no widgets. Tool-call state lives in a plain dict on
`StreamState` (`in_flight_tool_calls`, keyed by tool-call id). `tool.use`
fires in the stream loop once args parse and the tool-call id is known.
Result correlation pops the record from that dict.
Additionally, only the TUI has interactive HITL approval *widgets* that can be
rejected before any `ToolMessage` streams back, so only it needs
`_dispatch_terminal_tool_result_hooks` and the `completed_tool_result_ids`
duplicate-suppression set for synthetic middleware `ToolMessage` re-arrivals
after a resumed turn. The headless runner *does* have a HITL interrupt/resume
flow, but it never pre-dispatches terminal hooks before a resume — a rejection
arrives as a synthetic `ToolMessage` handled by the normal result path — so
there is no already-emitted terminal hook to suppress and therefore no
equivalent race.
Unifying these into a single function would require either a common
widget-or-dict abstraction (indirection with no behavioral gain) or pushing
widget-aware logic into this shared module (coupling headless mode to TUI
concepts it does not use). The split is deliberate: share everything that *can*
drift (payload shapes, arg parsing, status normalization, output truncation, and
the end-of-stream classification of tool calls that never emitted a `tool.use`)
here, and keep everything that *must* differ (when to dispatch, where args come
from, HITL handling) in each surface.
Parity contract for hook consumers
----------------------------------
A hook consumer can rely on these guarantees being identical across surfaces:
- **Payload shape**: every `tool.use`, `tool.result`, and `tool.error` payload
is built by the shared builders in this module, so field names and
truncation are the same.
- **Event completeness**: every executed tool emits `tool.result`, including
tools whose args never parsed or that carried no tool-call id (both surfaces
emit with `{}` args in that case). A tool whose stream is aborted before its
result — a cancelled turn or a mid-stream error — is also closed with a
terminal `tool.error`/`tool.result` (TUI via
`_dispatch_terminal_tool_result_hooks`, headless via
`_dispatch_orphaned_tool_result_hooks`), so no `tool.use` is left dangling.
- **Fire-once-per-id**: `tool.use` fires at most once per tool-call id on both
surfaces, each gated by a monotonic id set that is never discarded within a
turn (TUI via `displayed_tool_ids`, headless via `emitted_tool_use_ids`). A
redelivery of the same id's arg chunks *after* its result — non-standard for
`stream_mode="messages"` — is therefore ignored on both surfaces rather than
re-firing `tool.use`. Headless additionally keeps `in_flight_tool_calls` for
result correlation and the orphan drain; that map *is* cleared per result,
so it is deliberately not the fire-once guard.
- **`tool.error` co-firing**: whenever `tool_status` is `"error"`, both
surfaces emit `tool.error` alongside `tool.result`.
What is allowed to differ (and is not part of the contract):
- **Dispatch timing**: the TUI dispatches `tool.use` at widget mount;
headless dispatches it in the stream loop once args parse. A hook subscriber
should not assume `tool.use` fires at an identical point in the surface's
internal lifecycle — only that it fires before `tool.result` for the same
`tool_id`.
- **HITL rejection events**: both surfaces emit `tool.error`/`tool.result` for
a rejected tool, but by different routes. Headless (and the TUI on a resumed
turn) emits them from the synthetic `ToolMessage` on the normal result path;
the TUI *additionally* emits them via the dedicated
`_dispatch_terminal_tool_result_hooks` drain for tools rejected or cancelled
in the interactive UI *before* any `ToolMessage` streams back. Only that
pre-`ToolMessage` terminal drain is TUI-only — headless has no interactive
approval path that can reject a call before its result arrives.
When changing the dispatch or gating logic in one surface, verify the parity
contract still holds against the other surface.
The payload schema is documented in `deepagents_code.hooks`.
"""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypedDict
from deepagents_code.hooks import HOOK_TOOL_OUTPUT_LIMIT
if TYPE_CHECKING:
from collections.abc import Iterable
logger = logging.getLogger(__name__)
ToolStatus = Literal["success", "error"]
"""Terminal status of a tool call, mirroring `ToolMessage.status`."""
ToolCallBufferKey = int | str
"""Key for buffering an in-progress streamed tool call.
The streaming `index` (an `int`) when present, else the tool-call `id` (a
`str`), else a unique placeholder string (see `tool_call_buffer_key`). Exported
so both surfaces annotate their buffer maps identically rather than each
spelling the union — the exact drift this module exists to prevent.
"""
ProviderToolArgs = dict[str, Any] | list[Any] | str | int | float | bool | None
"""A whole tool-call arguments value delivered by a provider in one chunk.
Usually a `dict`; rarely a bare scalar or list. `None` means no whole value has
arrived yet (the args may instead be streaming as fragments in `args_parts`).
An `isinstance` check narrows it to the declared `dict[str, Any]` arg type at the
single `parse_args` return site.
"""
TOOL_OUTPUT_TRUNCATION_MARKER = "…[output truncated]"
"""Suffix appended to a `tool.result` `tool_output` that hit
`HOOK_TOOL_OUTPUT_LIMIT`, so a consumer can distinguish a capped result from a
genuinely short one. The marker is counted within the cap, so the final
`tool_output` length never exceeds `HOOK_TOOL_OUTPUT_LIMIT`."""
UNRENDERABLE_TOOL_OUTPUT = "<tool output could not be rendered>"
"""Sentinel `tool_output` used when formatting/coercing a tool result raises.
Lets both surfaces keep the terminal `tool.result` dispatch unconditional
without re-touching the offending content (whose `__str__`/`__repr__` may itself
raise), so a hook consumer still sees the result rather than a dropped event."""
def normalize_tool_status(raw_status: object, tool_name: str) -> ToolStatus:
"""Map a raw `ToolMessage.status` to the two-value hook domain, fail-closed.
`"error"` and `"success"` pass through. Any other *present* value — a future
provider status, an explicit `None`, or a typo — is unexpected and treated as
`"error"` (and logged), so an audit or notification hook is never told a
non-successful tool succeeded. Callers pass
`getattr(message, "status", "success")`, so a missing status arrives as
`"success"` and is not warned about.
Args:
raw_status: The raw `status` value read off the `ToolMessage`.
tool_name: The tool name, included in the warning for context.
Returns:
`"success"` or `"error"`.
"""
if raw_status == "error":
return "error"
if raw_status == "success":
return "success"
logger.warning(
"Unexpected ToolMessage.status %r for tool %s; treating as error",
raw_status,
tool_name,
)
return "error"
class ToolUsePayload(TypedDict):
"""`tool.use` hook payload (schema documented in `hooks`)."""
tool_name: str
"""The tool being invoked."""
tool_id: str
"""The tool-call id. Always a real id — a call with no id never produces a
`tool.use` (see `hooks`), so unlike `tool.result` this is never `None`."""
tool_args: dict[str, Any]
"""The parsed tool-call arguments."""
class ToolErrorPayload(TypedDict):
"""`tool.error` hook payload (schema documented in `hooks`)."""
tool_names: list[str]
"""Names of the tools whose calls failed or were rejected."""
class ToolResultPayload(TypedDict):
"""`tool.result` hook payload (schema documented in `hooks`)."""
tool_name: str
"""The tool that produced the result. Usually the real tool name; falls back
to `""` in the rare uncorrelated case where the `ToolMessage` carries no
`name` and no `tool.use` was correlated to supply one."""
tool_id: str | None
"""The tool-call id, or `None` when it could not be correlated."""
tool_args: dict[str, Any]
"""The parsed tool-call arguments, or `{}` when uncorrelated."""
tool_status: ToolStatus
"""`"success"`, or `"error"` for a failed, rejected, or cancelled call."""
tool_output: str
"""The tool's returned content, capped to `HOOK_TOOL_OUTPUT_LIMIT`. When the
full output exceeded the cap, the value ends with
`TOOL_OUTPUT_TRUNCATION_MARKER` so a consumer (e.g. a secret/policy scanner)
can tell a truncated result from a genuinely short one."""
def tool_call_buffer_key(
index: int | str | None, tool_id: str | None, count: int
) -> ToolCallBufferKey:
"""Compute a stable key for buffering an in-progress streamed tool call.
Prefers the streaming `index` (stable across fragments of one call), then
the tool-call `id`, falling back to a positional placeholder so unrelated
id-less calls don't collide.
Args:
index: The `index` field from the streamed tool-call chunk, if any.
Typed loosely because it is read from an untyped content-block dict.
tool_id: The tool-call `id` from the chunk, if any.
count: The current number of buffered calls, used to make the fallback
placeholder unique.
Returns:
The chunk `index`, else the `tool_id`, else a unique placeholder string.
"""
if index is not None:
return index
if tool_id is not None:
return tool_id
return f"unknown-{count}"
def _looks_structurally_complete(s: str) -> bool:
"""Return whether `s` is a balanced JSON container, string-state aware.
Scans for bracket balance while tracking whether the cursor is inside a
string literal (honoring backslash escapes), so a value is "complete" only
when every `{`/`[` is matched and the text does not end mid-string. Used
solely to decide whether a `json.loads` failure is worth warning about: a
balanced-but-unparseable value is malformed, whereas a partial stream — e.g.
a chunk boundary landing right after an inner `}` in `{"edits": [{"a": 1}` —
leaves an outer container open and returns `False`, so no false warning fires
on a healthy mid-stream fragment. A stray closer (`depth < 0`) is unbalanced
and can never be completed by more input, so it is reported complete
(malformed) too. Iterative, so it never raises `RecursionError` on
pathologically nested input.
Args:
s: The accumulated, stripped tool-call argument text.
Returns:
`True` if the brackets are matched (or over-closed) and the text does not
end inside a string; `False` while a container is still open.
"""
depth = 0
in_string = False
escaped = False
for ch in s:
if in_string:
if escaped:
escaped = False
elif ch == "\\":
escaped = True
elif ch == '"':
in_string = False
continue
if ch == '"':
in_string = True
elif ch in "{[":
depth += 1
elif ch in "}]":
depth -= 1
if depth < 0:
# More closers than openers: unbalanced, not a partial stream.
return True
return depth == 0 and not in_string
@dataclass
class ToolCallBuffer:
"""In-progress state for a single streamed tool call.
`args` and `args_parts` are two representations of the arguments used one at
a time, depending on whether the provider delivers the value whole or in
JSON string fragments.
"""
name: str | None = None
"""The tool name, once a chunk has supplied it."""
tool_id: str | None = None
"""The tool-call id, once a chunk has supplied it."""
args: ProviderToolArgs = None
"""A fully materialized arguments value delivered by a single chunk (dict or,
rarely, a scalar). Mutually exclusive with `args_parts` (see `__post_init__`
and `ingest`)."""
args_parts: list[str] = field(default_factory=list)
"""JSON string fragments accumulated across chunks, reassembled by
`parse_args` once the payload looks complete. Mutually exclusive with
`args` (see `__post_init__` and `ingest`)."""
displayed: bool = False
"""One-shot latch guarding the single "Calling tool" console line (used only
by the headless surface)."""
warned: bool = False
"""One-shot latch so a malformed-but-complete payload is logged at most once,
even though `parse_args` re-runs on every later chunk for the retained
buffer."""
def __post_init__(self) -> None:
"""Enforce the `args` XOR `args_parts` invariant at construction time.
`ingest` maintains this on every chunk. This guard catches an illegal
buffer built directly with both fields set; the fields are public and
mutable, though, so a caller can still reach the both-populated state by
assigning them after construction. `parse_args` re-checks the invariant
at read time (raising rather than silently reading `args` first and
masking a conflicting `args_parts`), so the illegal state fails loudly
wherever it originates.
Raises:
ValueError: If both `args` and `args_parts` are set.
"""
if self.args is not None and self.args_parts:
msg = "ToolCallBuffer cannot hold both args and args_parts"
raise ValueError(msg)
def _reset_for_new_call(self) -> None:
"""Discard retained state before this buffer is reused for another call."""
self.name = None
self.tool_id = None
self.args = None
self.args_parts = []
self.displayed = False
self.warned = False
def ingest(
self,
*,
name: str | None,
tool_id: str | None,
args: Any, # noqa: ANN401 # provider-shaped tool-call args chunk
) -> None:
"""Fold one streamed tool-call chunk's fields into the buffer.
A dict `args` replaces any accumulated fragments (the provider delivered
the whole value at once); a string `args` is appended as a fragment; any
other non-`None` value is stored as-is.
Args:
name: The tool name from this chunk, if present.
tool_id: The tool-call id from this chunk, if present.
args: The `args` field from this chunk (dict, string fragment, or
other scalar).
"""
# A differing id on the same buffer key means a *new* call has reused
# this streaming index. Indices restart per message, so a buffer retained
# from an earlier message or HITL-resume round (e.g. one whose args never
# parsed) can collide here. Reset the accumulated arg state so the new
# call's fragments never append onto the old call's leftover — which would
# otherwise leave both unparseable and silently drop the new call's
# `tool.use`. Per-call metadata is reset too, so stale name/display state
# cannot bleed into the new call if its first chunk only carries the id.
#
# Some providers deliver the new call's name/args before its replacement
# id. In that delayed-id shape, a retained `self.tool_id` is already stale
# even though this chunk has no id to compare against; clear it before
# folding the new name/args so parsed args cannot dispatch under the old
# call id.
#
# This assumes the standard LangChain streaming contract: a call's name
# arrives only on its first chunk, so id-less continuation chunks for the
# same call carry only args and keep accumulating below. A non-standard
# provider that *repeats* the name on an id-less continuation chunk would
# trip `delayed_id_reuses_index` mid-call and reset the accumulated args.
# That degrades gracefully rather than crashing: the call's `tool.use`
# (and its parsed args) is lost, but the tool still executes and its
# `tool.result` still fires via the uncorrelated `{}`-args path. No
# observed provider (Anthropic/OpenAI) streams that shape.
new_id_reuses_index = (
tool_id is not None and self.tool_id is not None and tool_id != self.tool_id
)
delayed_id_reuses_index = (
tool_id is None
and self.tool_id is not None
and name is not None
and self.name is not None
)
if new_id_reuses_index or delayed_id_reuses_index:
self._reset_for_new_call()
if name:
self.name = name
if tool_id:
self.tool_id = tool_id
# `args` (a whole value) and `args_parts` (JSON fragments) are mutually
# exclusive; each branch clears the counterpart so the buffer never holds
# both. Together with the `__post_init__` guard this keeps the invariant
# true for every buffer, so `parse_args` reads `args` first without
# masking a conflicting `args_parts` via read order. A provider streams a
# single call as either whole values or fragments, never a mix, so no
# real sequence loses data.
if isinstance(args, dict):
self.args = args
self.args_parts = []
elif isinstance(args, str):
# Append every non-empty fragment unconditionally. An earlier
# TUI-only version skipped a fragment equal to the immediately
# preceding one to dedup accidental redelivery; that guard was
# dropped because a stream can legitimately emit two identical
# consecutive deltas (e.g. two `", "` fragments) and skipping one
# corrupts the reassembled JSON. Standard `stream_mode="messages"`
# streaming never redelivers a fragment, so unconditional append is
# both lossless and correct in practice.
if args:
self.args = None
self.args_parts.append(args)
elif args is not None:
self.args = args
self.args_parts = []
def parse_args(self) -> dict[str, Any] | None:
"""Return the tool-call args once enough data has arrived, else `None`.
A non-object JSON value (a bare scalar or list — rare for tool calls) is
wrapped as `{"value": ...}` so a caller's `tool_args` is always an
object.
Returns:
Parsed tool-call arguments, or `None` when the args are not yet
complete (still streaming), empty, or structurally complete but
unparseable — malformed or too deeply nested (the
structurally-complete malformed case is logged once via
`warned`).
Raises:
ValueError: If both `args` and `args_parts` are populated, violating
the invariant `ingest`/`__post_init__` maintain. Guarded here so
a caller that assigned the public fields directly fails loudly
instead of silently dropping the fragments (read order would
otherwise return `args` and discard `args_parts`).
"""
if self.args is not None and self.args_parts:
msg = "ToolCallBuffer cannot hold both args and args_parts"
raise ValueError(msg)
if isinstance(self.args, dict):
# A whole-value dict delivered by the provider. The `isinstance`
# narrows `ProviderToolArgs` to the declared `dict[str, Any]` arg
# type, so this returns without a cast.
return self.args
if self.args is not None:
return {"value": self.args}
if not self.args_parts:
return None
joined = "".join(self.args_parts)
stripped = joined.strip()
if not stripped:
return None
# Cheap structural pre-check: bail while a JSON object/array is still
# open so we don't attempt to parse a partial streamed fragment. A
# well-formed object's closing brace is always its last char. The
# `"".join` still runs per fragment (so accumulation stays O(n^2) across
# the stream); the win is deferring the costlier `json.loads` until the
# value looks complete, rather than re-parsing the whole prefix on every
# fragment.
#
# A bracketed value with trailing junk after its close (e.g.
# `{"a": 1} x`) also returns here and is treated as still-streaming, so
# it never reaches the malformed warning below — the same deliberate
# trade as the non-bracketed-scalar exclusion documented at the end of
# this method. Real provider arg streams are pure JSON, so this shape
# does not occur in practice; if such a call still executes, its
# `tool.result` logs the correlation miss.
if stripped[0] in "{[" and not stripped.endswith(("}", "]")):
return None
try:
parsed = json.loads(joined)
except (json.JSONDecodeError, RecursionError):
# Args that are structurally balanced (all brackets matched, not
# ending mid-string) but still fail to parse are malformed, not
# mid-stream — surface them rather than silently dropping the
# tool.use hook. `_looks_structurally_complete` does a string-aware
# balance check rather than the cheaper "starts with {/[ and ends
# with }/]" heuristic, which false-positives on a normal chunk
# boundary inside nested args (e.g. `{"edits": [{"a": 1}`) and would
# warn on a perfectly healthy stream. `RecursionError` is caught
# alongside the decode error: pathologically nested model output
# makes `json.loads` exceed the interpreter recursion limit, and that
# is one malformed call to skip, not a reason to let the exception
# escape and abort the whole turn. `%r` escapes any control
# characters in the model-generated fragment. The `warned` latch
# keeps this to one line per call, since the buffer is retained and
# parse_args re-runs on every later chunk for the same call.
if (
stripped[0] in "{["
and _looks_structurally_complete(stripped)
and not self.warned
):
self.warned = True
logger.warning(
"Tool-call args look complete but failed to parse: %r",
joined[:200],
)
# A non-bracketed complete-but-malformed value is deliberately not
# warned here: it is indistinguishable from a still-streaming scalar
# fragment, so warning would be noisy. If such a call still executes,
# its `tool.result` logs the correlation miss at info; if it never
# executes there is no result to audit, so nothing is lost.
return None
if not isinstance(parsed, dict):
return {"value": parsed}
return parsed
class UnemittedToolCalls(NamedTuple):
"""Counts of buffered tool calls that never emitted a `tool.use`.
A `NamedTuple` so callers can still unpack positionally while the field names
keep the two same-typed slots from being load-bearing at every call site (a
bare `(int, int)` invites a silent transposition). Both surfaces log the two
counts as separate diagnostics.
"""
unparsed: int
"""Named buffers whose args never parsed."""
idless_parsed: int
"""Named buffers whose args parsed but whose `tool_id` stayed `None`."""
def count_unemitted_tool_calls(buffers: Iterable[ToolCallBuffer]) -> UnemittedToolCalls:
"""Classify buffered tool calls that never emitted a `tool.use`.
Both surfaces log the same end-of-stream diagnostic for tool calls still in
their buffer map when the stream ends: those whose args never parsed, and
those whose args parsed but whose `tool_id` stayed `None` (so `tool.use` was
gated out). Sharing the classification here keeps the two diagnostics from
drifting; each surface still emits its own log lines. `parse_args` is safe to
re-run (idempotent bar its one-shot `warned` latch).
Args:
buffers: The in-progress tool-call buffers remaining at stream end.
Returns:
The `unparsed` and `idless_parsed` counts (see `UnemittedToolCalls`).
"""
unparsed = 0
idless_parsed = 0
for buffer in buffers:
if buffer.name is None:
continue
if buffer.parse_args() is None:
unparsed += 1
elif buffer.tool_id is None:
idless_parsed += 1
return UnemittedToolCalls(unparsed, idless_parsed)
def build_tool_use_payload(
tool_name: str, tool_id: str, tool_args: dict[str, Any]
) -> ToolUsePayload:
"""Build the `tool.use` hook payload (schema documented in `hooks`).
Args:
tool_name: The tool being invoked.
tool_id: The tool-call id. Always a real id for an emitted `tool.use`
(a call with no id never produces one).
tool_args: The parsed tool-call arguments.
Returns:
The `tool.use` payload dict.
"""
return {
"tool_name": tool_name,
"tool_id": tool_id,
"tool_args": tool_args,
}
def build_tool_error_payload(tool_name: str) -> ToolErrorPayload:
"""Build the `tool.error` hook payload (schema documented in `hooks`).
Args:
tool_name: The tool whose call failed or was rejected.
Returns:
The `tool.error` payload dict.
"""
return {"tool_names": [tool_name]}
def build_tool_result_payload(
tool_name: str,
tool_id: str | None,
tool_args: dict[str, Any],
tool_status: ToolStatus,
tool_output: str,
) -> ToolResultPayload:
"""Build the `tool.result` hook payload (schema documented in `hooks`).
`tool_output` is capped to `HOOK_TOOL_OUTPUT_LIMIT` here so both surfaces
apply the identical cap regardless of where the raw output originates. When
the cap fires the value ends with `TOOL_OUTPUT_TRUNCATION_MARKER` (counted
within the cap, so the result never exceeds the limit) so a consumer can tell
a capped result from a short one. `tool_args` is intentionally not truncated
(see `HOOK_TOOL_OUTPUT_LIMIT`).
Args:
tool_name: The tool that produced the result.
tool_id: The tool-call id, or `None` when it could not be correlated.
tool_args: The parsed tool-call arguments, or `{}` when uncorrelated.
tool_status: `"success"` or `"error"`.
tool_output: The tool's returned content (capped in the payload).
Returns:
The `tool.result` payload dict with `tool_output` capped and marked when
truncation occurred.
"""
if len(tool_output) > HOOK_TOOL_OUTPUT_LIMIT:
# `max(..., 0)` guards a future `HOOK_TOOL_OUTPUT_LIMIT` set below the
# marker length: a negative `keep` would slice from the end and keep the
# tail instead of truncating. Positive with today's constants (2000 vs a
# ~19-char marker); this only future-proofs a constant change.
keep = max(HOOK_TOOL_OUTPUT_LIMIT - len(TOOL_OUTPUT_TRUNCATION_MARKER), 0)
capped_output = tool_output[:keep] + TOOL_OUTPUT_TRUNCATION_MARKER
else:
capped_output = tool_output
return {
"tool_name": tool_name,
"tool_id": tool_id,
"tool_args": tool_args,
"tool_status": tool_status,
"tool_output": capped_output,
}
libs/code/deepagents_code/app.py
+65 -27
View File
@@ -11930,7 +11930,12 @@ class DeepAgentsApp(App):
# Dispatch synchronously — the event loop is about to be torn down by
# super().exit(), so an async task would never complete.
from deepagents_code.hooks import _dispatch_hook_sync, _load_hooks
from deepagents_code.hooks import (
_dispatch_hook_sync,
_load_hooks,
drain_pending_hooks,
has_pending_hooks,
)
hooks = _load_hooks()
if hooks:
@@ -11961,38 +11966,71 @@ class DeepAgentsApp(App):
# finish persisting the in-flight run's trace instead of being
# SIGTERM'd mid-request.
agent_worker = self._agent_worker if self._agent_running else None
should_wait_for_agent = (
agent_worker is not None and not agent_worker.is_finished
)
should_drain_hooks = has_pending_hooks()
if agent_worker is not None and not agent_worker.is_finished:
if should_wait_for_agent or should_drain_hooks:
async def _graceful_exit() -> None:
from textual.worker import WorkerCancelled, WorkerFailed
try:
await asyncio.wait_for(
asyncio.shield(agent_worker.wait()),
timeout=_GRACEFUL_EXIT_WAIT_SECONDS,
)
except (asyncio.CancelledError, WorkerCancelled):
# Expected: exit() cancelled the worker above, so its
# cancellation handler ran to completion. Nothing to flag.
logger.debug(
"Agent worker cancelled cleanly before app exit",
exc_info=True,
)
except (TimeoutError, WorkerFailed):
# The worker did not finish within the window, so the
# in-flight run's server-side trace may be incomplete.
# Surface above debug so the loss isn't silent.
logger.warning(
"Agent worker did not finish persisting before app "
"exit; the in-flight run's trace may be incomplete",
exc_info=True,
)
except Exception:
logger.warning(
"Agent worker wait raised unexpectedly before app exit",
exc_info=True,
)
worker = agent_worker
if should_wait_for_agent and worker is not None:
try:
await asyncio.wait_for(
asyncio.shield(worker.wait()),
timeout=_GRACEFUL_EXIT_WAIT_SECONDS,
)
except (asyncio.CancelledError, WorkerCancelled):
# Expected: exit() cancelled the worker above, so
# its cancellation handler ran to completion.
logger.debug(
"Agent worker cancelled cleanly before app exit",
exc_info=True,
)
except (TimeoutError, WorkerFailed):
# The worker did not finish within the window, so
# the in-flight run's server-side trace may be
# incomplete. Surface above debug so the loss isn't
# silent.
logger.warning(
"Agent worker did not finish persisting before app "
"exit; the in-flight run's trace may be incomplete",
exc_info=True,
)
except Exception:
logger.warning(
"Agent worker wait raised unexpectedly before app exit",
exc_info=True,
)
try:
# Bound the drain so a hung hook subprocess can't stall
# an interactive quit indefinitely. Each hook is already
# capped by `hooks.HOOK_SUBPROCESS_TIMEOUT` in its own
# thread, but a slow/many-hook config could still exceed
# the graceful-exit budget; a dropped final tool.result is
# announced rather than letting the UI feel frozen. The
# headless surface leaves its drain unbounded on purpose —
# a script exit favors a complete audit trail over shutdown
# latency.
await asyncio.wait_for(
drain_pending_hooks(),
timeout=_GRACEFUL_EXIT_WAIT_SECONDS,
)
except TimeoutError:
logger.warning(
"Hook drain did not finish within %ss before app "
"exit; a final tool.result hook may be dropped",
_GRACEFUL_EXIT_WAIT_SECONDS,
)
except Exception:
logger.warning(
"Hook drain raised unexpectedly before app exit",
exc_info=True,
)
finally:
# This is the only call that stops the event loop, so it
# must run on every path the try/except can take, including
libs/code/deepagents_code/hooks.py
+156 -17
View File
@@ -15,6 +15,56 @@ If `events` is omitted or empty the hook receives **all** events.
Onboarding emits `user.name.set` with `{"name": "...", "assistant_id": "..."}`
after the user submits a non-empty preferred name.
`tool.use` fires before a tool call once its streamed arguments parse into a
complete value *and* its tool-call id is known; a call whose arguments never
parse, or that carries no id, is skipped. `tool.result` fires after every tool
call reaches a terminal state — successful execution, failure, or HITL
rejection/cancellation. The three blocks below show the payload *shapes*, not a
single sequence of events:
```jsonc
{"event": "tool.use", "tool_name": "write_file", "tool_id": "toolu_abc123",
"tool_args": {"file_path": "src/foo.py", "content": "..."}}
{"event": "tool.result", "tool_name": "write_file", "tool_id": "toolu_abc123",
"tool_args": {"file_path": "src/foo.py", "content": "..."},
"tool_status": "success", "tool_output": "Updated file src/foo.py"}
{"event": "tool.error", "tool_names": ["write_file"]}
```
`tool_args` is the parsed tool-call arguments; a non-object value (rare) is
wrapped as `{"value": ...}`. `tool_output` is the tool's returned content,
capped to `HOOK_TOOL_OUTPUT_LIMIT` characters (`tool_args` is not truncated); a
capped value ends with `…[output truncated]` so a consumer can tell a truncated
result from a short one.
`tool_status` is `"success"` or `"error"`; `"error"` covers both a tool that
raised and a call the user rejected or cancelled. Whenever a `tool.result` has
`tool_status: "error"`, `tool.error` (payload `{"tool_names": [<name>]}`) fires
alongside it, so existing `tool.error` hooks are unaffected.
`tool_args` is `{}` whenever a `tool.result` cannot be correlated back to a
`tool.use` — either because the call carried no id (then `tool_id` is `null`) or
because no `tool.use` fired for it (e.g. its args never parsed), in which case
`tool_id` may still be the real string id.
Ordering: the tool events (`tool.use`, `tool.result`, `tool.error`) are
dispatched fire-and-forget (see `dispatch_hook_fire_and_forget`) and every
matching hook command runs in its own subprocess. A `tool.use` is *dispatched*
before its `tool.result`, but the two run concurrently, so a hook subscribed to
both may observe them out of order, and events from parallel tool calls
interleave freely. Correlate by `tool_id` rather than relying on arrival order —
there is no cross-event delivery-ordering guarantee for the tool events. Most
non-tool events (`session.start`, `task.complete`, `session.end`, `user.prompt`,
`context.offload`, `context.compact`, `permission.request`) fire in program order.
They are dispatched with an awaited `dispatch_hook`, except `session.end` on the
interactive TUI, which is dispatched synchronously via `_dispatch_hook_sync` at
shutdown (the event loop is already tearing down); that path is still blocking and
in-order, so the program-order guarantee holds. `input.required` and
`user.name.set` are the exceptions with no program-order guarantee:
`user.name.set` is always dispatched fire-and-forget, and `input.required` is
fire-and-forget on the headless surface (awaited only in the interactive TUI).
"""
from __future__ import annotations
@@ -24,10 +74,36 @@ import json
import logging
import subprocess # noqa: S404
from concurrent.futures import ThreadPoolExecutor
from typing import Any
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Mapping
logger = logging.getLogger(__name__)
HOOK_TOOL_OUTPUT_LIMIT = 2000
"""Max characters of `tool_output` included in `tool.result` hook payloads.
Bounds payload size (data-amplification guard) while keeping enough of the
tool's output to be useful to audit/notification hooks. Applied in the single
shared builder `_tool_stream.build_tool_result_payload`, which both the
interactive and headless dispatch paths call, so the cap never drifts between
them. Only `tool_output` is capped; `tool_args` is passed through in full so hooks
that act on the arguments (e.g. a linter reading a `write_file` `content`) see
the exact value the tool received.
"""
HOOK_SUBPROCESS_TIMEOUT = 5
"""Seconds a single hook subprocess may run before it is killed.
Bounds how long one misbehaving hook can block the dispatch thread. Consumed in
code by the `subprocess.run` timeout and its timeout log message here, so those
two never drift. `app.py`'s graceful-exit comment names this symbol (rather than
a bare literal) so its prose can't go stale; note the drain itself is bounded
separately by `_GRACEFUL_EXIT_WAIT_SECONDS`, not by this value — a hook can run
up to this long while the aggregate drain gives up sooner.
"""
_hooks_config: list[dict[str, Any]] | None = None
"""Cached config — loaded lazily on first dispatch."""
@@ -84,8 +160,12 @@ def _load_hooks() -> list[dict[str, Any]]:
def _run_single_hook(command: list[str], event: str, payload_bytes: bytes) -> None:
"""Execute a single hook command, writing the JSON payload to its stdin.
Uses `subprocess.run` which automatically kills the child process on
timeout, preventing zombie/orphan process leaks.
On timeout `subprocess.run` kills and reaps only the direct hook process
(its `Popen.kill()` targets that one PID, never the process group), so any
grandchildren it spawned are left as orphans regardless — the timeout bounds
the hook process, not its whole descendant tree. `start_new_session=True`
does not change that; it only isolates the hook (and its descendants) into
their own session so a signal to our group doesn't reach them.
Args:
command: The command and arguments to run.
@@ -99,16 +179,26 @@ def _run_single_hook(command: list[str], event: str, payload_bytes: bytes) -> No
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
start_new_session=True,
timeout=5,
timeout=HOOK_SUBPROCESS_TIMEOUT,
check=False,
)
except subprocess.TimeoutExpired:
logger.warning("Hook command timed out (>5s) for event %s: %s", event, command)
logger.warning(
"Hook command timed out (>%ss) for event %s: %s",
HOOK_SUBPROCESS_TIMEOUT,
event,
command,
)
except (FileNotFoundError, PermissionError) as exc:
logger.warning("Hook command failed for event %s: %s — %s", event, command, exc)
except Exception:
logger.debug(
"Hook dispatch failed for event %s: %s",
# Unexpected failure (e.g. ENOEXEC for a non-executable hook file, an
# embedded null byte, or fd/memory exhaustion). These are the failures
# we understand least, so surface them at warning — the expected
# timeout / not-found / permission cases above are also warnings, and a
# silent debug here would hide a hook that never fires.
logger.warning(
"Hook dispatch failed unexpectedly for event %s: %s",
event,
command,
exc_info=True,
@@ -122,8 +212,8 @@ def _dispatch_hook_sync(
Iterates over all configured hooks, skipping those whose event filter
does not match or whose `command` is missing/invalid. Matching hooks are
executed concurrently with a 5-second timeout per command. Errors are caught
per-hook and logged without propagating.
executed concurrently, each bounded by `HOOK_SUBPROCESS_TIMEOUT` per command.
Errors are caught per-hook and logged without propagating.
Args:
event: Dotted event name (e.g. `'session.start'`).
@@ -134,6 +224,15 @@ def _dispatch_hook_sync(
for hook in hooks:
command = hook.get("command")
if not isinstance(command, list) or not command:
# A misconfigured `command` (missing, a bare string instead of an
# argv list, or empty) means this hook can never fire. Warn rather
# than silently skip so the config mistake is greppable instead of
# looking like the hook simply never matched.
logger.warning(
"Skipping hook with invalid `command` for event %s: %r",
event,
command,
)
continue
events = hook.get("events")
@@ -158,26 +257,30 @@ def _dispatch_hook_sync(
future.result()
async def dispatch_hook(event: str, payload: dict[str, Any]) -> None:
async def dispatch_hook(event: str, payload: Mapping[str, Any]) -> None:
"""Fire matching hook commands with `payload` serialized as JSON on stdin.
The `event` name is automatically injected into the payload under the
`"event"` key so callers don't need to duplicate it.
The blocking subprocess work is offloaded to a thread so the caller's
event loop is never stalled. Matching hooks run concurrently, each with
a 5-second timeout. Errors are logged and never propagated.
event loop is never stalled. Matching hooks run concurrently, each bounded
by `HOOK_SUBPROCESS_TIMEOUT`. Errors are logged and never propagated.
Args:
event: Dotted event name (e.g. `'session.start'`).
payload: Arbitrary JSON-serializable dict sent on the command's stdin.
payload: Arbitrary JSON-serializable mapping sent on the command's stdin.
"""
try:
hooks = _load_hooks()
if not hooks:
return
payload_bytes = json.dumps({"event": event, **payload}).encode()
# `default=str` degrades a non-JSON-serializable value (e.g. a
# provider-delivered whole-value arg object) to its string form rather
# than raising and dropping the entire hook event — the invocation stays
# auditable even if one field isn't natively serializable.
payload_bytes = json.dumps({"event": event, **payload}, default=str).encode()
await asyncio.to_thread(_dispatch_hook_sync, event, payload_bytes, hooks)
except Exception:
logger.warning(
@@ -187,7 +290,7 @@ async def dispatch_hook(event: str, payload: dict[str, Any]) -> None:
)
def dispatch_hook_fire_and_forget(event: str, payload: dict[str, Any]) -> None:
def dispatch_hook_fire_and_forget(event: str, payload: Mapping[str, Any]) -> None:
"""Schedule `dispatch_hook` as a background task with a strong reference.
Use this instead of bare `create_task(dispatch_hook(...))` to prevent the
@@ -197,13 +300,49 @@ def dispatch_hook_fire_and_forget(event: str, payload: dict[str, Any]) -> None:
Args:
event: Dotted event name (e.g. `'session.start'`).
payload: Arbitrary JSON-serializable dict sent on the command's stdin.
payload: Arbitrary JSON-serializable mapping sent on the command's stdin.
"""
try:
loop = asyncio.get_running_loop()
except RuntimeError:
logger.debug("No running event loop; skipping hook for %s", event)
# A dropped hook is an audit/notification gap, so surface it at warning
# rather than debug. In the streaming paths a loop is always running, so
# this fires only from an unexpected sync call site.
logger.warning("No running event loop; skipping hook for %s", event)
return
task = loop.create_task(dispatch_hook(event, payload))
_background_tasks.add(task)
task.add_done_callback(_background_tasks.discard)
def has_pending_hooks() -> bool:
"""Return whether fire-and-forget hook tasks are still in flight."""
return any(not task.done() for task in _background_tasks)
async def drain_pending_hooks() -> None:
"""Await all in-flight fire-and-forget hook tasks.
Call this before the event loop tears down (e.g. at the end of a headless
run driven by `asyncio.run`) so background dispatches — most importantly the
final `tool.result` — are not cancelled mid-flight and silently dropped.
Each task's exceptions are already swallowed inside `dispatch_hook`, and any
stragglers are collected with `return_exceptions=True`, so this never
raises.
Precondition: this snapshots the in-flight set once and awaits it, so any
hook scheduled *after* the snapshot (during the await) is not drained. Call
it only once no further dispatches are possible. The headless caller invokes
it after `_run_agent_loop` has fully returned. The `app.py` graceful-exit
caller cancels the agent worker first, whose cancel handler
(`_handle_interrupt_cleanup`) schedules its terminal `tool.result` hooks
*synchronously* before this snapshot runs — see the ordering comment there —
so they are captured; a hook scheduled after a slow async write would not
be.
"""
# Snapshot: tasks remove themselves from the set via their done-callback as
# they finish, so iterating the live set while gathering would mutate it.
pending = list(_background_tasks)
if not pending:
return
await asyncio.gather(*pending, return_exceptions=True)
libs/code/deepagents_code/non_interactive.py
+322 -52
View File
@@ -40,6 +40,18 @@ from rich.style import Style
from rich.text import Text
from deepagents_code._cli_context import CLIContext
from deepagents_code._tool_stream import (
UNRENDERABLE_TOOL_OUTPUT,
ToolCallBuffer,
ToolCallBufferKey,
ToolStatus,
build_tool_error_payload,
build_tool_result_payload,
build_tool_use_payload,
count_unemitted_tool_calls,
normalize_tool_status,
tool_call_buffer_key,
)
from deepagents_code._version import __version__
from deepagents_code.agent import DEFAULT_AGENT_NAME
from deepagents_code.config import (
@@ -50,10 +62,15 @@ from deepagents_code.config import (
settings,
)
from deepagents_code.file_ops import FileOpTracker
from deepagents_code.hooks import dispatch_hook, dispatch_hook_fire_and_forget
from deepagents_code.hooks import (
dispatch_hook,
dispatch_hook_fire_and_forget,
drain_pending_hooks,
)
from deepagents_code.model_config import ModelConfigError
from deepagents_code.sessions import generate_thread_id
from deepagents_code.textual_adapter import SessionStats, print_usage_table
from deepagents_code.tool_display import format_tool_message_content
from deepagents_code.unicode_security import (
check_url_safety,
detect_dangerous_unicode,
@@ -230,6 +247,24 @@ async def _terminate_startup_process(proc: Process) -> None:
)
@dataclass(frozen=True)
class InFlightToolCall:
"""A tool call whose `tool.use` has fired but whose result has not arrived.
Bundling the name and args into one record keeps them structurally in
lock-step: a single `dict[str, InFlightToolCall]` cannot represent an id with
args but no name (or vice versa), which two parallel dicts could. That
removes the desync failure mode where an orphaned drain would emit a
`tool.error` with an empty `tool_name`.
"""
name: str
"""The tool name, carried so an orphaned call can be closed with a name."""
args: dict[str, Any]
"""The parsed tool-call arguments, for correlating the matching result."""
@dataclass
class StreamState:
"""Mutable state accumulated while iterating over the agent stream."""
@@ -249,11 +284,30 @@ class StreamState:
full_response: list[str] = field(default_factory=list)
"""Accumulated text fragments from the AI message stream."""
tool_call_buffers: dict[int | str, dict[str, str | None]] = field(
tool_call_buffers: dict[ToolCallBufferKey, ToolCallBuffer] = field(
default_factory=dict
)
"""Maps a tool-call index or ID to its name/ID metadata for in-progress
tool calls."""
"""Maps a tool-call index or ID to its in-progress buffer: name, ID,
accumulated argument fragments, and the display latch."""
in_flight_tool_calls: dict[str, InFlightToolCall] = field(default_factory=dict)
"""Maps in-flight tool-call IDs (those whose `tool.use` has fired) to their
name and parsed arguments, so the matching `tool.result` can be correlated.
Entries are removed when the result arrives; any still present when the
stream aborts are closed by `_dispatch_orphaned_tool_result_hooks`. One
record per id keeps name and args structurally in lock-step (see
`InFlightToolCall`)."""
displayed_tool_call_ids: set[str] = field(default_factory=set)
"""Tool-call IDs whose non-interactive call line has already been printed."""
emitted_tool_use_ids: set[str] = field(default_factory=set)
"""Tool-call IDs for which a `tool.use` has been dispatched.
Monotonic: never cleared within a run, so `tool.use` fires at most once per
id even if a call's arg chunks are redelivered after its result (mirrors the
TUI's `displayed_tool_ids`). Result correlation and the orphan drain use the
separate `in_flight_tool_calls`, which *is* cleared per result."""
pending_interrupts: dict[str, HITLRequest] = field(default_factory=dict)
"""Maps interrupt IDs to their validated HITL requests that are awaiting
@@ -412,28 +466,77 @@ def _process_ai_message(
chunk_name = block.get("name")
chunk_id = block.get("id")
chunk_index = block.get("index")
chunk_args = block.get("args")
if chunk_index is not None:
buffer_key: int | str = chunk_index
elif chunk_id is not None:
buffer_key = chunk_id
else:
buffer_key = f"unknown-{len(state.tool_call_buffers)}"
buffer_key = tool_call_buffer_key(
chunk_index, chunk_id, len(state.tool_call_buffers)
)
buffer = state.tool_call_buffers.setdefault(buffer_key, ToolCallBuffer())
buffer.ingest(name=chunk_name, tool_id=chunk_id, args=chunk_args)
if buffer_key not in state.tool_call_buffers:
state.tool_call_buffers[buffer_key] = {"name": None, "id": None}
if chunk_name:
state.tool_call_buffers[buffer_key]["name"] = chunk_name
buffer_name = buffer.name
buffer_id = buffer.tool_id
already_displayed = (
isinstance(buffer_id, str)
and buffer_id in state.displayed_tool_call_ids
)
if (
isinstance(buffer_name, str)
and not buffer.displayed
and not already_displayed
):
if state.spinner:
state.spinner.stop()
if state.quiet:
continue
if state.full_response:
_write_newline()
console.print(
f"[dim]🔧 Calling tool: {escape_markup(chunk_name)}[/dim]",
highlight=False,
if not state.quiet:
if state.full_response:
_write_newline()
console.print(
f"[dim]🔧 Calling tool: {escape_markup(buffer_name)}[/dim]",
highlight=False,
)
buffer.displayed = True
if isinstance(buffer_id, str):
state.displayed_tool_call_ids.add(buffer_id)
elif isinstance(buffer_id, str) and buffer.displayed:
state.displayed_tool_call_ids.add(buffer_id)
# Gate tool.use on a resolved tool id so this surface matches the
# interactive one, which dispatches at widget-mount time in
# `textual_adapter.execute_task_textual`. Both gate on a resolved
# tool-call id and fire at most once per id via a monotonic id set
# (`emitted_tool_use_ids` here, `displayed_tool_ids` there) that is
# never cleared within the run — see the "fire-once-per-id" clause of
# the parity contract in `_tool_stream`. Gating on the monotonic set
# rather than `in_flight_tool_calls` (which is cleared per result)
# means a redelivery of a completed call's arg chunks does not
# re-fire `tool.use` and spawn a spurious orphan.
parsed_args = buffer.parse_args()
if (
isinstance(buffer_name, str)
and buffer_id is not None
and parsed_args is not None
and buffer_id not in state.emitted_tool_use_ids
):
dispatch_hook_fire_and_forget(
"tool.use",
build_tool_use_payload(buffer_name, buffer_id, parsed_args),
)
state.emitted_tool_use_ids.add(buffer_id)
state.in_flight_tool_calls[buffer_id] = InFlightToolCall(
buffer_name, parsed_args
)
if (
isinstance(buffer_id, str)
and parsed_args is not None
and buffer_id in state.emitted_tool_use_ids
):
# Drop the buffer so a later turn that reuses this streaming
# index (indices restart per message, per LangChain streaming
# semantics) starts fresh rather than reusing this call's state.
# This also clears a redelivered completed call's recreated
# buffer, whose `tool.use` is already suppressed by the
# monotonic `emitted_tool_use_ids`.
state.tool_call_buffers.pop(buffer_key, None)
def _process_message_chunk(
@@ -471,6 +574,32 @@ def _process_message_chunk(
if isinstance(message_obj, AIMessage):
_process_ai_message(message_obj, state, console)
elif isinstance(message_obj, ToolMessage):
tool_id = getattr(message_obj, "tool_call_id", None)
correlated_tool_name = ""
# Args come from the matching tool.use. They default to {} when the call
# had no id to correlate on, or no tool.use fired (e.g. its args never
# parsed). The seam is intentional — without a correlated id we cannot
# pair them — but log the correlation miss so a lost pairing is not
# completely silent.
if isinstance(tool_id, str):
in_flight = state.in_flight_tool_calls.pop(tool_id, None)
tool_args = in_flight.args if in_flight is not None else None
correlated_tool_name = in_flight.name if in_flight is not None else ""
state.displayed_tool_call_ids.discard(tool_id)
if tool_args is None:
# Warning, not info/debug: a real-id result with no matching
# tool.use means a hook consumer sees a `tool.result` with empty
# args for a tool that actually executed (its args never parsed) —
# degraded audit fidelity worth surfacing at default log levels,
# consistent with the rest of this feature's severity philosophy.
logger.warning(
"tool.result for %s has no correlated tool.use args; "
"sending empty tool_args",
tool_id,
)
tool_args = {}
else:
tool_args = {}
record = file_op_tracker.complete_with_message(message_obj)
if record and record.diff:
if state.spinner:
@@ -480,6 +609,51 @@ def _process_message_chunk(
f"[dim]📝 {escape_markup(record.display_path)}[/dim]",
highlight=False,
)
tool_name = getattr(message_obj, "name", "")
if not tool_name:
tool_name = correlated_tool_name
# Normalize to the two-value hook domain, fail-closed: an unexpected
# provider status is logged and treated as an error (see
# `normalize_tool_status`) rather than silently reported as success.
tool_status: ToolStatus = normalize_tool_status(
getattr(message_obj, "status", "success"), tool_name
)
# Format the content the same way the interactive surface does so
# `tool_output` is identical across surfaces for list/structured content
# (e.g. multimodal or MCP tools returning content blocks) rather than a
# raw Python list repr here vs. extracted text there. Truncation to
# HOOK_TOOL_OUTPUT_LIMIT happens inside build_tool_result_payload.
# Guard formatting so a formatter error can't skip the tool.result
# dispatch below; on failure use a sentinel (see the except) rather than
# re-touching the offending content, so the dispatch stays unconditional.
try:
tool_content = format_tool_message_content(message_obj.content)
tool_output = str(tool_content) if tool_content else ""
except Exception:
# Guard formatting *and* the str() coercion together: a pathological
# __str__ must not re-raise past the fallback and skip the
# tool.result dispatch below. Use a sentinel rather than re-touching
# the offending content, so the dispatch stays unconditional.
logger.exception("Failed to format tool output")
tool_output = UNRENDERABLE_TOOL_OUTPUT
# Headless always dispatches tool.result for every ToolMessage — there
# are no widgets to skip. The TUI handles ToolMessages in three branches
# in `textual_adapter.execute_task_textual`: the widget-backed path and
# an `else` for unmounted tools both dispatch (mirroring this
# always-dispatch behavior), while the `completed_tool_result_ids` branch
# suppresses a duplicate rather than dispatching. See the parity contract
# in `_tool_stream` for the full guarantee.
if tool_status == "error":
dispatch_hook_fire_and_forget(
"tool.error",
build_tool_error_payload(tool_name),
)
dispatch_hook_fire_and_forget(
"tool.result",
build_tool_result_payload(
tool_name, tool_id, tool_args, tool_status, tool_output
),
)
if state.spinner:
state.spinner.start()
@@ -773,6 +947,45 @@ async def _stream_agent(
state.spinner.stop()
def _dispatch_orphaned_tool_result_hooks(state: StreamState, tool_output: str) -> None:
"""Close out `tool.use` events that never received a `ToolMessage`.
On a normally-completing run every `tool.use` is followed by a `ToolMessage`
that drains `in_flight_tool_calls`, so this is a no-op. When the stream is
aborted mid-flight (e.g. a provider error between the tool call and its
result), any id still present had its `tool.use` dispatched with no terminal
event; emit `tool.error` + a `tool_status="error"` `tool.result` for each so
the headless surface upholds the same "every `tool.use` is closed" guarantee
as the TUI's `_dispatch_terminal_tool_result_hooks`.
Args:
state: The stream state whose in-flight tool maps are drained.
tool_output: Terminal output recorded on each synthesized `tool.result`.
"""
if state.in_flight_tool_calls:
# A non-empty in-flight map here means real tool results were lost to a
# mid-stream abort (a clean run drains every id via its result). Surface
# it at warning — matching the TUI's backstop for the same class — so an
# operator can tell a clean run from one that synthesized error closes,
# rather than the drain being silent (degraded audit fidelity).
logger.warning(
"Stream ended with %d in-flight tool call(s) that never received a "
"result; closing each with a synthetic tool.error/tool.result",
len(state.in_flight_tool_calls),
)
for tool_id, in_flight in list(state.in_flight_tool_calls.items()):
dispatch_hook_fire_and_forget(
"tool.error", build_tool_error_payload(in_flight.name)
)
dispatch_hook_fire_and_forget(
"tool.result",
build_tool_result_payload(
in_flight.name, tool_id, in_flight.args, "error", tool_output
),
)
state.in_flight_tool_calls.clear()
async def _run_agent_loop(
agent: Any, # noqa: ANN401
message: str,
@@ -846,41 +1059,86 @@ async def _run_agent_loop(
start_time = time.monotonic()
# Initial stream
await _stream_agent(
agent, stream_input, config, state, console, file_op_tracker, context
)
# The internal default applies when --max-turns is omitted, guarding
# against unbounded runaway loops in scripts that forgot to set one.
effective_limit = max_turns if max_turns is not None else _MAX_HITL_ITERATIONS
# The initial stream above counts as turn 1; each HITL resume is a further
# turn. Raise before starting a resume that would exceed the budget so the
# user-facing count matches the flag's semantics.
turns = 1
while state.interrupt_occurred:
if turns >= effective_limit:
limit_source = (
f"--max-turns {max_turns}"
if max_turns is not None
else f"the internal safety default of {_MAX_HITL_ITERATIONS}"
)
msg = (
f"Exceeded {effective_limit} agentic turns ({limit_source}). "
"The agent may be stuck retrying rejected commands. "
"Increase --max-turns or break the task into smaller steps."
)
raise HITLIterationLimitError(msg)
turns += 1
state.interrupt_occurred = False
state.hitl_response.clear()
_process_hitl_interrupts(state, console)
stream_input = Command(resume=state.hitl_response)
try:
# Initial stream
await _stream_agent(
agent, stream_input, config, state, console, file_op_tracker, context
)
# The internal default applies when --max-turns is omitted, guarding
# against unbounded runaway loops in scripts that forgot to set one.
effective_limit = max_turns if max_turns is not None else _MAX_HITL_ITERATIONS
# The initial stream above counts as turn 1; each HITL resume is a
# further turn. Raise before starting a resume that would exceed the
# budget so the user-facing count matches the flag's semantics.
turns = 1
while state.interrupt_occurred:
if turns >= effective_limit:
limit_source = (
f"--max-turns {max_turns}"
if max_turns is not None
else f"the internal safety default of {_MAX_HITL_ITERATIONS}"
)
msg = (
f"Exceeded {effective_limit} agentic turns ({limit_source}). "
"The agent may be stuck retrying rejected commands. "
"Increase --max-turns or break the task into smaller steps."
)
raise HITLIterationLimitError(msg)
turns += 1
state.interrupt_occurred = False
state.hitl_response.clear()
_process_hitl_interrupts(state, console)
stream_input = Command(resume=state.hitl_response)
await _stream_agent(
agent, stream_input, config, state, console, file_op_tracker, context
)
finally:
# Close out any `tool.use` with no matching `ToolMessage` — e.g. a stream
# aborted by a provider error mid-tool. On a clean run every id was
# already drained by its result, so this is a no-op. Guarded so a
# dispatch problem can never mask the exception propagating from the
# stream (this runs on the error path too).
try:
_dispatch_orphaned_tool_result_hooks(
state, "Stream ended before tool result"
)
except Exception:
logger.warning(
"Orphaned tool.result drain failed unexpectedly", exc_info=True
)
# Surface any buffered tool call whose args never parsed: it never
# entered `in_flight_tool_calls` (so the orphan drain above skips it) and
# would otherwise be dropped with `state` at scope exit with no trace.
# Info, not warning — some of these may still have executed (their
# `tool.result` fired with `{}` args and logged a correlation miss); this
# only asserts the args never parsed. Guarded so a logging failure can
# never mask an exception propagating from the stream.
try:
# Two distinct reasons a buffered call never fired tool.use — args
# that never parsed, and args that parsed but carried no id (so
# tool.use was gated out). The classification is shared with the TUI
# via `count_unemitted_tool_calls`; each surface logs its own lines.
unemitted = count_unemitted_tool_calls(state.tool_call_buffers.values())
if unemitted.unparsed:
logger.info(
"Stream ended with %d tool call(s) whose arguments never "
"parsed; no tool.use was emitted for them",
unemitted.unparsed,
)
if unemitted.idless_parsed:
logger.info(
"Stream ended with %d tool call(s) whose arguments parsed but "
"carried no tool-call id; no tool.use was emitted for them",
unemitted.idless_parsed,
)
except Exception:
logger.warning(
"Unparsed tool-call buffer check failed unexpectedly",
exc_info=True,
)
wall_time = time.monotonic() - start_time
if state.full_response:
@@ -1379,3 +1637,15 @@ async def run_non_interactive(
return 1
else:
return 0
finally:
# Fire-and-forget hooks (tool.use/tool.result) run as background tasks;
# await them here so the final tool.result is not cancelled when
# asyncio.run tears the loop down. Never return from this block — that
# would swallow the exit code determined above. drain_pending_hooks is
# documented never to raise, but guard it anyway (mirroring app.py's
# shutdown drain) so a future contract break can't replace the exit code
# with an exception escaping the finally.
try:
await drain_pending_hooks()
except Exception:
logger.warning("Hook drain raised unexpectedly before exit", exc_info=True)
libs/code/deepagents_code/textual_adapter.py
+489 -113
View File
@@ -5,7 +5,6 @@ from __future__ import annotations
import asyncio
import contextlib
import json
import logging
import time
import uuid
@@ -56,10 +55,25 @@ from deepagents_code._session_stats import (
SpinnerStatus as SpinnerStatus,
format_token_count as format_token_count,
)
from deepagents_code._tool_stream import (
UNRENDERABLE_TOOL_OUTPUT,
ToolCallBuffer,
ToolCallBufferKey,
ToolStatus,
build_tool_error_payload,
build_tool_result_payload,
build_tool_use_payload,
count_unemitted_tool_calls,
normalize_tool_status,
tool_call_buffer_key,
)
from deepagents_code.config import build_stream_config, get_glyphs
from deepagents_code.file_ops import FileOpTracker
from deepagents_code.formatting import format_duration
from deepagents_code.hooks import dispatch_hook
from deepagents_code.hooks import (
dispatch_hook,
dispatch_hook_fire_and_forget,
)
from deepagents_code.input import MediaTracker, parse_file_mentions
from deepagents_code.media_utils import create_multimodal_content
from deepagents_code.tool_display import format_tool_message_content
@@ -79,6 +93,81 @@ _hitl_adapter_cache: TypeAdapter | None = None
_ASK_USER_UNSUPPORTED_ERROR = "ask_user not supported by this UI"
def _dispatch_tool_use_hook(
tool_name: str, tool_id: str, tool_args: dict[str, Any]
) -> None:
"""Dispatch a `tool.use` hook with the payload documented in `hooks`."""
dispatch_hook_fire_and_forget(
"tool.use", build_tool_use_payload(tool_name, tool_id, tool_args)
)
def _dispatch_tool_error_hook(tool_name: str) -> None:
"""Dispatch a `tool.error` hook with the payload documented in `hooks`."""
dispatch_hook_fire_and_forget("tool.error", build_tool_error_payload(tool_name))
def _dispatch_tool_result_hook(
tool_name: str,
tool_id: str | None,
tool_args: dict[str, Any],
tool_status: ToolStatus,
tool_output: str,
) -> None:
"""Dispatch a `tool.result` hook with the payload documented in `hooks`.
`tool_output` is truncated to `HOOK_TOOL_OUTPUT_LIMIT` inside the shared
payload builder.
"""
dispatch_hook_fire_and_forget(
"tool.result",
build_tool_result_payload(
tool_name, tool_id, tool_args, tool_status, tool_output
),
)
def _dispatch_terminal_tool_result_hooks(
tool_messages: dict[str, ToolCallMessage],
tool_output: str,
) -> list[str]:
"""Emit terminal `tool.error`/`tool.result` for still-pending tool widgets.
Every widget in `tool_messages` already had its `tool.use` dispatched (that
is when the widget is mounted), so any tool that reaches a terminal outcome
*without* a streamed `ToolMessage` — a HITL rejection, a cancelled turn, or
an aborted stream — would otherwise leave its `tool.use` unterminated. This
closes each one with a `tool_status="error"` result carrying the widget's
real `tool_name`/`args`, so the "every `tool.use` is closed by a matching
terminal event" guarantee holds on those paths too.
TUI-only: the headless surface reaches the equivalent state through
`_run_agent_loop`'s orphan drain rather than widgets.
Args:
tool_messages: Map of tool-call id to its widget for the pending tools.
tool_output: Terminal output string recorded on each `tool.result`.
Returns:
The tool-call ids that received terminal hooks. Callers track these
(via `completed_tool_result_ids`) so a later synthetic `ToolMessage`
— when the turn still resumes, e.g. alongside an answered `ask_user`
— does not double-dispatch.
"""
dispatched: list[str] = []
for tool_id, tool_msg in list(tool_messages.items()):
_dispatch_tool_error_hook(tool_msg.tool_name)
_dispatch_tool_result_hook(
tool_msg.tool_name,
tool_id,
tool_msg.args,
"error",
tool_output,
)
dispatched.append(tool_id)
return dispatched
def _get_hitl_request_adapter(hitl_request_type: type) -> TypeAdapter:
"""Return a cached `TypeAdapter(HITLRequest)`.
@@ -370,6 +459,11 @@ class TextualUIAdapter:
Args:
error: Error text to display in each pending tool widget.
"""
# Each pending widget already had its `tool.use` dispatched at mount, so
# emit terminal hooks before dropping them — otherwise an aborted stream
# leaves those `tool.use` events unterminated for audit consumers. Runs
# before the widget updates so a `set_error` failure can't skip it.
_dispatch_terminal_tool_result_hooks(self._current_tool_messages, error)
for tool_msg in list(self._current_tool_messages.values()):
tool_msg.set_error(error)
self._current_tool_messages.clear()
@@ -620,7 +714,12 @@ async def execute_task_textual(
file_op_tracker = FileOpTracker(assistant_id=assistant_id, backend=backend)
displayed_tool_ids: set[str] = set()
tool_call_buffers: dict[str | int, dict] = {}
tool_call_buffers: dict[ToolCallBufferKey, ToolCallBuffer] = {}
# Tool-call ids that already received terminal hooks before a resumed
# `ToolMessage` can stream. When the turn still resumes, middleware
# synthetic messages would otherwise re-dispatch `tool.result`; this set
# suppresses those duplicates.
completed_tool_result_ids: set[str] = set()
# Track pending text and assistant messages PER NAMESPACE to avoid interleaving
# when multiple subagents stream in parallel
@@ -779,23 +878,55 @@ async def execute_task_textual(
if tool_id not in displayed_tool_ids:
if adapter._set_spinner:
await adapter._set_spinner(None)
tool_args = {
"questions": validated_ask_user[
"questions"
]
}
tool_msg = ToolCallMessage(
"ask_user",
{
"questions": validated_ask_user[
"questions"
]
},
tool_args,
)
try:
await adapter._mount_message(tool_msg)
except Exception:
# Mount failed (e.g. a torn-down
# DOM during shutdown). tool.use
# is dispatched only on mount
# success (below), so a failed
# mount leaves no unterminated
# tool.use to orphan if the turn
# is then cancelled before the
# ask_user resolution loop runs.
# The id is left unlatched so a
# re-observed interrupt can retry
# the mount; the question is still
# asked and closed by the
# resolution loop, which
# dispatches the terminal
# tool.result independently of
# this widget.
logger.exception(
"Failed to mount ask_user "
"tool row for %s",
tool_id,
)
else:
# Fire tool.use and latch the id
# together, only once the widget
# is mounted, so the "every
# tool.use is closed" guarantee
# holds with no widget-less orphan
# on the mount-failure path.
# Gating on mount success also
# keeps tool.use fire-once: a
# failed mount never fires it, and
# a successful mount latches the
# id so a re-observed interrupt is
# skipped.
_dispatch_tool_use_hook(
"ask_user", tool_id, tool_args
)
displayed_tool_ids.add(tool_id)
adapter._current_tool_messages[
tool_id
@@ -900,31 +1031,96 @@ async def execute_task_textual(
if isinstance(message, ToolMessage):
tool_name = getattr(message, "name", "")
tool_status = getattr(message, "status", "success")
tool_content = format_tool_message_content(message.content)
# Normalize to the two-value hook domain, fail-closed: an
# unexpected provider status is logged and treated as an
# error (see `normalize_tool_status`) rather than silently
# reported as success.
tool_status: ToolStatus = normalize_tool_status(
getattr(message, "status", "success"), tool_name
)
# Guard formatting *and* the str() coercion so a
# pathological __str__ on the content can't re-raise and
# skip the tool.result dispatch below. On failure use a
# sentinel rather than re-touching the offending content,
# so the terminal dispatch is genuinely unconditional.
try:
tool_content = format_tool_message_content(message.content)
output_str = str(tool_content) if tool_content else ""
except Exception:
logger.exception("Failed to format tool output")
output_str = UNRENDERABLE_TOOL_OUTPUT
record = file_op_tracker.complete_with_message(message)
# Update tool call status with output
tool_id = getattr(message, "tool_call_id", None)
if tool_id and tool_id in adapter._current_tool_messages:
# Pop before widget calls so the dict drains even
# Pop before the widget calls so the dict drains even
# if set_success/set_error raises.
tool_msg = adapter._current_tool_messages.pop(tool_id)
output_str = str(tool_content) if tool_content else ""
if tool_status == "success":
tool_msg.set_success(output_str)
else:
tool_msg.set_error(output_str or "Error")
await dispatch_hook(
"tool.error",
{"tool_names": [tool_msg._tool_name]},
)
elif tool_id:
logger.debug(
"ToolMessage tool_call_id=%s not in "
"_current_tool_messages; spinner gating "
"may be stale",
# Dispatch the terminal hooks *before* touching the
# widget: a render failure must never drop this tool's
# tool.result/tool.error (which would leave its
# tool.use unterminated). The headless path likewise
# dispatches without depending on any widget.
if tool_status == "error":
_dispatch_tool_error_hook(tool_msg.tool_name)
_dispatch_tool_result_hook(
tool_msg.tool_name,
tool_id,
tool_msg.args,
tool_status,
output_str,
)
# Update the widget last, guarded: a set_success/
# set_error failure must not abort the turn and drop
# the remaining tools' hooks.
try:
if tool_status == "success":
tool_msg.set_success(output_str)
else:
tool_msg.set_error(output_str or "Error")
except Exception:
logger.exception(
"Failed to update tool row for %s", tool_id
)
elif tool_id and tool_id in completed_tool_result_ids:
# This is a middleware synthetic ToolMessage for a
# tool whose terminal hooks already fired while the
# turn was resolving interrupts. Its widget was
# cleared, so it lands here — consume the id and skip
# re-dispatch to avoid a duplicate tool.result (with
# mismatched `{}` args).
completed_tool_result_ids.discard(tool_id)
else:
# The tool call was never mounted — either it has no
# tool_call_id, or its streamed args never parsed so
# no tool.use fired and no widget exists. Still emit
# tool.result (with {} args, since without a widget
# we lack the parsed args) so audit hooks observe
# every executed tool, matching the headless path.
# tool_id may be None here, mirroring headless.
# Reciprocal: headless always dispatches tool.result
# from `_process_message_chunk` since it has no
# widget concept; see `non_interactive.py`. The
# parity contract is documented in `_tool_stream`.
if tool_id:
# Warning, not info/debug: a real-id result with
# no mounted widget (its args never parsed, so no
# tool.use fired) means a hook consumer sees a
# `tool.result` with empty args for a tool that
# actually executed — degraded audit fidelity worth
# surfacing at default log levels, matching the
# headless path.
logger.warning(
"ToolMessage tool_call_id=%s not in "
"_current_tool_messages; no correlated "
"tool.use, sending empty tool_args",
tool_id,
)
if tool_status == "error":
_dispatch_tool_error_hook(tool_name)
_dispatch_tool_result_hook(
tool_name, tool_id, {}, tool_status, output_str
)
# Show file operation results - always show diffs in chat
@@ -1061,84 +1257,33 @@ async def execute_task_textual(
chunk_id = block.get("id")
chunk_index = block.get("index")
buffer_key: str | int
if chunk_index is not None:
buffer_key = chunk_index
elif chunk_id is not None:
buffer_key = chunk_id
else:
buffer_key = f"unknown-{len(tool_call_buffers)}"
buffer_key = tool_call_buffer_key(
chunk_index, chunk_id, len(tool_call_buffers)
)
buffer = tool_call_buffers.setdefault(
buffer_key,
{
"name": None,
"id": None,
"args": None,
"args_parts": [],
},
buffer_key, ToolCallBuffer()
)
buffer.ingest(
name=chunk_name, tool_id=chunk_id, args=chunk_args
)
if chunk_name:
buffer["name"] = chunk_name
if chunk_id:
buffer["id"] = chunk_id
if isinstance(chunk_args, dict):
buffer["args"] = chunk_args
buffer["args_parts"] = []
elif isinstance(chunk_args, str):
if chunk_args:
parts: list[str] = buffer.setdefault(
"args_parts", []
)
if not parts or chunk_args != parts[-1]:
parts.append(chunk_args)
elif chunk_args is not None:
buffer["args"] = chunk_args
buffer_name = buffer.get("name")
buffer_id = buffer.get("id")
buffer_name = buffer.name
buffer_id = buffer.tool_id
if buffer_name is None:
continue
# Resolve the tool arguments. String fragments are
# accumulated in `args_parts` and joined + parsed
# once the buffer holds a complete JSON value. Re-
# joining and re-parsing the whole prefix on every
# fragment is O(n^2) and ran on the UI event loop for
# large `edit_file` blobs. Each `continue` below
# leaves the buffer in `tool_call_buffers` so the next
# `parse_args` reassembles streamed JSON string
# fragments, deferring the parse until the value
# looks complete — which avoids re-parsing the whole
# prefix on every fragment (costly on the UI event
# loop for large `edit_file` blobs) — and returns
# None while still incomplete. Each `continue` leaves
# the buffer in `tool_call_buffers` so the next
# fragment keeps accumulating; it is popped only after
# a successful parse + mount.
direct_args = buffer.get("args")
if isinstance(direct_args, dict):
parsed_args = direct_args
elif direct_args is not None:
parsed_args = {"value": direct_args}
else:
parts = buffer.get("args_parts") or []
if not parts:
continue
joined = "".join(parts)
stripped = joined.strip()
if not stripped:
continue
# Objects/arrays can be large (e.g. `edit_file`
# blobs), so defer parsing until the closing
# bracket arrives. Scalars are always small and
# never end in `}`/`]`, so parse them eagerly
# rather than leaving them stuck unparsed.
if stripped[0] in "{[" and not stripped.endswith(
("}", "]")
):
continue
try:
parsed_args = json.loads(joined)
except json.JSONDecodeError:
continue
if not isinstance(parsed_args, dict):
parsed_args = {"value": parsed_args}
# a successful parse + mount below.
parsed_args = buffer.parse_args()
if parsed_args is None:
continue
# Flush pending text before tool call
pending_text = pending_text_by_namespace.get(ns_key, "")
@@ -1183,16 +1328,43 @@ async def execute_task_textual(
buffer_name,
repr(parsed_args)[:200],
)
# Dispatch tool.use once the streamed call has a
# resolved id and parsed args. The headless
# surface dispatches from the stream loop
# instead; see the "Gate tool.use" comment in
# `non_interactive._process_ai_message`. Both
# gate on a resolved tool-call id and fire at
# most once per id — the parity contract is
# documented in `_tool_stream`.
_dispatch_tool_use_hook(
buffer_name, buffer_id, parsed_args
)
tool_msg = ToolCallMessage(buffer_name, parsed_args)
await adapter._mount_message(tool_msg)
try:
await adapter._mount_message(tool_msg)
except Exception:
# tool.use already fired. If the mount raises
# (e.g. mounting into a torn-down DOM during
# shutdown), still track the pending call so
# the later real ToolMessage remains
# authoritative for tool.result status/output.
# If the stream ends first, the terminal
# drains close this tool.use from the same
# pending map.
logger.exception(
"Failed to mount tool widget for %s",
buffer_id,
)
else:
# Mark running so the group row reflects live
# progress; the row itself is hidden inside
# the group, so this drives state, not a
# visible per-tool spinner.
tool_msg.set_running()
adapter._current_tool_messages[buffer_id] = tool_msg
# Mark running so the group row reflects live
# progress; the row itself is hidden inside the
# group, so this drives state, not a visible
# per-tool spinner.
tool_msg.set_running()
tool_call_buffers.pop(buffer_key, None)
if buffer_id is not None:
tool_call_buffers.pop(buffer_key, None)
if getattr(message, "chunk_position", None) == "last":
pending_text = pending_text_by_namespace.get(ns_key, "")
@@ -1253,6 +1425,7 @@ async def execute_task_textual(
for interrupt_id, ask_req in list(pending_ask_user.items()):
questions = ask_req["questions"]
tool_args = {"questions": questions}
if adapter._request_ask_user:
if adapter._set_spinner:
@@ -1306,11 +1479,22 @@ async def execute_task_textual(
answers = result.get("answers", [])
if isinstance(answers, list):
resume_payload[interrupt_id] = {"answers": answers}
output = "User answered"
tool_msg = adapter._current_tool_messages.pop(
tool_id, None
)
_dispatch_tool_result_hook(
"ask_user", tool_id, tool_args, "success", output
)
completed_tool_result_ids.add(tool_id)
if tool_msg is not None:
tool_msg.set_success("User answered")
try:
tool_msg.set_success(output)
except Exception:
logger.exception(
"Failed to update ask_user row for %s",
tool_id,
)
else:
logger.warning(
"ask_user tool_id %s missing from "
@@ -1329,13 +1513,23 @@ async def execute_task_textual(
"answers": ["" for _ in questions],
}
any_rejected = True
output = "invalid ask_user answers payload"
tool_msg = adapter._current_tool_messages.pop(
tool_id, None
)
_dispatch_tool_error_hook("ask_user")
_dispatch_tool_result_hook(
"ask_user", tool_id, tool_args, "error", output
)
completed_tool_result_ids.add(tool_id)
if tool_msg is not None:
tool_msg.set_error(
"invalid ask_user answers payload"
)
try:
tool_msg.set_error(output)
except Exception:
logger.exception(
"Failed to update ask_user row for %s",
tool_id,
)
elif result_type == "cancelled":
resume_payload[interrupt_id] = {
"status": "cancelled",
@@ -1346,8 +1540,20 @@ async def execute_task_textual(
# resume so the agent can react to the failure.
ask_user_cancelled = True
tool_msg = adapter._current_tool_messages.pop(tool_id, None)
output = "Question cancelled"
_dispatch_tool_error_hook("ask_user")
_dispatch_tool_result_hook(
"ask_user", tool_id, tool_args, "error", output
)
completed_tool_result_ids.add(tool_id)
if tool_msg is not None:
tool_msg.set_rejected()
try:
tool_msg.set_rejected()
except Exception:
logger.exception(
"Failed to update ask_user row for %s",
tool_id,
)
else:
logger.warning(
"ask_user tool_id %s missing from "
@@ -1365,8 +1571,19 @@ async def execute_task_textual(
}
any_rejected = True
tool_msg = adapter._current_tool_messages.pop(tool_id, None)
_dispatch_tool_error_hook("ask_user")
_dispatch_tool_result_hook(
"ask_user", tool_id, tool_args, "error", error_text
)
completed_tool_result_ids.add(tool_id)
if tool_msg is not None:
tool_msg.set_error(error_text)
try:
tool_msg.set_error(error_text)
except Exception:
logger.exception(
"Failed to update ask_user row for %s",
tool_id,
)
else:
logger.warning(
"ask_user interrupt received but no UI callback is "
@@ -1379,8 +1596,22 @@ async def execute_task_textual(
}
tool_id = ask_req["tool_call_id"]
tool_msg = adapter._current_tool_messages.pop(tool_id, None)
_dispatch_tool_error_hook("ask_user")
_dispatch_tool_result_hook(
"ask_user",
tool_id,
tool_args,
"error",
_ASK_USER_UNSUPPORTED_ERROR,
)
completed_tool_result_ids.add(tool_id)
if tool_msg is not None:
tool_msg.set_error(_ASK_USER_UNSUPPORTED_ERROR)
try:
tool_msg.set_error(_ASK_USER_UNSUPPORTED_ERROR)
except Exception:
logger.exception(
"Failed to update ask_user row for %s", tool_id
)
for interrupt_id, hitl_request in list(pending_interrupts.items()):
action_requests = hitl_request["action_requests"]
@@ -1517,7 +1748,6 @@ async def execute_task_textual(
)
for tool_msg in tool_msgs:
tool_msg.set_rejected(reason=reject_message)
adapter._current_tool_messages.clear()
# Bare reject aborts the turn and shows the
# canned "Command rejected" banner so the user
# can redirect. When a reason is supplied, the
@@ -1525,6 +1755,13 @@ async def execute_task_textual(
# agent: keep `any_rejected=False` so the
# stream resumes and the banner is suppressed.
if reject_message is None:
completed_tool_result_ids.update(
_dispatch_terminal_tool_result_hooks(
adapter._current_tool_messages,
"Tool approval rejected",
)
)
adapter._current_tool_messages.clear()
any_rejected = True
else:
logger.warning(
@@ -1539,6 +1776,12 @@ async def execute_task_textual(
adapter._current_tool_messages.values()
):
tool_msg.set_rejected()
completed_tool_result_ids.update(
_dispatch_terminal_tool_result_hooks(
adapter._current_tool_messages,
"Tool approval rejected",
)
)
adapter._current_tool_messages.clear()
any_rejected = True
else:
@@ -1553,6 +1796,12 @@ async def execute_task_textual(
adapter._current_tool_messages.values()
):
tool_msg.set_rejected()
completed_tool_result_ids.update(
_dispatch_terminal_tool_result_hooks(
adapter._current_tool_messages,
"Tool approval rejected",
)
)
adapter._current_tool_messages.clear()
any_rejected = True
@@ -1586,6 +1835,32 @@ async def execute_task_textual(
stream_input = Command(resume=resume_payload)
else:
# Clean stream end. Any tool still in `_current_tool_messages`
# had its `tool.use` dispatched at mount but never received a
# `ToolMessage` (e.g. a custom/remote graph that ends the turn
# after emitting an unexecuted tool call). Close each one with a
# terminal hook so the "every `tool.use` is terminated" guarantee
# does not depend on the graph raising. This mirrors the headless
# `_dispatch_orphaned_tool_result_hooks`, which likewise closes
# orphans hooks-only (no widget mutation) on every loop exit —
# the widget keeps its rendered state; only the audit stream and
# the cross-turn `_current_tool_messages` tracking are settled.
if adapter._current_tool_messages:
logger.info(
"Stream ended with %d un-resulted tool call(s); "
"closing with terminal hooks",
len(adapter._current_tool_messages),
)
_dispatch_terminal_tool_result_hooks(
adapter._current_tool_messages,
"Stream ended before tool result",
)
adapter._current_tool_messages.clear()
# The end-of-stream diagnostic for buffered tool calls that never
# fired a `tool.use` runs in the `finally` below, not here, so it
# fires on cancel and mid-stream error too (not only this clean
# end) — mirroring the headless surface, whose identical
# diagnostic lives in `_run_agent_loop`'s `finally`.
await dispatch_hook("task.complete", {"thread_id": thread_id})
break
@@ -1615,6 +1890,67 @@ async def execute_task_textual(
except Exception: # drain must not mask the original error
logger.exception("Failed to drain assistant streams on exit")
# Self-contained backstop for the "every `tool.use` is terminated" hook
# guarantee. The clean-end branch, HITL-reject branches, and interrupt
# cleanup each already drained `_current_tool_messages` and cleared it, so
# this is a no-op on those paths. The one path it covers is a non-cancel
# mid-stream error propagating to the caller: without it, the tools that
# fired `tool.use` would be terminated only by the caller's
# `finalize_pending_tools_with_error`, leaving the hook guarantee dependent
# on the caller rather than owned here (a future second caller, or a
# missing adapter, would leak an unterminated `tool.use`). Runs before the
# exception reaches the caller, whose `finalize_pending_tools_with_error`
# then finds an empty dict and no-ops, so no `tool.result` is dispatched
# twice. Fail-loud and guarded so a dispatch problem can never mask the
# error propagating from the stream.
if adapter._current_tool_messages:
logger.warning(
"Turn exited with %d un-terminated tool call(s); closing with "
"terminal hooks as a backstop",
len(adapter._current_tool_messages),
)
try:
adapter.finalize_pending_tools_with_error(
"Agent error before tool result"
)
except Exception:
logger.warning(
"Backstop terminal tool close failed unexpectedly",
exc_info=True,
)
# Surface any buffered tool call that never mounted and never fired a
# `tool.use`, so it would otherwise vanish with `tool_call_buffers` at turn
# end with no trace. Two distinct cases (args that never parsed, and args
# that parsed but carried no tool-call id) are classified by the shared
# `count_unemitted_tool_calls`. In the `finally` so it fires on every exit
# path — clean end, cancel, and mid-stream error — matching the headless
# surface. Info, not warning: nothing executed for these and the
# precondition (exiting mid-tool-call) is unusual; it only needs to be
# greppable. Guarded so a logging failure can never mask a propagating
# exception (`parse_args`, re-run inside the count, can raise on the
# invariant-violating both-fields-set buffer).
try:
unemitted = count_unemitted_tool_calls(tool_call_buffers.values())
if unemitted.unparsed:
logger.info(
"Stream ended with %d tool call(s) whose arguments never "
"parsed; no tool.use was emitted for them",
unemitted.unparsed,
)
if unemitted.idless_parsed:
logger.info(
"Stream ended with %d tool call(s) whose arguments parsed "
"but carried no tool-call id; no tool.use was emitted for "
"them",
unemitted.idless_parsed,
)
except Exception:
logger.warning(
"Unparsed tool-call buffer check failed unexpectedly",
exc_info=True,
)
# Update token count and return stats. Persistence is handled inside the
# graph by `ResumeStateMiddleware.after_model`, so this only refreshes UI.
turn_stats.wall_time_seconds = time.monotonic() - start_time
@@ -1722,6 +2058,35 @@ async def _handle_interrupt_cleanup(
adapter._current_tool_messages,
)
# Close out any tool whose `tool.use` fired but whose `ToolMessage` never
# arrived because the turn was cancelled: emit terminal hooks before the
# widgets are dropped, so a cancel path leaves no unterminated `tool.use`
# (mirroring the HITL-reject branches). The turn does not resume from here,
# so the returned ids need not be tracked for dedup.
#
# Dispatched *before* the `aupdate_state` writes below (not alongside the
# `set_rejected` loop after them): those writes await a possibly-slow remote
# checkpointer, and on an interactive quit the graceful-exit drain in
# `app.py` snapshots the in-flight hook tasks right after cancelling this
# worker. Scheduling the fire-and-forget hooks here — synchronously, as soon
# as cancellation is observed — guarantees they are in that snapshot and get
# drained, rather than being scheduled after a slow write and cancelled at
# loop teardown (a silent audit gap). It reads `tool_msg.args`/`tool_name`,
# both available regardless of the widget's rejected state.
#
# Guarded because this now sits *before* the recovery-state write below: the
# dispatch never raises by construction today (pure payload builders, and
# `dispatch_hook_fire_and_forget` swallows serialization inside its task), but
# this function's whole contract is best-effort-must-not-propagate, so a
# future change here must never skip the `aupdate_state` save or escape the
# cancel handler.
try:
_dispatch_terminal_tool_result_hooks(
adapter._current_tool_messages, "Turn cancelled"
)
except Exception:
logger.warning("Terminal tool.result dispatch failed on cancel", exc_info=True)
# Save accumulated state before marking tools as rejected (best-effort).
# State update failures shouldn't prevent cleanup.
from langsmith import tracing_context
@@ -1766,9 +2131,20 @@ async def _handle_interrupt_cleanup(
)
)
# Mark tools as rejected AFTER saving state
# Mark tools as rejected AFTER saving state. Terminal hooks for these were
# already dispatched before the state writes above (see the comment there).
# Guard each `set_rejected` — it does DOM work that can raise during
# app-exit teardown — so a failure can't skip the `clear()` below. If it
# did, `_current_tool_messages` would stay populated and the caller's
# `finally` backstop would re-dispatch a duplicate terminal hook for every
# id already closed at the top of this function.
for tool_msg in list(adapter._current_tool_messages.values()):
tool_msg.set_rejected()
try:
tool_msg.set_rejected()
except Exception:
logger.exception(
"Failed to mark tool row rejected during interrupt cleanup"
)
adapter._current_tool_messages.clear()
# Keep the token count marked stale whenever interrupted state was captured,
libs/code/deepagents_code/widgets/messages.py
+25
View File
@@ -1353,6 +1353,13 @@ class ToolCallMessage(Vertical):
Args:
result: Tool output/result to display
"""
if self._status in {"rejected", "skipped"}:
# A rejected tool (or one skipped due to a sibling rejection) never
# legitimately becomes successful. A resumed turn can still stream a
# synthetic ToolMessage for such a tool (see the reasoned-reject path
# in `textual_adapter`); ignore it so the row keeps its terminal
# rejected/skipped state instead of flipping.
return
elapsed = time() - self._start_time if self._start_time is not None else None
self._stop_animation()
self._status = "success"
@@ -1401,6 +1408,13 @@ class ToolCallMessage(Vertical):
Args:
error: Error message
"""
if self._status in {"rejected", "skipped"}:
# A rejected/skipped tool never legitimately errors. A resumed turn
# can stream a synthetic error ToolMessage for a reasoned-reject tool
# (see `textual_adapter`); ignore it so the row keeps its rejected
# state rather than flipping to "Error" (which also left the stale
# `rejected` CSS class alongside `error`).
return
self._stop_animation()
self._status = "error"
# For shell commands, prepend the full command so users can see what failed
@@ -2635,6 +2649,17 @@ class ToolCallMessage(Vertical):
"""Public read-only accessor for the underlying tool name."""
return self._tool_name
@property
def args(self) -> dict[str, Any]:
"""Public read-only accessor for the parsed tool-call arguments.
Returns a shallow copy so a consumer (e.g. a hook payload built from
`args`) cannot rebind the widget's top-level keys by reference. Nested
mutable values are shared, not deep-copied, so callers must treat them as
read-only and must not deep-mutate a returned nested value.
"""
return dict(self._args)
@property
def is_success(self) -> bool:
"""Whether the tool completed successfully."""
libs/code/tests/unit_tests/test__tool_stream.py
+503
View File
@@ -0,0 +1,503 @@
"""Tests for the shared streaming tool-call buffer and hook-payload builders.
`_tool_stream` is the single source of truth for reassembling streamed tool-call
arguments and building `tool.use` / `tool.result` / `tool.error` payloads across
both execution surfaces, so its contract is exercised directly here (the two
surfaces additionally exercise it end-to-end in their own suites).
"""
from __future__ import annotations
import pytest
from deepagents_code._tool_stream import (
TOOL_OUTPUT_TRUNCATION_MARKER,
ToolCallBuffer,
build_tool_error_payload,
build_tool_result_payload,
build_tool_use_payload,
count_unemitted_tool_calls,
normalize_tool_status,
tool_call_buffer_key,
)
from deepagents_code.hooks import HOOK_TOOL_OUTPUT_LIMIT
class TestToolCallBufferKey:
"""Precedence of the buffer key: index, then id, then placeholder."""
def test_index_preferred_over_id(self) -> None:
"""A present index wins even when an id is also available."""
assert tool_call_buffer_key(0, "toolu_1", 3) == 0
def test_falls_back_to_id_when_no_index(self) -> None:
"""The tool id keys the buffer when no streaming index is present."""
assert tool_call_buffer_key(None, "toolu_1", 3) == "toolu_1"
def test_placeholder_when_no_index_or_id(self) -> None:
"""A count-based placeholder keeps unrelated id-less calls distinct."""
assert tool_call_buffer_key(None, None, 2) == "unknown-2"
assert tool_call_buffer_key(None, None, 5) == "unknown-5"
def test_zero_index_is_not_treated_as_missing(self) -> None:
"""Index 0 is a valid key, not a falsy stand-in for absent."""
assert tool_call_buffer_key(0, None, 0) == 0
class TestToolCallBufferConstruction:
"""The `args` XOR `args_parts` invariant is enforced at construction."""
def test_both_args_and_args_parts_rejected(self) -> None:
"""Constructing with both populated raises, making the state unrepresentable.
`ingest` maintains the XOR on every chunk; `__post_init__` guarantees no
buffer can be built in the illegal both-set state that `parse_args` would
otherwise mask via read order.
"""
with pytest.raises(ValueError, match="cannot hold both"):
ToolCallBuffer(args={"a": 1}, args_parts=['{"a":'])
def test_only_args_allowed(self) -> None:
"""A whole value alone is a legal buffer."""
assert ToolCallBuffer(args={"a": 1}).args == {"a": 1}
def test_only_args_parts_allowed(self) -> None:
"""Fragments alone are a legal buffer."""
assert ToolCallBuffer(args_parts=['{"a":']).args_parts == ['{"a":']
class TestToolCallBufferIngest:
"""Folding streamed chunk fields into the buffer."""
def test_name_and_id_captured(self) -> None:
"""Name and id from a chunk populate the buffer."""
buffer = ToolCallBuffer()
buffer.ingest(name="write_file", tool_id="toolu_1", args=None)
assert buffer.name == "write_file"
assert buffer.tool_id == "toolu_1"
def test_dict_args_replace_accumulated_fragments(self) -> None:
"""A whole-value dict delivery discards any partial fragments."""
buffer = ToolCallBuffer(args_parts=['{"partial": '])
buffer.ingest(name=None, tool_id=None, args={"path": "foo.py"})
assert buffer.args == {"path": "foo.py"}
assert buffer.args_parts == []
def test_string_fragments_accumulate(self) -> None:
"""String fragments append in order."""
buffer = ToolCallBuffer()
buffer.ingest(name=None, tool_id=None, args='{"a":')
buffer.ingest(name=None, tool_id=None, args=" 1}")
assert buffer.args_parts == ['{"a":', " 1}"]
def test_empty_string_fragment_skipped(self) -> None:
"""An empty fragment carries no payload and is not appended."""
buffer = ToolCallBuffer()
buffer.ingest(name=None, tool_id=None, args="")
assert buffer.args_parts == []
def test_adjacent_identical_fragments_preserved(self) -> None:
"""Repeated identical fragments are real content, not de-duplicated."""
buffer = ToolCallBuffer()
for fragment in ('{"content": "', "hi", "hi", '"}'):
buffer.ingest(name=None, tool_id=None, args=fragment)
assert buffer.parse_args() == {"content": "hihi"}
def test_non_string_non_dict_scalar_stored(self) -> None:
"""A non-None scalar that is neither dict nor str is stored as-is."""
buffer = ToolCallBuffer()
buffer.ingest(name=None, tool_id=None, args=7)
assert buffer.args == 7
def test_same_tool_id_keeps_accumulating(self) -> None:
"""Repeated chunks of one call (same id, then id-less) keep appending."""
buffer = ToolCallBuffer()
buffer.ingest(name="write_file", tool_id="toolu_1", args='{"a":')
buffer.ingest(name=None, tool_id="toolu_1", args=" 1}")
buffer.ingest(name=None, tool_id=None, args="")
assert buffer.parse_args() == {"a": 1}
def test_new_tool_id_resets_stale_call_state(self) -> None:
"""A differing id (reused streaming index) discards old call state.
Indices restart per message, so a buffer retained from an earlier call
(e.g. one whose args never parsed) can be handed to a new call via the
same key. The new id must reset the old call's arguments and metadata so
they cannot leak into chunks for the new call.
"""
buffer = ToolCallBuffer(
name="read_file",
tool_id="toolu_a",
args_parts=["{bad"],
displayed=True,
)
buffer.ingest(name=None, tool_id="toolu_b", args='{"x": 1}')
assert buffer.tool_id == "toolu_b"
assert buffer.name is None
assert buffer.displayed is False
assert buffer.parse_args() == {"x": 1}
def test_idless_named_chunk_resets_stale_tool_id(self) -> None:
"""A delayed-id new call must not parse under the prior call's id.
A retained malformed buffer can be reused by the next call's stream index
before the provider emits the new id. The new name is the first boundary
signal available, so it must clear the old id before args are parsed.
"""
buffer = ToolCallBuffer(
name="read_file",
tool_id="toolu_a",
args_parts=["{bad"],
displayed=True,
)
buffer.ingest(name="write_file", tool_id=None, args='{"x": 1}')
assert buffer.tool_id is None
assert buffer.name == "write_file"
assert buffer.displayed is False
assert buffer.parse_args() == {"x": 1}
def test_delayed_tool_id_attaches_to_idless_new_call_args(self) -> None:
"""The real id attaches without discarding args collected before it."""
buffer = ToolCallBuffer(
name="read_file",
tool_id="toolu_a",
args_parts=["{bad"],
displayed=True,
)
buffer.ingest(name="write_file", tool_id=None, args='{"x": 1}')
buffer.ingest(name=None, tool_id="toolu_b", args="")
assert buffer.tool_id == "toolu_b"
assert buffer.name == "write_file"
assert buffer.parse_args() == {"x": 1}
def test_new_tool_id_resets_warned_latch(self) -> None:
"""A reused index resets the `warned` latch.
The new call's own malformed payload is still surfaced once, rather than
being suppressed by the previous call's latch.
"""
buffer = ToolCallBuffer(tool_id="toolu_a", args_parts=["{bad json}"])
assert buffer.parse_args() is None # sets warned for call a
assert buffer.warned is True
buffer.ingest(name="write_file", tool_id="toolu_b", args="{also bad}")
assert buffer.warned is False
def test_string_fragment_clears_prior_whole_value(self) -> None:
"""A fragment supersedes a prior whole value.
Enforces the `args` XOR `args_parts` invariant at write time rather than
leaving both populated for `parse_args` read-order to disambiguate.
"""
buffer = ToolCallBuffer(args={"stale": True})
buffer.ingest(name=None, tool_id=None, args='{"real":')
assert buffer.args is None
assert buffer.args_parts == ['{"real":']
def test_scalar_clears_prior_fragments(self) -> None:
"""A whole scalar value discards accumulated fragments (XOR invariant)."""
buffer = ToolCallBuffer(args_parts=['{"partial":'])
buffer.ingest(name=None, tool_id=None, args=7)
assert buffer.args == 7
assert buffer.args_parts == []
class TestNormalizeToolStatus:
"""Fail-closed mapping of a raw `ToolMessage.status` to the hook domain."""
def test_success_passthrough(self) -> None:
assert normalize_tool_status("success", "read_file") == "success"
def test_error_passthrough(self) -> None:
assert normalize_tool_status("error", "read_file") == "error"
def test_unexpected_status_treated_as_error_and_warns(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""An unknown present status fails closed to error and is logged."""
with caplog.at_level("WARNING", logger="deepagents_code._tool_stream"):
assert normalize_tool_status("cancelled", "execute") == "error"
assert any("Unexpected ToolMessage.status" in r.message for r in caplog.records)
def test_none_status_treated_as_error(self) -> None:
"""An explicit `None` status is unexpected, not a silent success."""
assert normalize_tool_status(None, "execute") == "error"
def test_success_default_not_warned(self, caplog: pytest.LogCaptureFixture) -> None:
"""The missing-status caller default (`"success"`) must not warn."""
with caplog.at_level("WARNING", logger="deepagents_code._tool_stream"):
normalize_tool_status("success", "read_file")
assert not any(
"Unexpected ToolMessage.status" in r.message for r in caplog.records
)
class TestToolCallBufferParseArgs:
"""Argument reassembly and completeness gating."""
def test_dict_returned_as_is(self) -> None:
buffer = ToolCallBuffer(args={"path": "foo.py"})
assert buffer.parse_args() == {"path": "foo.py"}
def test_scalar_wrapped(self) -> None:
buffer = ToolCallBuffer(args=42)
assert buffer.parse_args() == {"value": 42}
def test_fragments_joined_and_parsed(self) -> None:
buffer = ToolCallBuffer(args_parts=['{"command": "uv run', ' pytest"}'])
assert buffer.parse_args() == {"command": "uv run pytest"}
def test_incomplete_object_returns_none(self) -> None:
buffer = ToolCallBuffer(args_parts=['{"command": "uv run'])
assert buffer.parse_args() is None
def test_non_object_json_wrapped(self) -> None:
buffer = ToolCallBuffer(args_parts=["[1, 2, 3]"])
assert buffer.parse_args() == {"value": [1, 2, 3]}
def test_empty_and_whitespace_return_none(self) -> None:
assert ToolCallBuffer().parse_args() is None
assert ToolCallBuffer(args_parts=[" "]).parse_args() is None
def test_malformed_complete_json_warns_once(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""A complete-but-invalid payload logs exactly once via the latch.
The buffer is retained across chunks on failure, so `parse_args` runs
again on every later chunk; the `warned` latch keeps that from spamming
an identical warning per fragment.
"""
buffer = ToolCallBuffer(args_parts=["{bad json}"])
with caplog.at_level("WARNING", logger="deepagents_code._tool_stream"):
assert buffer.parse_args() is None
assert buffer.parse_args() is None
assert buffer.parse_args() is None
warnings = [
r
for r in caplog.records
if "look complete but failed to parse" in r.message
]
assert len(warnings) == 1
assert buffer.warned is True
def test_midstream_nested_json_does_not_warn(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""A partial nested payload that happens to end in `}` is not warned.
A chunk boundary landing right after an inner object closes leaves the
outer container open (`{"edits": [{"a": 1}`). The old "starts with {/[
and ends with }/]" heuristic mistook this for a complete-but-malformed
value and logged a WARNING on a perfectly healthy stream. The
string-aware balance check treats it as still-incomplete: no warning,
`warned` stays unset, and the next fragment can still complete it.
"""
buffer = ToolCallBuffer(args_parts=['{"edits": [{"a": 1}'])
with caplog.at_level("WARNING", logger="deepagents_code._tool_stream"):
assert buffer.parse_args() is None
assert buffer.warned is False
assert not any(
"look complete but failed to parse" in r.message for r in caplog.records
)
# The completing fragments still parse once they arrive.
buffer.ingest(name=None, tool_id=None, args=', {"b": 2}]}')
assert buffer.parse_args() == {"edits": [{"a": 1}, {"b": 2}]}
def test_trailing_brace_inside_open_string_not_warned(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""A `}` that lives inside an unterminated string is not "complete".
The payload ends in `}` (so it clears the cheap pre-check and reaches
`json.loads`), but that brace is inside an open string literal, so the
outer object is still unbalanced. The string-aware balance check must
report it incomplete — no warning — rather than treating the trailing
brace as a real close.
"""
buffer = ToolCallBuffer(args_parts=['{"content": "a } b}'])
with caplog.at_level("WARNING", logger="deepagents_code._tool_stream"):
assert buffer.parse_args() is None
assert buffer.warned is False
assert not any(
"look complete but failed to parse" in r.message for r in caplog.records
)
def test_pathologically_nested_json_is_skipped_not_raised(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""Deeply nested model output is one skipped call, not an escaped error.
`json.loads` raises `RecursionError` (not `JSONDecodeError`) on input
nested past the interpreter limit. That must be caught like any other
malformed-but-complete payload — returning `None` and logging once —
rather than escaping `parse_args` and aborting the whole turn.
"""
depth = 100_000
nested = "[" * depth + "]" * depth
buffer = ToolCallBuffer(args_parts=[nested])
with caplog.at_level("WARNING", logger="deepagents_code._tool_stream"):
assert buffer.parse_args() is None
assert buffer.warned is True
assert any(
"look complete but failed to parse" in r.message for r in caplog.records
)
def test_over_closed_json_warns_via_balance_check(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""A payload with more closers than openers is complete-but-malformed.
`{"a": 1}}` ends in `}` (so it clears the cheap pre-check and reaches
`json.loads`, which rejects the trailing brace). The string-aware balance
scan hits the `depth < 0` branch and reports the value complete — a stray
closer can never be finished by more input — so the failed parse is
warned once rather than mistaken for a still-open mid-stream fragment.
"""
buffer = ToolCallBuffer(args_parts=['{"a": 1}}'])
with caplog.at_level("WARNING", logger="deepagents_code._tool_stream"):
assert buffer.parse_args() is None
assert buffer.warned is True
assert any(
"look complete but failed to parse" in r.message for r in caplog.records
)
def test_both_args_and_args_parts_raises_at_read(self) -> None:
"""Violating the XOR invariant after construction fails loudly at read.
`__post_init__` only guards construction, and the fields are public and
mutable, so a direct assignment can reach the illegal both-populated
state. `parse_args` re-checks the invariant and raises rather than
silently reading `args` first and discarding the accumulated
`args_parts`.
"""
buffer = ToolCallBuffer(args={"a": 1})
buffer.args_parts = ['{"b":', " 2}"]
with pytest.raises(ValueError, match="cannot hold both"):
buffer.parse_args()
class TestPayloadBuilders:
"""Fixed-shape hook payloads and the output truncation invariant."""
def test_tool_use_payload_shape(self) -> None:
assert build_tool_use_payload("write_file", "toolu_1", {"path": "f"}) == {
"tool_name": "write_file",
"tool_id": "toolu_1",
"tool_args": {"path": "f"},
}
def test_tool_error_payload_shape(self) -> None:
assert build_tool_error_payload("execute") == {"tool_names": ["execute"]}
def test_tool_result_payload_shape(self) -> None:
assert build_tool_result_payload(
"write_file", "toolu_1", {"path": "f"}, "success", "ok"
) == {
"tool_name": "write_file",
"tool_id": "toolu_1",
"tool_args": {"path": "f"},
"tool_status": "success",
"tool_output": "ok",
}
def test_tool_output_truncated_to_limit(self) -> None:
"""`tool_output` is capped at `HOOK_TOOL_OUTPUT_LIMIT` and marked."""
payload = build_tool_result_payload(
"read_file", "toolu_1", {}, "success", "x" * (HOOK_TOOL_OUTPUT_LIMIT + 500)
)
# The marker is counted within the cap, so the total length never exceeds
# the limit, and its presence lets a consumer tell it was truncated.
assert len(payload["tool_output"]) == HOOK_TOOL_OUTPUT_LIMIT
assert payload["tool_output"].endswith(TOOL_OUTPUT_TRUNCATION_MARKER)
def test_tool_output_at_limit_not_marked(self) -> None:
"""Output exactly at the cap is passed through unmarked (no truncation)."""
exact = "x" * HOOK_TOOL_OUTPUT_LIMIT
payload = build_tool_result_payload(
"read_file", "toolu_1", {}, "success", exact
)
assert payload["tool_output"] == exact
assert not payload["tool_output"].endswith(TOOL_OUTPUT_TRUNCATION_MARKER)
def test_tool_output_one_over_limit_is_marked(self) -> None:
"""The first length that trips the cap (`LIMIT + 1`) is truncated + marked.
Pins the exact `>` boundary: `LIMIT` is passed through (test above) and
`LIMIT + 1` must be the first value that truncates, guarding against a
`>=`/`>` off-by-one in the builder.
"""
payload = build_tool_result_payload(
"read_file", "toolu_1", {}, "success", "x" * (HOOK_TOOL_OUTPUT_LIMIT + 1)
)
assert len(payload["tool_output"]) == HOOK_TOOL_OUTPUT_LIMIT
assert payload["tool_output"].endswith(TOOL_OUTPUT_TRUNCATION_MARKER)
def test_tool_args_not_truncated(self) -> None:
"""`tool_args` is passed through in full; only output is capped."""
big_content = "y" * (HOOK_TOOL_OUTPUT_LIMIT + 500)
payload = build_tool_result_payload(
"write_file", "toolu_1", {"content": big_content}, "success", ""
)
assert payload["tool_args"]["content"] == big_content
class TestCountUnemittedToolCalls:
"""Classification of buffered tool calls that never fired a `tool.use`."""
def test_empty_iterable_counts_zero(self) -> None:
"""No buffers means nothing unemitted."""
assert count_unemitted_tool_calls([]) == (0, 0)
def test_unnamed_buffer_ignored(self) -> None:
"""A buffer that never even got a name is not counted either way.
Without a name no tool call was ever recognizable, so it is neither an
unparsed nor an id-less unemitted call.
"""
unnamed = ToolCallBuffer()
unnamed.ingest(name=None, tool_id="t", args='{"a": ')
assert count_unemitted_tool_calls([unnamed]) == (0, 0)
def test_classifies_unparsed_and_idless_parsed(self) -> None:
"""Named buffers split into unparsed vs parsed-but-id-less; others skip.
A named buffer whose args parsed *and* carries an id would have emitted a
`tool.use`, so it is counted in neither bucket.
"""
unparsed = ToolCallBuffer()
unparsed.ingest(name="f", tool_id="t1", args='{"a": ') # never closes
parsed_with_id = ToolCallBuffer()
parsed_with_id.ingest(name="g", tool_id="t2", args='{"a": 1}')
idless_parsed = ToolCallBuffer()
idless_parsed.ingest(name="h", tool_id=None, args='{"b": 2}')
assert count_unemitted_tool_calls(
[unparsed, parsed_with_id, idless_parsed]
) == (1, 1)
def test_asymmetric_counts_pin_slot_order(self) -> None:
"""Distinct per-bucket counts catch a swapped-counter transposition.
The other cases all assert symmetric tuples (`(0, 0)`, `(1, 1)`), which a
swap of the two return values would pass unchanged. Two unparsed buffers
against one id-less-parsed buffer pins which count is which — both via the
positional tuple and the named fields.
"""
unparsed_a = ToolCallBuffer()
unparsed_a.ingest(name="a", tool_id="t1", args='{"x": ') # never closes
unparsed_b = ToolCallBuffer()
unparsed_b.ingest(name="b", tool_id="t2", args='{"y": ') # never closes
idless_parsed = ToolCallBuffer()
idless_parsed.ingest(name="c", tool_id=None, args='{"z": 3}')
result = count_unemitted_tool_calls([unparsed_a, unparsed_b, idless_parsed])
assert result == (2, 1)
assert result.unparsed == 2
assert result.idless_parsed == 1
libs/code/tests/unit_tests/test_app.py
+139
View File
@@ -10840,6 +10840,20 @@ class TestTerminalBackgroundSync:
class TestExitGracefulWorkerHandoff:
"""Verify `exit()` defers teardown for an in-flight agent worker."""
@pytest.fixture(autouse=True)
def _reset_background_hooks(self) -> Iterator[None]:
"""Isolate the module-global hook task set across tests.
These tests read the real `has_pending_hooks()`; a stray undone task
left in `_background_tasks` by an earlier test would otherwise push
`test_synchronous_when_worker_finished` onto the graceful path and flake.
"""
import deepagents_code.hooks as hooks_mod
hooks_mod._background_tasks.clear()
yield
hooks_mod._background_tasks.clear()
async def test_defers_when_agent_worker_unfinished(self) -> None:
"""A running, unfinished worker arms a deferred graceful exit."""
app = DeepAgentsApp()
@@ -10998,6 +11012,131 @@ class TestExitGracefulWorkerHandoff:
assert app._graceful_exit_task is None
worker.wait.assert_not_awaited()
async def test_drains_pending_hooks_before_exit_without_worker(self) -> None:
"""Pending fire-and-forget hooks defer teardown until they drain."""
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
with (
patch("deepagents_code.hooks.has_pending_hooks", return_value=True),
patch(
"deepagents_code.hooks.drain_pending_hooks",
new_callable=AsyncMock,
) as drain_hooks,
patch.object(App, "exit") as super_exit,
):
app.exit()
super_exit.assert_not_called()
assert app._graceful_exit_task is not None
await app._graceful_exit_task
drain_hooks.assert_awaited_once()
super_exit.assert_called_once()
async def test_drains_pending_hooks_after_waiting_for_agent(self) -> None:
"""Both the agent-worker wait and the hook drain run before teardown.
Covers the combined path: an in-flight agent worker AND pending hooks
must each be awaited inside the single graceful-exit task before the
event loop tears down.
"""
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
app._agent_running = True
worker = MagicMock()
worker.is_finished = False
worker.wait = AsyncMock()
app._agent_worker = worker
with (
patch("deepagents_code.hooks.has_pending_hooks", return_value=True),
patch(
"deepagents_code.hooks.drain_pending_hooks",
new_callable=AsyncMock,
) as drain_hooks,
patch.object(App, "exit") as super_exit,
):
app.exit()
assert app._graceful_exit_task is not None
await app._graceful_exit_task
worker.wait.assert_awaited_once()
drain_hooks.assert_awaited_once()
super_exit.assert_called_once()
async def test_drains_pending_hooks_when_worker_wait_fails(self) -> None:
"""The hook drain still runs when the agent-worker wait errors.
The drain sits in its own try/except after the worker-wait block, so a
worker that fails to persist must not skip draining pending hooks — the
final tool.result would otherwise be dropped whenever shutdown coincides
with a worker error.
"""
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
app._agent_running = True
worker = MagicMock()
worker.is_finished = False
worker.wait = AsyncMock(side_effect=RuntimeError("boom"))
app._agent_worker = worker
with (
patch("deepagents_code.hooks.has_pending_hooks", return_value=True),
patch(
"deepagents_code.hooks.drain_pending_hooks",
new_callable=AsyncMock,
) as drain_hooks,
patch.object(App, "exit") as super_exit,
):
app.exit()
assert app._graceful_exit_task is not None
await app._graceful_exit_task
worker.wait.assert_awaited_once()
drain_hooks.assert_awaited_once()
super_exit.assert_called_once()
async def test_hook_drain_timeout_is_bounded_and_logged(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""A hung hook drain is bounded by the graceful-exit budget and logged.
A slow/hanging hook subprocess must not stall an interactive quit; the
drain is wrapped in `asyncio.wait_for`, so teardown proceeds after the
timeout and the possibly-dropped final tool.result is announced rather
than the UI silently hanging.
"""
async def hang() -> None:
await asyncio.Event().wait()
app = DeepAgentsApp()
async with app.run_test() as pilot:
await pilot.pause()
with (
caplog.at_level(logging.WARNING, logger="deepagents_code.app"),
patch("deepagents_code.hooks.has_pending_hooks", return_value=True),
patch("deepagents_code.hooks.drain_pending_hooks", new=hang),
patch("deepagents_code.app._GRACEFUL_EXIT_WAIT_SECONDS", 0.01),
patch.object(App, "exit") as super_exit,
):
app.exit()
assert app._graceful_exit_task is not None
await app._graceful_exit_task
super_exit.assert_called_once()
assert any(
"Hook drain did not finish" in record.message
and record.levelno == logging.WARNING
for record in caplog.records
)
async def test_second_exit_force_quits_pending_graceful_exit(self) -> None:
"""A second exit() during a pending graceful exit force-quits.
libs/code/tests/unit_tests/test_hooks.py
+253 -10
View File
@@ -2,10 +2,12 @@
from __future__ import annotations
import asyncio
import json
import logging
import subprocess
from typing import TYPE_CHECKING, Any
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import pytest
@@ -190,32 +192,46 @@ class TestDispatchHook:
mock_run.assert_called_once()
async def test_hook_without_command_skipped(self):
"""Hook entry missing 'command' is silently skipped."""
async def test_hook_without_command_skipped(self, caplog):
"""Hook entry missing 'command' is skipped and the misconfig is warned."""
hooks_mod._hooks_config = [{"events": ["session.start"]}]
with patch("deepagents_code.hooks.subprocess.run") as mock_run:
with (
patch("deepagents_code.hooks.subprocess.run") as mock_run,
caplog.at_level(logging.WARNING, logger="deepagents_code.hooks"),
):
await hooks_mod.dispatch_hook("session.start", {})
mock_run.assert_not_called()
# The config mistake must be greppable, not look like a hook that simply
# never matched.
assert "invalid `command`" in caplog.text
async def test_hook_with_string_command_skipped(self):
"""Hook with string command (not list) is skipped."""
async def test_hook_with_string_command_skipped(self, caplog):
"""Hook with string command (not list) is skipped and warned."""
hooks_mod._hooks_config = [{"command": "echo hello"}]
with patch("deepagents_code.hooks.subprocess.run") as mock_run:
with (
patch("deepagents_code.hooks.subprocess.run") as mock_run,
caplog.at_level(logging.WARNING, logger="deepagents_code.hooks"),
):
await hooks_mod.dispatch_hook("session.start", {})
mock_run.assert_not_called()
assert "invalid `command`" in caplog.text
async def test_hook_with_empty_command_list_skipped(self):
"""Hook with empty command list is skipped."""
async def test_hook_with_empty_command_list_skipped(self, caplog):
"""Hook with empty command list is skipped and warned."""
hooks_mod._hooks_config = [{"command": []}]
with patch("deepagents_code.hooks.subprocess.run") as mock_run:
with (
patch("deepagents_code.hooks.subprocess.run") as mock_run,
caplog.at_level(logging.WARNING, logger="deepagents_code.hooks"),
):
await hooks_mod.dispatch_hook("session.start", {})
mock_run.assert_not_called()
assert "invalid `command`" in caplog.text
async def test_timeout_does_not_propagate(self):
"""TimeoutExpired is caught and logged, not raised."""
@@ -261,6 +277,31 @@ class TestDispatchHook:
# Should not raise.
await hooks_mod.dispatch_hook("session.start", {})
async def test_generic_error_logged_at_warning(
self, caplog: pytest.LogCaptureFixture
) -> None:
"""Unexpected subprocess failures surface at WARNING, not hidden at DEBUG.
The catch-all handles the least-understood failures (e.g. ENOEXEC for a
non-executable hook file), so a silent debug log would hide a hook that
never fires.
"""
hooks_mod._hooks_config = [{"command": ["bad"]}]
with (
patch(
"deepagents_code.hooks.subprocess.run",
side_effect=RuntimeError("unexpected"),
),
caplog.at_level("WARNING", logger="deepagents_code.hooks"),
):
await hooks_mod.dispatch_hook("session.start", {})
assert any(
"failed unexpectedly" in r.getMessage() and r.levelname == "WARNING"
for r in caplog.records
)
async def test_multiple_hooks_dispatched(self):
"""All matching hooks fire, not just the first."""
hooks_mod._hooks_config = [
@@ -315,6 +356,57 @@ class TestDispatchHook:
# Should not raise despite non-serializable payload.
await hooks_mod.dispatch_hook("session.start", {"bad": object()})
async def test_dispatch_hook_stringifies_non_serializable_value(self):
"""A non-JSON-serializable value is stringified and still delivered.
Locks in the `default=str` behavior: the subprocess must still run with
the value coerced to its string form, rather than the whole event being
dropped (which is what happens if `default=str` is removed and the
serialization error is swallowed by the outer guard).
"""
hooks_mod._hooks_config = [{"command": ["cat"]}]
class _Widget:
def __str__(self) -> str:
return "STRINGIFIED_WIDGET"
with patch("deepagents_code.hooks.subprocess.run") as mock_run:
await hooks_mod.dispatch_hook("tool.use", {"tool_args": _Widget()})
mock_run.assert_called_once()
stdin_bytes = mock_run.call_args[1]["input"]
assert b"STRINGIFIED_WIDGET" in stdin_bytes
async def test_dispatch_hook_drops_event_when_default_str_raises(self, caplog):
"""If `default=str` itself raises, the event is dropped (documented gap).
`json.dumps(default=str)` degrades a non-serializable value to its string
form, but if that value's own `__str__` raises, serialization fails and
the outer guard drops the whole event rather than delivering a partial
payload. This is a known non-guarantee — only `tool_output` is
sentinel-protected upstream, not `tool_args` — so pin it here so any
future change to the guard is a conscious one.
"""
hooks_mod._hooks_config = [{"command": ["cat"]}]
class _Exploding:
def __str__(self) -> str:
msg = "cannot stringify"
raise RuntimeError(msg)
__repr__ = __str__
with (
patch("deepagents_code.hooks.subprocess.run") as mock_run,
caplog.at_level(logging.WARNING, logger="deepagents_code.hooks"),
):
await hooks_mod.dispatch_hook("tool.use", {"tool_args": _Exploding()})
# Serialization failed, so the subprocess never ran — the whole event is
# dropped rather than partially delivered.
mock_run.assert_not_called()
assert "Unexpected error in dispatch_hook" in caplog.text
# ---------------------------------------------------------------------------
# dispatch_hook_fire_and_forget
@@ -355,3 +447,154 @@ class TestDispatchHookFireAndForget:
# Call from sync context with no running loop — should not raise
hooks_mod.dispatch_hook_fire_and_forget("session.start", {})
assert len(hooks_mod._background_tasks) == 0
# ---------------------------------------------------------------------------
# drain_pending_hooks
# ---------------------------------------------------------------------------
class TestDrainPendingHooks:
"""Test draining of in-flight fire-and-forget hook tasks."""
async def test_drains_pending_task_before_returning(self):
"""A scheduled hook runs to completion before drain returns."""
hooks_mod._hooks_config = [{"command": ["echo"]}]
with patch("deepagents_code.hooks.subprocess.run") as mock_run:
hooks_mod.dispatch_hook_fire_and_forget("tool.result", {"tool_name": "x"})
assert len(hooks_mod._background_tasks) == 1
await hooks_mod.drain_pending_hooks()
assert hooks_mod._background_tasks == set()
mock_run.assert_called_once()
async def test_no_pending_tasks_is_noop(self):
"""Draining with nothing pending returns immediately."""
await hooks_mod.drain_pending_hooks()
assert hooks_mod._background_tasks == set()
async def test_real_dispatch_then_real_drain_completes_inflight_task(self):
"""End-to-end: the real drain runs a real, still-in-flight dispatch to done.
Composes the real `dispatch_hook_fire_and_forget` with the real
`drain_pending_hooks` (only `subprocess.run` is patched, to record). A
freshly scheduled task has not run yet — no `await` has ceded control — so
the hook subprocess has not been invoked at dispatch time; the drain is
what carries it to completion. This is the composed seam that the
per-surface tests (which patch the dispatch capture point) and the
instant-return drain test do not exercise together: it pins that the drain
is load-bearing for an in-flight hook rather than one already finished.
"""
hooks_mod._hooks_config = [{"command": ["echo"]}]
ran = {"count": 0}
def record_run(*_args: object, **_kwargs: object) -> MagicMock:
ran["count"] += 1
return MagicMock(returncode=0, stdout="", stderr="")
with patch("deepagents_code.hooks.subprocess.run", side_effect=record_run):
hooks_mod.dispatch_hook_fire_and_forget("tool.result", {"tool_name": "x"})
# Real task scheduled but not yet executed (no await has occurred), so
# the hook subprocess has not been called.
assert ran["count"] == 0
assert len(hooks_mod._background_tasks) == 1
await hooks_mod.drain_pending_hooks()
assert ran["count"] == 1
assert hooks_mod._background_tasks == set()
async def test_drain_swallows_task_exceptions(self):
"""A task that raises does not propagate out of drain."""
async def _boom() -> None:
await asyncio.sleep(0)
msg = "boom"
raise RuntimeError(msg)
loop = asyncio.get_running_loop()
task = loop.create_task(_boom())
hooks_mod._background_tasks.add(task)
task.add_done_callback(hooks_mod._background_tasks.discard)
# Must not raise despite the task erroring.
await hooks_mod.drain_pending_hooks()
assert hooks_mod._background_tasks == set()
async def test_drain_snapshots_once_and_ignores_later_scheduled_hooks(self):
"""A hook scheduled *during* the drain await is not awaited by that drain.
`drain_pending_hooks` snapshots the in-flight set once; its documented
precondition is that no further dispatches happen during the await. Pin
that snapshot-once behavior: a task that schedules another task while the
drain is in flight leaves the second one un-awaited by the same drain
call, so a change to loop-until-empty semantics fails here.
"""
loop = asyncio.get_running_loop()
second_done = False
async def _second() -> None:
nonlocal second_done
await asyncio.sleep(0.05)
second_done = True
async def _first() -> None:
# Yield first so this runs inside the drain's gather, then schedule a
# new hook task *after* the drain has already snapshotted the set.
await asyncio.sleep(0)
second = loop.create_task(_second())
hooks_mod._background_tasks.add(second)
second.add_done_callback(hooks_mod._background_tasks.discard)
first = loop.create_task(_first())
hooks_mod._background_tasks.add(first)
first.add_done_callback(hooks_mod._background_tasks.discard)
await hooks_mod.drain_pending_hooks()
# The drain awaited `first` (now done) but not the task it spawned.
assert first.done()
assert not second_done
assert hooks_mod._background_tasks # the second task is still tracked
# Clean up the straggler so it does not leak into other tests.
await asyncio.gather(*hooks_mod._background_tasks, return_exceptions=True)
assert second_done
# ---------------------------------------------------------------------------
# has_pending_hooks
# ---------------------------------------------------------------------------
class TestHasPendingHooks:
"""`has_pending_hooks` gates the TUI's drain-on-exit, so verify it directly.
A wrong predicate here would silently skip the graceful-exit drain and drop
the final `tool.result`, which a mock-only test could not catch.
"""
async def test_false_when_no_tasks(self):
"""No scheduled hooks means nothing to wait for."""
assert hooks_mod.has_pending_hooks() is False
async def test_true_while_task_in_flight_then_false_after_drain(self):
"""Reports True with a live task pending and False once it drains."""
async def _slow() -> None:
await asyncio.sleep(0.05)
loop = asyncio.get_running_loop()
task = loop.create_task(_slow())
hooks_mod._background_tasks.add(task)
task.add_done_callback(hooks_mod._background_tasks.discard)
assert hooks_mod.has_pending_hooks() is True
await hooks_mod.drain_pending_hooks()
assert hooks_mod.has_pending_hooks() is False
assert hooks_mod._background_tasks == set()
libs/code/tests/unit_tests/test_messages.py
+90
View File
@@ -753,6 +753,96 @@ class TestToolCallMessageDuration:
assert "Took" not in getattr(content, "plain", str(content))
class TestToolCallMessageTerminalStateGuards:
"""A rejected/skipped row must not flip to success/error on a resumed turn."""
async def test_set_success_noop_on_rejected_row(self) -> None:
"""A resumed synthetic success ToolMessage keeps a rejected row rejected.
After a reasoned reject the turn can resume and stream a synthetic
ToolMessage for the rejected tool; `set_success` must be ignored so the
row keeps its terminal rejected state instead of flipping.
"""
app = _tool_msg_app("execute", {"command": "rm -rf /"})
async with app.run_test() as pilot:
await pilot.pause()
app.msg.set_rejected()
assert app.msg._status == "rejected"
app.msg.set_success("done")
await pilot.pause()
assert app.msg._status == "rejected"
assert app.msg.is_success is False
async def test_set_error_noop_on_rejected_row(self) -> None:
"""A resumed synthetic error ToolMessage keeps a rejected row rejected."""
app = _tool_msg_app("execute", {"command": "rm -rf /"})
async with app.run_test() as pilot:
await pilot.pause()
app.msg.set_rejected()
app.msg.set_error("boom")
await pilot.pause()
assert app.msg._status == "rejected"
async def test_set_success_noop_on_skipped_row(self) -> None:
"""A skipped row (sibling rejection) stays skipped, not flipped to success.
The guard names both `rejected` and `skipped`; a tool skipped because a
sibling was rejected can still receive a synthetic success ToolMessage on
the resumed turn, which must be ignored.
"""
app = _tool_msg_app("execute", {"command": "ls"})
async with app.run_test() as pilot:
await pilot.pause()
app.msg.set_skipped()
assert app.msg._status == "skipped"
app.msg.set_success("done")
await pilot.pause()
assert app.msg._status == "skipped"
assert app.msg.is_success is False
async def test_set_error_noop_on_skipped_row(self) -> None:
"""A skipped row keeps its terminal state instead of flipping to error."""
app = _tool_msg_app("execute", {"command": "ls"})
async with app.run_test() as pilot:
await pilot.pause()
app.msg.set_skipped()
app.msg.set_error("boom")
await pilot.pause()
assert app.msg._status == "skipped"
class TestToolCallMessageArgs:
"""The public `args` accessor must not expose internal widget state."""
def test_args_returns_shallow_copy(self) -> None:
"""Rebinding top-level keys of the returned dict must not affect `_args`.
Hook payloads are built directly from `tool_msg.args`, so the copy is a
load-bearing safety contract: a consumer that reassigns its payload's
top-level keys must not corrupt the widget's stored arguments by
reference. The copy is shallow — nested mutable values are shared (see
`test_args_nested_values_are_shared`) — which is sufficient because the
only consumer serializes the payload rather than deep-mutating it.
"""
msg = ToolCallMessage("write_file", {"file_path": "a.py", "content": "x"})
returned = msg.args
returned["file_path"] = "hacked.py"
returned["injected"] = True
assert msg.args == {"file_path": "a.py", "content": "x"}
assert msg._args == {"file_path": "a.py", "content": "x"}
def test_args_nested_values_are_shared(self) -> None:
"""The copy is shallow: nested mutables are shared, not deep-copied.
Pins the documented boundary of the `args` accessor so a future reader
does not mistake the shallow copy for a deep one.
"""
msg = ToolCallMessage("edit_file", {"edits": [{"old": "a"}]})
returned = msg.args
returned["edits"][0]["old"] = "mutated"
assert msg._args["edits"][0]["old"] == "mutated"
class TestToolCallMessageTodos:
"""Tests for `write_todos` output formatting."""
libs/code/tests/unit_tests/test_non_interactive.py
+1335
View File
File diff suppressed because it is too large Load Diff
libs/code/tests/unit_tests/test_textual_adapter.py
+2753 -6
View File
File diff suppressed because it is too large Load Diff