diff --git a/libs/code/deepagents_code/agent.py b/libs/code/deepagents_code/agent.py index 91d0d2c2151..041f6d846a1 100644 --- a/libs/code/deepagents_code/agent.py +++ b/libs/code/deepagents_code/agent.py @@ -1929,12 +1929,11 @@ def _should_interrupt_tool_call( Returns: `True` to interrupt, or `False` for Auto/YOLO bypass. """ - from deepagents_code.hooks.server_middleware import pre_tool_behavior + from deepagents_code.hooks.server_middleware import hook_decided_permission tool_call = getattr(request, "tool_call", None) tool_call_id = str(tool_call.get("id") or "") if isinstance(tool_call, dict) else "" - hook_behavior = pre_tool_behavior(getattr(request, "state", None), tool_call_id) - if hook_behavior in {"allow", "deny"}: + if hook_decided_permission(getattr(request, "state", None), tool_call_id): return False runtime = getattr(request, "runtime", None) diff --git a/libs/code/deepagents_code/auto_mode.py b/libs/code/deepagents_code/auto_mode.py index c97e924a1b4..5751e796b14 100644 --- a/libs/code/deepagents_code/auto_mode.py +++ b/libs/code/deepagents_code/auto_mode.py @@ -2733,12 +2733,12 @@ async def aafter_model( ) if ai_message is None or not ai_message.tool_calls: return {"_auto_decision_plan": None} - from deepagents_code.hooks.server_middleware import pre_tool_behavior + from deepagents_code.hooks.server_middleware import hook_decided_permission hook_bypass_ids = { _tool_call_id(call) for call in ai_message.tool_calls - if pre_tool_behavior(state, _tool_call_id(call)) in {"allow", "deny"} + if hook_decided_permission(state, _tool_call_id(call)) } thread_key = _thread_key(runtime) # Derive the emission scope once for the whole node run. Deriving it diff --git a/libs/code/deepagents_code/hooks/server_middleware.py b/libs/code/deepagents_code/hooks/server_middleware.py index fcf1152897a..fcfe8e6b0bf 100644 --- a/libs/code/deepagents_code/hooks/server_middleware.py +++ b/libs/code/deepagents_code/hooks/server_middleware.py @@ -16,6 +16,7 @@ from datetime import UTC, datetime, timedelta from typing import ( TYPE_CHECKING, + Annotated, Any, Literal, NotRequired, @@ -35,6 +36,7 @@ AgentMiddleware, AgentState, ContextT, + PrivateStateAttr, ResponseT, hook_config, ) @@ -93,19 +95,60 @@ _INVOCATION_NAMESPACE = UUID("f2896d18-cf2a-4e7d-b11a-d5b10fc0e335") PreToolBehavior: TypeAlias = Literal["allow", "deny", "none"] +_DEFAULT_DENY_REASON = "Blocked by PreToolUse hook" -class _PreToolState(TypedDict): - behavior: PreToolBehavior - reason: str | None +class _PreToolDenied(TypedDict): + """Outcome for a call a hook refused. A denial always carries a reason.""" + + behavior: Literal["deny"] + reason: str + context: list[str] + + +class _PreToolPassed(TypedDict): + """Outcome for a call a hook allowed or had no opinion on.""" + + behavior: Literal["allow", "none"] context: list[str] +_PreToolState: TypeAlias = _PreToolDenied | _PreToolPassed + + class ServerHooksState(AgentState[Any]): - """Agent state extensions for server-owned hook middleware.""" + """Agent state extensions for server-owned hook middleware. + + Both fields are per-turn bookkeeping owned by `ServerHooksMiddleware` and + marked `PrivateStateAttr`: they are omitted from the public graph I/O schema, + and `SubAgentMiddleware` strips them from subagent result merges so parallel + `task` calls cannot produce two concurrent writes to these `LastValue` + channels. + + `PrivateStateAttr` only omits the fields from the input and output schemas; + the channels stay ordinary checkpointed `LastValue` channels visible to every + node, so values still flow from `after_model` to `wrap_tool_call` and survive + interrupt/resume. + + Note: + If either field ever needs a reducer, the reducer must be placed *after* + `PrivateStateAttr` in the `Annotated` metadata. LangGraph only inspects + the last metadata entry when detecting reducers, so a reducer added + before the marker is silently ignored. + """ - _hooks_stop_continuation_count: NotRequired[int] - _hooks_pre_tool_outcomes: NotRequired[dict[str, _PreToolState]] + _hooks_stop_continuation_count: NotRequired[Annotated[int, PrivateStateAttr]] + """Stop-hook continuations in the current turn; reset to 0 when the loop ends.""" + + _hooks_pre_tool_outcomes: NotRequired[ + Annotated[dict[str, _PreToolState], PrivateStateAttr] + ] + """Pre-execution hook verdicts keyed by tool-call id. + + A full snapshot of the *current* turn's calls, not an accumulator: every + `_after_model` replaces the whole dict (including with `{}`) so stale ids + cannot survive into a later turn. + """ class _SessionHookGate(TypedDict): @@ -392,7 +435,7 @@ def _after_model( reason = ( permission.reason or decision.stop_reason - or "Blocked by PreToolUse hook" + or _DEFAULT_DENY_REASON ) elif permission.behavior == "ask": blocked = _ask_permission_via_hitl(call, permission) @@ -408,11 +451,19 @@ def _after_model( ) elif permission.behavior == "allow": behavior = "allow" - outcomes[call.id] = { - "behavior": behavior, - "reason": reason, - "context": hook_context, - } + if behavior == "deny": + outcomes[call.id] = { + "behavior": "deny", + # Every deny path above resolves a reason; the guard keeps the + # "a denial always explains itself" invariant checkable here. + "reason": reason if reason is not None else _DEFAULT_DENY_REASON, + "context": hook_context, + } + else: + outcomes[call.id] = { + "behavior": behavior, + "context": hook_context, + } return {_PRE_TOOL_STATE_KEY: outcomes} def _maybe_post_tool_use( @@ -541,19 +592,23 @@ def _event_enabled(gate: _SessionHookGate | None, event: HookEvent) -> bool: return gate is not None and event.value in gate["events"] -def pre_tool_behavior(state: object, tool_call_id: str) -> PreToolBehavior | None: - """Return the replayed pre-execution hook behavior for one call.""" +def hook_decided_permission(state: object, tool_call_id: str) -> bool: + """Report whether a pre-execution hook already settled permission for a call. + + Args: + state: Agent state carrying the current turn's hook outcomes. + tool_call_id: Tool call to look up. + + Returns: + `True` when a hook explicitly allowed or denied the call, so stock + approval flows must not prompt again. `False` when no hook ran, the hook + expressed no opinion, or no outcome was recorded -- in every one of those + cases normal approval still applies. + """ outcome = _pre_tool_state(state, tool_call_id) if outcome is None: - return None - behavior = outcome.get("behavior") - if behavior == "allow": - return "allow" - if behavior == "deny": - return "deny" - if behavior == "none": - return "none" - return None + return False + return outcome.get("behavior") in {"allow", "deny"} def _pre_tool_state(state: object, tool_call_id: str) -> Mapping[str, object] | None: diff --git a/libs/code/tests/unit_tests/hooks/test_server_lifecycle.py b/libs/code/tests/unit_tests/hooks/test_server_lifecycle.py index 0a929ea6389..8d060fd4bf9 100644 --- a/libs/code/tests/unit_tests/hooks/test_server_lifecycle.py +++ b/libs/code/tests/unit_tests/hooks/test_server_lifecycle.py @@ -7,18 +7,26 @@ import sys from datetime import UTC, datetime, timedelta from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NotRequired from unittest.mock import MagicMock from uuid import uuid4 import pytest +from deepagents import create_deep_agent +from deepagents.middleware import CompiledSubAgent, SubAgent +from deepagents.middleware._state import private_state_field_names +from langchain.agents import create_agent +from langchain.agents.middleware.types import AgentMiddleware, AgentState from langchain_core.language_models.fake_chat_models import GenericFakeChatModel -from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langchain_core.runnables import RunnableLambda +from langchain_core.tools import tool from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import START, StateGraph from langgraph.types import Command from pydantic import BaseModel +from deepagents_code._cli_context import CLIContextSchema from deepagents_code.agent import _should_interrupt_tool_call, create_cli_agent from deepagents_code.approval_mode import ApprovalMode from deepagents_code.hooks.client import fulfill_hook_invocation @@ -73,7 +81,11 @@ from deepagents_code.hooks.transcript import SUBAGENT_TRANSCRIPT_ID_METADATA_KEY if TYPE_CHECKING: - from langchain_core.runnables import RunnableConfig + from collections.abc import Callable, Sequence + + from langchain_core.language_models import LanguageModelInput + from langchain_core.runnables import Runnable, RunnableConfig + from langchain_core.tools import BaseTool from deepagents_code._cli_context import CLIContext @@ -82,6 +94,317 @@ class _ReplayState(BaseModel): completed: bool +class _ToolCallingFakeChatModel(GenericFakeChatModel): + def bind_tools( + self, + tools: Sequence[dict[str, Any] | type | Callable[..., Any] | BaseTool], + *, + tool_choice: str | None = None, + **kwargs: Any, + ) -> Runnable[LanguageModelInput, AIMessage]: + _ = tools, tool_choice, kwargs + return self + + +class _PublicHookState(AgentState[Any]): + """Mirror of `ServerHooksState`'s keys *without* the privacy marker. + + Used only by the `CompiledSubAgent` fixtures below, which need to write the + hook keys from outside the middleware. `_hook_state_keys` checks this mirror + against `ServerHooksState` so it cannot silently drift. + """ + + _hooks_stop_continuation_count: NotRequired[int] + _hooks_pre_tool_outcomes: NotRequired[dict[str, Any]] + + +def _hook_state_keys() -> frozenset[str]: + """Return the private hook keys, asserting the local mirror matches.""" + private_fields = private_state_field_names(ServerHooksState) + hook_keys = frozenset(name for name in private_fields if name.startswith("_hooks_")) + mirrored = frozenset(_PublicHookState.__annotations__) & hook_keys + assert mirrored == hook_keys, ( + f"_PublicHookState is missing hook keys {sorted(hook_keys - mirrored)}; " + "update the fixture when ServerHooksState gains a private field." + ) + return hook_keys + + +def _hook_state_subagent(*, name: str, content: str) -> CompiledSubAgent: + """Subagent that writes both hook keys directly, as a bare compiled runnable. + + `CompiledSubAgent` is the weaker of the two outbound layers: a raw `SubAgent` + also gets filtered by its own graph's output schema, so only this shape + exercises `SubAgentMiddleware`'s explicit `private_state_keys` strip. + """ + + def finish(_state: _PublicHookState) -> dict[str, Any]: + return { + "_hooks_stop_continuation_count": 1, + "_hooks_pre_tool_outcomes": {name: {"behavior": "none", "context": []}}, + "messages": [AIMessage(content=content)], + } + + return CompiledSubAgent( + name=name, + description=f"Return {content}.", + runnable=RunnableLambda(finish), + ) + + +def _real_hook_subagent(*, name: str, content: str, cwd: Path) -> SubAgent: + """Subagent built the way production builds them, with its own hook middleware. + + This is the shape that actually triggered the reported crash: every real + subagent carries `ServerHooksMiddleware`, whose `_after_model` writes + `_hooks_pre_tool_outcomes` unconditionally -- even with no hooks configured -- + so two parallel `task` calls both write that channel in one step. + """ + middleware: list[AgentMiddleware[Any, Any]] = [ + ServerHooksMiddleware(cwd=cwd, emit_stop=False) + ] + return SubAgent( + name=name, + description=f"Return {content}.", + system_prompt=f"Say {content}.", + model=_ToolCallingFakeChatModel( + messages=iter([AIMessage(content=content)]), + ), + middleware=middleware, + ) + + +def test_server_hook_state_fields_are_private() -> None: + private_fields = private_state_field_names(ServerHooksState) + + assert "_hooks_pre_tool_outcomes" in private_fields + assert "_hooks_stop_continuation_count" in private_fields + # `private_state_field_names` skips schemas whose annotations cannot be + # resolved, which would silently return an empty set and revert the fix. + assert _hook_state_keys() == { + "_hooks_pre_tool_outcomes", + "_hooks_stop_continuation_count", + } + + +def test_task_omits_private_server_hook_state_from_subagent_update( + tmp_path: Path, +) -> None: + """A single `task` must not clobber the parent's hook state. + + One subagent cannot trip `InvalidUpdateError`, so this covers the silent half + of the bug. Built through `create_deep_agent` so the private-key derivation in + `deepagents.graph` is exercised rather than reimplemented. + """ + model = _ToolCallingFakeChatModel( + messages=iter( + [ + AIMessage( + content="", + tool_calls=[ + { + "name": "task", + "args": { + "description": "Run the child", + "subagent_type": "child", + }, + "id": "call-child", + "type": "tool_call", + } + ], + ), + AIMessage(content="parent complete"), + ] + ) + ) + agent = create_deep_agent( + model=model, + middleware=[ServerHooksMiddleware(cwd=tmp_path)], + subagents=[_hook_state_subagent(name="child", content="child complete")], + ) + + result = agent.invoke({"messages": [HumanMessage(content="delegate")]}) + + assert "_hooks_pre_tool_outcomes" not in result + assert "_hooks_stop_continuation_count" not in result + tool_messages = [ + message for message in result["messages"] if isinstance(message, ToolMessage) + ] + assert len(tool_messages) == 1 + assert tool_messages[0].content == "child complete" + + +@pytest.mark.parametrize("subagent_kind", ["compiled", "real"]) +def test_parallel_tasks_do_not_merge_subagent_server_hook_state( + tmp_path: Path, + subagent_kind: str, +) -> None: + """Two `task` calls completing in one step must not both write hook channels. + + Covers both subagent shapes: `compiled` writes the keys by hand and exercises + `SubAgentMiddleware`'s strip, while `real` carries its own + `ServerHooksMiddleware` and reproduces the production trigger + (`_hooks_pre_tool_outcomes`, written unconditionally by `_after_model`). + """ + model = _ToolCallingFakeChatModel( + messages=iter( + [ + AIMessage( + content="", + tool_calls=[ + { + "name": "task", + "args": { + "description": "Run the first child", + "subagent_type": "first", + }, + "id": "call-first", + "type": "tool_call", + }, + { + "name": "task", + "args": { + "description": "Run the second child", + "subagent_type": "second", + }, + "id": "call-second", + "type": "tool_call", + }, + ], + ), + AIMessage(content="parent complete"), + ] + ) + ) + subagents: list[Any] = ( + [ + _hook_state_subagent(name="first", content="first complete"), + _hook_state_subagent(name="second", content="second complete"), + ] + if subagent_kind == "compiled" + else [ + _real_hook_subagent(name="first", content="first complete", cwd=tmp_path), + _real_hook_subagent(name="second", content="second complete", cwd=tmp_path), + ] + ) + checkpointer = InMemorySaver() + agent = create_deep_agent( + model=model, + middleware=[ServerHooksMiddleware(cwd=tmp_path)], + subagents=subagents, + checkpointer=checkpointer, + ) + config: RunnableConfig = {"configurable": {"thread_id": str(uuid4())}} + + result = agent.invoke( + {"messages": [HumanMessage(content="run both children")]}, + config=config, + ) + + tool_messages = { + message.tool_call_id: message.content + for message in result["messages"] + if isinstance(message, ToolMessage) + } + assert tool_messages == { + "call-first": "first complete", + "call-second": "second complete", + } + state = agent.get_state(config).values + assert "first" not in state.get("_hooks_pre_tool_outcomes", {}) + assert "second" not in state.get("_hooks_pre_tool_outcomes", {}) + assert "_hooks_stop_continuation_count" not in state + + +@pytest.mark.parametrize("resume_round_trip", [False, True]) +def test_pretool_deny_blocks_tool_through_real_graph( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + resume_round_trip: bool, +) -> None: + """A `deny` must survive the real node-to-node channel and block the tool. + + The other deny tests call `_after_model`/`wrap_tool_call` directly and copy the + update between them by hand, so none of them would notice if the outcome stopped + reaching the tools node. This drives a compiled graph instead, which is what + marking the state private could plausibly have broken. + """ + executed: list[str] = [] + + @tool + def danger(target: str) -> str: + """Do something that hooks should be able to block.""" + executed.append(target) + return f"ran on {target}" + + model = _ToolCallingFakeChatModel( + messages=iter( + [ + AIMessage( + content="", + tool_calls=[ + { + "name": "danger", + "args": {"target": "prod"}, + "id": "call-danger", + "type": "tool_call", + } + ], + ), + AIMessage(content="stopped"), + ] + ) + ) + + def _deny(*_args: object, **_kwargs: object) -> PreToolUseDecision: + return PreToolUseDecision( + event=HookEvent.PRE_TOOL_USE, + permission=PermissionEffect(behavior="deny", reason="blocked by policy"), + ) + + monkeypatch.setattr( + "deepagents_code.hooks.server_middleware._invoke_hook", + _deny, + ) + agent = create_agent( + model=model, + tools=[danger], + middleware=[ServerHooksMiddleware(cwd=tmp_path)], + context_schema=CLIContextSchema, + checkpointer=InMemorySaver(), + ) + context = CLIContextSchema( + hooks_snapshot_id="snap", + hooks_server_events=[HookEvent.PRE_TOOL_USE.value], + thread_id="t1", + approval_mode=ApprovalMode.MANUAL.value, + ) + config: RunnableConfig = {"configurable": {"thread_id": str(uuid4())}} + + if resume_round_trip: + # Prove the outcome survives a checkpoint round trip, not just one step. + agent.invoke( + {"messages": [HumanMessage(content="go")]}, + config=config, + context=context, + ) + result = agent.invoke(None, config=config, context=context) + else: + result = agent.invoke( + {"messages": [HumanMessage(content="go")]}, + config=config, + context=context, + ) + + assert executed == [] + denied = [ + message for message in result["messages"] if isinstance(message, ToolMessage) + ] + assert len(denied) == 1 + assert denied[0].status == "error" + assert "blocked by policy" in str(denied[0].content) + + def _request(event: PreToolUseEvent | None = None) -> HookInvocationRequest: invocation = HookInvocation( context=HookContext(