diff --git a/libs/code/THREAT_MODEL.md b/libs/code/THREAT_MODEL.md index 509dc45a10..a00b933902 100644 --- a/libs/code/THREAT_MODEL.md +++ b/libs/code/THREAT_MODEL.md @@ -183,6 +183,7 @@ #### DC5: Offloaded Conversation History - **Fields**: Timestamped, formatted conversation messages written by `offload.offload_messages_to_backend`. +- **Producers**: Three paths write this data — automatic trigger-based compaction, the model-initiated `compact_conversation` tool (HITL-gated, see TB2), and the explicit `/offload` command. On a server-backed agent the last runs the `offload` operation graph, which writes outside the tool-approval path. Only the *forced* paths (`/offload`, whether via the operation graph or the seeded tool call) wrap the backend in `offload_middleware._ArchiveReadGuard`, which fails closed rather than truncating existing history when its prerequisite read fails; the automatic and model-initiated paths write through the raw backend on the SDK's own code path. The guard is applied per write site rather than by the backend's type, so a new write site does not inherit it — see the `_guarded_backend()` call sites. - **Storage**: Sandbox backend filesystem at path `/conversation_history/{thread_id}.md` where `thread_id` is a UUID7 (via `sessions.generate_thread_id`). - **Access**: Accessible within the sandbox session; depends on provider access controls. - **Encryption**: Depends on sandbox provider storage backend. @@ -222,6 +223,10 @@ - **Outside**: Once the user clicks "approve" (interactive) or a command passes the allow-list check (non-interactive), the tool executes with no further framework-level gating. - **Crossing mechanism**: LangGraph HITL interrupt routed through `RemoteAgent` SSE stream. - **Key note**: `auto_approve` mode bypasses all HITL approval prompts while still displaying Unicode/URL warnings. +- **Key note**: This boundary gates the *model-initiated* `compact_conversation` tool. Whether the explicit `/offload` command crosses it depends on the agent: + - **Server-backed agent**: `/offload` does *not* cross this boundary. It runs the separate `offload` operation graph (`offload_middleware.create_forced_compaction_graph`), which has no tool node and no HITL middleware, and therefore no tool-approval interrupt it can route — the slash command is the authorization. That graph still dispatches `PreCompact` — and `PreToolUse`, which the same `ServerHooksMiddleware` boundary raises for the forced call — against an in-memory forced tool call, so a hook can veto it, but the archive write reaches `backend.awrite()` without traversing the tool-approval path. Note that the hook boundary itself can still raise interrupts on this path: hook invocations pause the graph and the client fulfills them, and a `PreToolUse` hook returning an `ask` permission makes `ServerHooksMiddleware` raise a `HITLRequest` interrupt that nothing here can answer — the client reports the failure and leaves the run paused rather than proceeding. See DF25. + - **Local in-process `Pregel` agent (including ACP mode)**: `/offload` *does* cross this boundary. `app._drive_local_seeded_compaction` seeds a `compact_conversation` tool call through the agent's own HITL-gated `ToolNode`, and the **client approves the resulting interrupt itself** — that self-approval is the trust decision on this path, scoped to exactly one forced `compact_conversation` call keyed on the seeded call id. Every other tool requested during that run is rejected by the compaction middleware independently of HITL configuration. +- **Key note**: Only the *pre* hook events fire for `/offload` on the operation-graph path. `PostToolUse`/`PostToolUseFailure`, which `ServerHooksMiddleware.awrap_tool_call` raises for a model-initiated `compact_conversation`, do not fire — the graph runs no `ToolNode`. A matcher on `compact_conversation` therefore observes automatic compaction but not the explicit command. Likewise an allowing `PreToolUse` hook's `additionalContext` is discarded (logged, not injected): there is no tool result to carry it. #### TB3: Tool Result → LLM Context @@ -252,6 +257,8 @@ - **Inside**: Server bound to `127.0.0.1` by default; `client/launch/server.py:_DEFAULT_HOST = "127.0.0.1"`. `RemoteAgent` only connects to the URL returned by `ServerProcess.url`. Server is ephemeral — started at session start, stopped at session end. Binds a free ephemeral port by default (`client/launch/server.py:_EPHEMERAL_PORT`); an explicit port is honored but still falls back to a free port if occupied. - **Outside**: `LANGGRAPH_AUTH_TYPE=noop` disables all LangGraph server authentication. Any process on localhost that discovers the port can submit requests, read thread state, or inject messages. - **Crossing mechanism**: HTTP POST/GET to `http://127.0.0.1:{port}` using `langgraph.pregel.remote.RemoteGraph`. +- **Key note**: For the default built-in `graph_ref`, `generate_langgraph_json` registers two graphs, `agent` and `offload`, so the surface is both. (A custom `graph_ref` registers `agent` alone.) The `offload` graph accepts exactly one writable **run-input** channel, `messages`, enforced by an explicit `StateGraph(input_schema=offload_middleware._OffloadInput)`; LangGraph drops every other key from run input before the node sees state. This scopes run input only — an out-of-run state update (`POST /threads/{id}/state`) resolves its writes against the full state schema and is unaffected. Note the restriction comes from the input schema and *not* from the `PrivateStateAttr` / `OmitFromInput` markers on the state schema's other channels: those are honored by `create_agent`, not by a raw `StateGraph`, so without `_OffloadInput` a local caller could set the compaction cutoff via `_summarization_event` or inflate the thread's recorded spend via the additive `_session_cost_usd`. +- **Key note**: The `messages` channel cannot be dropped from the input schema — the driver has to replay the thread's messages as run input — and against the real server that replay is **authoritative for the channel: it replaces the conversation rather than merging into it**. Measured, not inferred: streaming `{"messages": []}` at an 8-message thread leaves it with 0 messages. (The `add_messages` reducer predicts the opposite, and an in-process checkpointer does behave that way, so this is only observable against a live server.) A localhost caller that reaches the `offload` graph can therefore truncate or rewrite a thread's conversation, not merely append to it. This remains a narrower escalation than the pre-existing consequence of `noop` auth, which already permits injecting messages on the `agent` graph. Client-side, the driver defends the two consequences: it re-reads thread state immediately before the run so a stale snapshot is never written back, and the node raises on an empty `messages` rather than reporting the resulting wipe as "already compact". #### TB11: Config File → Code Execution diff --git a/libs/code/deepagents_code/_cli_context.py b/libs/code/deepagents_code/_cli_context.py index 422c43774d..d7a0802cd4 100644 --- a/libs/code/deepagents_code/_cli_context.py +++ b/libs/code/deepagents_code/_cli_context.py @@ -140,10 +140,14 @@ class CLIContext(TypedDict, total=False): """Current user-turn ID for binding trusted interactive responses.""" offload_tool_call_id: str | None - """The sole tool-call ID authorized during a server-driven `/offload` run. + """The sole tool-call ID authorized during a seeded `/offload` run. This is set by the client, not graph state, so model-generated calls cannot grant themselves permission to execute during the hidden compaction turn. + + Only the seeded driver sets it (local in-process agents). A server-backed + `/offload` runs the dedicated operation graph, which has no tool node and so + nothing to authorize; it leaves this `None`. """ hooks_snapshot_id: str | None diff --git a/libs/code/deepagents_code/agent.py b/libs/code/deepagents_code/agent.py index 2f335a62ab..234d072ee4 100644 --- a/libs/code/deepagents_code/agent.py +++ b/libs/code/deepagents_code/agent.py @@ -100,7 +100,11 @@ _artifacts_root, _offload_fallback_root, ) -from deepagents_code.offload_middleware import _create_cli_compaction_middleware +from deepagents_code.offload_middleware import ( + OffloadServerResources, + _create_cli_compaction_middleware, + attach_offload_resources, +) from deepagents_code.plugins.adapters.skills_middleware import PluginSkillsMiddleware from deepagents_code.project_utils import ProjectContext, get_server_project_context from deepagents_code.reliable_rubric import ReliableRubricMiddleware @@ -2857,7 +2861,7 @@ def _subagent_cli_middleware( trusted_root, narrow_allow_list = auto_mode_config # An explicit argument wins; otherwise the env var / `config.toml` # preference is read here, where agent construction already runs off the - # blockbuster-guarded server loop (see `server_graph._make_graph`). + # blockbuster-guarded server loop (see `server_graph._make_graphs`). classifier_model = ( auto_classifier_model if auto_classifier_model is not None @@ -2887,7 +2891,21 @@ def _subagent_cli_middleware( from deepagents_code.hooks.server_middleware import ServerHooksMiddleware hooks_cwd = Path(effective_cwd) if effective_cwd is not None else Path.cwd() - agent_middleware.append(ServerHooksMiddleware(cwd=hooks_cwd, mcp_tools=mcp_tools)) + server_hooks_middleware = ServerHooksMiddleware(cwd=hooks_cwd, mcp_tools=mcp_tools) + agent_middleware.append(server_hooks_middleware) + + # The dedicated server-side `/offload` graph has no model or tool nodes, so + # it reuses these exact instances through the shared composite backend: the + # same summarizer/backend setup, and the same lifecycle implementation for + # the `PreCompact`/`PreToolUse` dispatch it runs against an in-memory forced + # tool call. + attach_offload_resources( + composite_backend, + OffloadServerResources( + compaction=compaction_middleware, + hooks=server_hooks_middleware, + ), + ) if fs_tools is not None: # `fs_tools` is an explicit allowlist here (`--allow-fs-tools all` and an diff --git a/libs/code/deepagents_code/app.py b/libs/code/deepagents_code/app.py index 6c376885fb..562bd44d97 100644 --- a/libs/code/deepagents_code/app.py +++ b/libs/code/deepagents_code/app.py @@ -453,6 +453,45 @@ def _warn_discarded_goal_channels(state_values: dict[str, Any]) -> list[str]: user why an unrelated next turn might fail and how to recover. """ +_OFFLOAD_REBIND_WARNING = ( + "The thread could not be re-associated with the main agent. Commands that " + "write thread state (/goal, /rubric) may error until you send a new message." +) +"""Shown when the post-`/offload` graph rebind fails. + +Distinct from `_OFFLOAD_WEDGE_WARNING`: nothing is wedged, but the thread is +still bound to the `offload` graph, so an out-of-run +`aupdate_state(as_node="model")` resolves against a graph with no `model` node +until the next agent turn rebinds it. Naming the affected commands beats letting +an unrelated `/goal` fail with no explanation. + +Deliberately says nothing about whether the offload itself succeeded: it is +mounted after a success, after a genuine no-op, and after an unconfirmed or +hook-stopped attempt, and each of those mounts its own outcome line first. +""" + +_OFFLOAD_MAX_RESUME_ROUNDS = 10 +"""Bound on interrupt fulfill/resume rounds during an `/offload`. + +A hook that keeps raising fresh interrupt payloads would otherwise loop the +fulfill/resume cycle forever. On the operation-graph path an `ask` permission is +*not* what this bounds — that exits immediately through the `unresolvable` +branch, since that graph has no HITL middleware to route it. On the seeded path +the agent graph does have HITL middleware, so its approval retries are exactly +what the bound is for. + +Shared by both drivers so the bound is one number rather than two literals that +can drift apart, but the two treat exhaustion differently and deliberately so: +the operation graph runs its hook boundary *before* compaction, so hitting the +cap means nothing committed and the round is a failure, while the seeded +driver's compaction has already committed by then and the cap only leaves a +messy tail. +""" + + +class _MissingOffloadGraphError(Exception): + """The server does not expose the optional `offload` operation graph.""" + def _summarization_cutoff(event: Any) -> int: # noqa: ANN401 """Return the absolute cutoff index of a `_summarization_event`. @@ -464,11 +503,11 @@ def _summarization_cutoff(event: Any) -> int: # noqa: ANN401 Returns: The `cutoff_index`, or `0` when the event is missing or malformed. """ - if isinstance(event, dict): - cutoff = event.get("cutoff_index") - if isinstance(cutoff, int): - return cutoff - return 0 + from deepagents_code.offload_middleware import _event_cutoff + + # Shared with the operation graph's own no-advance check so the client's + # "did the cutoff move?" test and the node's cannot drift apart. + return _event_cutoff(event) def _effective_conversation(messages: list[Any], event: Any) -> list[Any]: # noqa: ANN401 @@ -546,11 +585,14 @@ def _is_tool_message(msg: Any) -> bool: # noqa: ANN401 def _find_compaction_failure(messages: list[Any]) -> str | None: """Return a persisted forced-compaction failure message, if present. - `/offload` primarily detects tool failures from the live message stream, - but a stream hiccup (or an update-injected `ToolMessage` that never surfaces - on the `messages` stream) can drop that signal even though the failure - `ToolMessage` still lands in durable state. Scanning committed state closes - that gap so a genuine failure is not misreported as "nothing to offload". + This covers the seeded driver used for local in-process agents, which is + the only `/offload` path that produces a `ToolMessage` at all. It primarily + detects tool failures from the live message stream, but a stream hiccup (or + an update-injected `ToolMessage` that never surfaces on the `messages` + stream) can drop that signal even though the failure `ToolMessage` still + lands in durable state. Scanning committed state closes that gap so a + genuine failure is not misreported as "nothing to offload". The server-side + operation graph raises instead, so it never reaches this scan. The caller passes only the messages produced by the *current* `/offload` attempt (the tail after the pre-seed prefix). This matters because the @@ -2851,6 +2893,14 @@ class DeepAgentsApp(App): Hydration now runs on every scroll-offset delta, so a persistent failure would otherwise notify on every scroll tick.""" + _offload_rebind_failed: bool = False + """Set when the post-`/offload` main-graph rebind failed. + + Recorded by `_drive_offload_operation_graph` and consumed by + `_handle_offload`, because only the caller knows whether it is about to + report the offload as having finished -- `_OFFLOAD_REBIND_WARNING` says it + did, and pairing that with a failure message would contradict it.""" + BINDINGS: ClassVar[list[BindingType]] = [ Binding("escape", "interrupt", "Interrupt", show=False, priority=True), Binding( @@ -14683,13 +14733,14 @@ async def _get_conversation_token_count(self) -> int | None: async def _handle_offload(self) -> None: """Offload older messages to free context window space. - Runs offload SERVER-SIDE by driving the agent's own - `compact_conversation` tool (with `force=True`) instead of - reimplementing summarization + persistence client-side. This keeps the - offloaded archive in the agent's composite backend so it is readable - via `read_file` in every run mode (server, sandbox, in-process). The - client only seeds the tool call, approves the resulting HITL interrupt, - drains the run, and renders the persisted `_summarization_event`. + Server-backed agents run the dedicated `offload` operation graph (no + model node, no synthetic tool call); local in-process `Pregel` agents + drive the agent's own `compact_conversation` tool (with `force=True`) + via a seeded tool call instead. Either way the offloaded archive lands + in the agent's composite backend so it is readable via `read_file` in + every run mode (server, sandbox, in-process). The seeded path's tool is + HITL-gated with no approval bypass, so its driver answers the resulting + interrupt itself, approving only the forced `compact_conversation` call. """ from langchain_core.messages.utils import count_tokens_approximately @@ -14733,15 +14784,51 @@ async def _handle_offload(self) -> None: _effective_conversation(before_messages, prior_event) ) + # Local `Pregel` agents have no server operation graph, so they + # drive the seeded `compact_conversation` tool call in-process + # instead. Custom server graphs likewise omit the optional + # operation graph and fall back to this established path when the + # server reports that `offload` is absent. + local_seeded = self._remote_agent() is None # Own the seeded tool-call id here so a failed run can clean up the # committed-but-unanswered seed (see `_remove_unanswered_offload_seed`). - seed_tool_call_id = str(uuid.uuid4()) + seed_tool_call_id = str(uuid.uuid4()) if local_seeded else None + # Stale from a previous `/offload` otherwise; only the operation + # graph writes it, and only the reporting paths below read it. + self._offload_rebind_failed = False try: - tool_error = await self._drive_server_side_compaction( - config, seed_tool_call_id - ) + if seed_tool_call_id is not None: + tool_error = await self._drive_local_seeded_compaction( + config, seed_tool_call_id + ) + else: + try: + tool_error = await self._drive_offload_operation_graph( + config, state_values + ) + except _MissingOffloadGraphError: + # `generate_langgraph_json()` preserves custom graph + # references without requiring an undocumented paired + # operation graph. They used seeded compaction before + # the operation graph existed, so retain that behavior + # rather than making `/offload` depend on a graph the + # custom server never registered. + local_seeded = True + seed_tool_call_id = str(uuid.uuid4()) + tool_error = await self._drive_local_seeded_compaction( + config, seed_tool_call_id + ) except ClientHookStopError: + # The seeded driver mounts the stop reason itself before + # raising; the operation graph's hook fulfillment has no such + # guarantee, so say something rather than clearing the spinner + # with no output at all. + if not local_seeded: + await self._mount_message( + ErrorMessage("Offload stopped by a hook.") + ) + await self._warn_if_offload_rebind_failed() return except Exception as stream_error: # A server graph can checkpoint the tool-node update before a @@ -14758,19 +14845,26 @@ async def _handle_offload(self) -> None: "Failed to reconcile state after offload stream error", exc_info=True, ) - if not await self._remove_unanswered_offload_seed( - config, seed_tool_call_id + if ( + seed_tool_call_id is not None + and not await self._remove_unanswered_offload_seed( + config, seed_tool_call_id + ) ): await self._mount_message(ErrorMessage(_OFFLOAD_WEDGE_WARNING)) raise stream_error from state_error reconciled_event = new_state.get("_summarization_event") if _summarization_cutoff(reconciled_event) <= prior_cutoff: - # Compaction did not commit, so the seeded tool call was - # never answered. Remove it before re-raising so a failed - # `/offload` cannot wedge the thread with a dangling + # Compaction did not commit. The seeded driver may have left + # its tool call unanswered; remove it before re-raising so a + # failed `/offload` cannot wedge the thread with a dangling # `tool_use` that the model API rejects on the next turn. - if not await self._remove_unanswered_offload_seed( - config, seed_tool_call_id + # The operation graph commits no seed to clean up. + if ( + seed_tool_call_id is not None + and not await self._remove_unanswered_offload_seed( + config, seed_tool_call_id + ) ): await self._mount_message(ErrorMessage(_OFFLOAD_WEDGE_WARNING)) raise @@ -14787,6 +14881,28 @@ async def _handle_offload(self) -> None: # (the archive now lives in the agent's own backend, not a # client-local directory the server can never read). new_state = await self._get_thread_state_values(self._lc_thread_id) + if not new_state and not local_seeded: + # On the operation-graph path this read is the *only* + # evidence of the outcome -- there is no `ToolMessage` to + # fall back on. `_get_thread_state_values` collapses a + # missing snapshot (a 404 after the run rebound the thread, + # a server restart, an un-flushed checkpoint) to `{}`, which + # is indistinguishable from "no event" and would be reported + # as the actively-wrong "already compact". + logger.warning( + "Offload completed but the thread state read came back " + "empty; cannot confirm the result" + ) + await self._mount_message( + ErrorMessage( + "Offload finished, but its result could not be " + "confirmed — the thread state could not be read " + "back. Run /context to check whether the " + "conversation was compacted." + ) + ) + await self._warn_if_offload_rebind_failed() + return # The compaction run's summary model spend is priced and committed by # the graph, so the state just read is the complete total. self._sync_session_cost_from_state(new_state) @@ -14795,50 +14911,78 @@ async def _handle_offload(self) -> None: if new_event is None or new_cutoff <= prior_cutoff: # A failure and a genuine no-op both leave `_summarization_event` - # unchanged. Stream-based detection can miss the failure - # `ToolMessage` (e.g. an update-injected message that never - # surfaces on the `messages` stream), so cross-check committed - # state before concluding there was nothing to do. + # unchanged. On the seeded path, stream-based detection can miss + # the failure `ToolMessage` (e.g. an update-injected message that + # never surfaces on the `messages` stream), so cross-check + # committed state before concluding there was nothing to do. The + # operation graph raises on failure and so never reaches here + # with one pending; for it this branch is a true no-op. current_messages = new_state.get("messages", [])[len(before_messages) :] failure = _find_compaction_failure(current_messages) if failure is not None: await self._mount_message(ErrorMessage(failure)) return - # A no-op still commits the synthetic assistant seed and its - # tool result. Restore the exact pre-run conversation so an - # operation reported as doing nothing truly changes nothing. - await self._remove_offload_artifacts( - config, current_messages, prior_event - ) + if local_seeded: + # A no-op seeded run still commits the synthetic assistant + # seed and its tool result. Restore the exact pre-run + # conversation so an operation reported as doing nothing + # truly changes nothing. The operation graph commits no + # such artifacts. + await self._remove_offload_artifacts( + config, current_messages, prior_event + ) # `force=True` bypasses the eligibility gate, so this branch is # reached when there is nothing older than the retention window - # to summarize (effective cutoff 0). It also absorbs the - # degenerate chained case where only the prior summary would be - # re-summarized (effective cutoff 1 -> new_cutoff == prior_cutoff - # via `_compute_state_cutoff`): a fresh event may commit but the - # absolute cutoff does not advance, so "nothing to offload" is - # the correct, if conservative, report. + # to summarize (effective cutoff 0), or in the degenerate + # chained case where only the prior summary would be + # re-summarized (effective cutoff 1 -> the absolute cutoff would + # not advance). The operation graph detects that second case + # itself and returns before spending a model call, so on that + # path the report and the committed state now agree: nothing was + # written. The seeded path still commits a replacement event + # there, which is what `_remove_offload_artifacts` above undoes. await self._mount_message( AppMessage( "Nothing to offload \u2014 the conversation is already " "compact.", ), ) + await self._warn_if_offload_rebind_failed() return + if not local_seeded: + # The seeded driver fires this from inside its drain, keyed on + # the compaction tool result. The operation graph produces no + # tool result, so fire it here instead -- at the same point in + # the lifecycle (compaction has committed), so a configured + # `SessionStart` hook sees `/offload` exactly as it sees + # automatic compaction. + # + # A stop is deliberately not returned on: compaction has already + # committed to the checkpoint, so returning here would leave the + # user with only "stopped by a hook" while their conversation was + # in fact compacted and the status bar kept pre-offload counts. + # `_run_session_start_hook` mounts the stop reason itself, so the + # user sees both facts in lifecycle order. + from deepagents_code.hooks.models.domain import SessionStartCause + + await self._run_session_start_hook(SessionStartCause.COMPACT) + archive_path = ( new_event.get("file_path") if isinstance(new_event, dict) else getattr(new_event, "file_path", None) ) - # Recompute the post-offload size from the ORIGINAL pre-seed - # messages plus the new event. `_effective_conversation` yields + # Recompute the post-offload size from the ORIGINAL pre-run messages + # plus the new event. `_effective_conversation` yields # `[summary, *before_messages[new_cutoff:]]` — the compacted - # conversation without the tool's own machinery (the seeded tool + # conversation without the seeded driver's machinery (the seeded tool # call, the tool result, and the trailing model turn), all of which - # land in `new_state["messages"]` at/after `new_cutoff`. Counting - # `before_messages` keeps this token figure consistent with the - # message counts below and avoids understating the reduction. + # land in `new_state["messages"]` at/after `new_cutoff`. The + # operation graph commits no such artifacts, so for that path this is + # simply the pre-run conversation. Counting `before_messages` keeps + # this token figure consistent with the message counts below and + # avoids understating the reduction. # # This is a client-side approximation for the status bar and is # deliberately not the persisted `_context_tokens` (refreshed from @@ -14900,11 +15044,19 @@ async def _handle_offload(self) -> None: ) ) + await self._warn_if_offload_rebind_failed() self._on_tokens_update(tokens_after) except Exception as exc: # surface offload errors to user + from deepagents_code.client.remote_client import format_agent_exception + logger.exception("Offload failed") - await self._mount_message(ErrorMessage(f"Offload failed: {exc}")) + # The operation graph reports failure by raising, so a server-backed + # `/offload` surfaces a `RemoteException` here. `str()` renders that + # as a raw dict repr; `format_agent_exception` unwraps it. + await self._mount_message( + ErrorMessage(f"Offload failed: {format_agent_exception(exc)}") + ) finally: self._set_agent_running(False) try: @@ -14912,20 +15064,398 @@ async def _handle_offload(self) -> None: except Exception: # best-effort spinner cleanup logger.exception("Failed to dismiss spinner after offload") - async def _drive_server_side_compaction( + async def _warn_if_offload_rebind_failed(self) -> None: + """Mount `_OFFLOAD_REBIND_WARNING` if the post-`/offload` rebind failed. + + Called from every path that returns after the operation graph has run + and does *not* already tell the user the thread may need to reset: a + success, a genuine no-op, an unconfirmed result, and a hook stop. The + drain-error paths are excluded because they say so themselves — and + because the rebind is deliberately skipped while the run is still + suspended, so the flag cannot be set on them. + """ + if not self._offload_rebind_failed: + return + self._offload_rebind_failed = False + await self._mount_message(ErrorMessage(_OFFLOAD_REBIND_WARNING)) + + async def _drive_offload_operation_graph( + self, config: RunnableConfig, state_values: dict[str, Any] + ) -> str | None: + """Run the explicit server-side `/offload` operation graph. + + The operation graph shares the interactive agent's checkpoint and + composite backend but has no model node or HITL middleware. The slash + command is the user's authorization, so what it persists is the + summarization event and the run's cost drain — never an assistant tool + call or tool result in the conversation. (It does build a forced tool + call for the hook dispatch, but only in memory.) + + Args: + config: Config with `configurable.thread_id`. + state_values: Pre-run thread state, refreshed by the re-read below + before use. Only `messages` is replayed as run input, and that + replay *replaces* the channel on a real server (see the comment + at the `stream_input` construction), so it must be current. + Every other channel is deliberately withheld, because replaying + the cost channels corrupts them. + + Returns: + An error string when the run could not be driven to completion + (an interrupt this client cannot answer, a hook that outlasted + `_OFFLOAD_MAX_RESUME_ROUNDS`, or a stream that ended without the + node ever reporting an update), otherwise `None`. Compaction + failures inside the node are *not* reported this way: the graph + has no tool node to carry a message, so it raises and the caller + reconciles against the checkpointed `_summarization_event`. The + `str | None` shape matches the seeded driver so + `_handle_offload` handles both paths uniformly. + + Raises: + _MissingOffloadGraphError: If a custom server graph does not + register the optional `offload` operation graph. + RuntimeError: If the app is not connected to its server graph, or if + current thread state could not be read before the run -- + replaying the pre-run snapshot would delete messages committed + since it was taken. + """ + from langgraph.types import Command + from langgraph_sdk.errors import NotFoundError + + from deepagents_code.config import settings + from deepagents_code.hooks.interrupt import is_hook_interrupt_payload + from deepagents_code.offload_middleware import ( + OFFLOAD_GRAPH_NODE, + _OffloadInput, + ) + + remote = self._remote_agent() + if remote is None: + msg = "The explicit /offload operation requires the server graph." + raise RuntimeError(msg) + # Re-read rather than trusting the caller's snapshot. That snapshot is + # taken before `_set_agent_running(True)`, so a turn committed in + # between would be missing from it -- and because the replay below + # *replaces* the `messages` channel, a stale list is not a stale read + # but a destructive write. Do NOT fall back to the caller's copy when + # the re-read cannot produce current state: the gap that motivates the + # re-read (an externally managed/shared remote thread, or a turn + # completing concurrently with the initial read) is exactly when the + # caller's copy is missing messages, so replaying it would delete the + # messages the re-read was meant to preserve. Abort instead. + thread_id = self._lc_thread_id + if thread_id is None: + await remote.aensure_thread(dict(config)) + else: + try: + fresh_values = await self._get_thread_state_values(thread_id) + except Exception as exc: + msg = ( + "Could not refresh thread state before /offload; not " + "replaying the stale pre-run snapshot." + ) + raise RuntimeError(msg) from exc + # `_get_thread_state_values` collapses a missing snapshot (a 404 + # after a rebind, a server restart, an un-flushed checkpoint) to + # `{}`, and the fallback snapshot was read before this run marked + # the agent busy, so it cannot be trusted to fill in. + if not fresh_values: + msg = ( + "Could not read current thread state before /offload " + "(the state read came back empty); not replaying the " + "stale pre-run snapshot." + ) + raise RuntimeError(msg) + state_values = fresh_values + + stream_context = CLIContext( + model=self._effective_model_spec(), + model_params=self._model_params_override or {}, + profile_overrides=self._profile_override or {}, + model_context_limit=settings.model_context_limit, + thread_id=self._lc_thread_id, + ) + self._hooks.apply_graph_context(stream_context) + # Replay the thread's messages -- and *only* the messages -- as the run + # input. Against a real server the run input is *authoritative* for this + # channel: it replaces the checkpointed conversation rather than merging + # into it. Measured, not assumed -- streaming `{"messages": []}` against + # a live server takes an 8-message thread to 0 and still reports + # "already compact", while replaying the list leaves all 8 intact. + # (`messages` carries the `add_messages` reducer, so a plain reducer + # reading predicts the opposite, and an in-process checkpointer does + # resolve it that way -- which is why unit tests alone cannot catch + # this. `test_offload_server_side.py` pins the real behavior.) + # `astream(None)` is not an option either -- it never starts a run on + # `RemoteAgent`. + # + # Two consequences follow, and both are defended rather than assumed: + # the snapshot must be *fresh* (a stale or partial read would be written + # back over the live conversation), so it is re-read here rather than + # reusing `_handle_offload`'s pre-run read; and an empty replay is + # destructive, so the node raises on an empty `messages` instead of + # reporting a no-op. + # + # Known cost of the replay, accepted rather than fixed: the round-trip + # is not byte-identical. Re-serializing each `AIMessage` adds + # `invalid_tool_calls` and `usage_metadata` keys inside + # `additional_kwargs`, which is also where provider-specific fields + # (Anthropic `cache_control`, thinking blocks, citations) live, and the + # rewrite compounds across repeated offloads. Removing it needs the node + # to read `messages` from the checkpoint the way it already reads + # `_summarization_event`, which this client cannot arrange alone. + # + # Nothing else may be replayed. Every other channel either reaches the + # node through the checkpoint anyway or is actively corrupted by a + # replay: `_session_cost_usd` reduces with `operator.add`, so echoing the + # checkpointed total back as input *doubles* the thread's persisted cost + # on every `/offload`, and `_session_cost_transfers` (`operator.or_`) + # resurrects already-settled subagent transfers. `_summarization_event` + # is private to the schema (see `offload_middleware._OffloadState`) and + # its embedded summary message is a shape the server cannot deserialize. + # Annotated with the graph's own input type so dropping `messages` is a + # type error here rather than a silent conversation wipe there. + replay_messages = state_values.get("messages", []) + if not replay_messages: + # `_handle_offload` already rejected an empty thread, so an empty + # list here means the state re-read above came back partial. Writing + # it would truncate the conversation and report success. + msg = ( + "Offload aborted: the thread's messages could not be read back " + "before the run. Nothing was changed; try again." + ) + raise RuntimeError(msg) + stream_input: _OffloadInput = {"messages": replay_messages} + + # `stream_mode=["updates"]` only: unlike the seeded driver, this path + # cannot itemize the offload under the `offload` kind in `/cost`. The + # summarizer runs as a direct model invoke inside the node and passes its + # own `config`, which replaces the ambient run config the `messages` + # stream is built from, so no message chunk ever reaches this client. + # The spend itself is not lost: the process-wide recorder captures the + # request and the graph's + # `CostTrackingMiddleware` prices it onto `_session_cost_usd`, which + # `_sync_session_cost_from_state` reads back. Only the per-kind token + # breakdown is unavailable, and recovering it needs a usage channel on + # the checkpoint that does not exist yet. + offload_agent = remote.for_graph("offload") + # Configured `PreCompact`/`PreToolUse` server hooks interrupt the graph + # at the hook boundary rather than returning a decision; fulfill each + # interrupt against the client hook engine and resume, mirroring the + # seeded driver's drain loop. + resume_input: Any = cast("Any", stream_input) + drain_error: str | None = None + # Interrupt detection is all this loop does, so a chunk shape it does + # not recognize is indistinguishable from a clean run: `pending` stays + # empty, the loop exits with no error, and the caller reads an + # unadvanced event as "already compact" while the run is still paused + # server-side. Track whether the node ever reported an update so that + # case becomes a reported failure instead. + saw_node_update = False + try: + rounds = 0 + while True: + pending: dict[str, Any] = {} + unresolvable: list[str] = [] + async for chunk in offload_agent.astream( + resume_input, + stream_mode=["updates"], + # Load-bearing, despite this graph having no subgraphs. + # `RemoteAgent.astream` only yields the documented + # `(namespace, mode, data)` 3-tuple when it is set; without + # it the chunk arrives as `("updates", {...}, None)`, so the + # unpacking below binds `mode` to the payload dict and every + # chunk falls through the `mode != "updates"` filter. That + # discards interrupts silently: hook fulfillment never runs, + # the run stays paused server-side, and the caller reads an + # unadvanced event as "already compact". Unit tests cannot + # see it -- they feed the loop chunks directly. The + # `saw_node_update` check below is what makes a regression + # here loud instead of silent. + subgraphs=True, + config=config, + context=stream_context, + # Accepted for signature parity and documented as ignored by + # `RemoteAgent.astream` (the server owns durability), so + # this is not a guarantee that the event is committed when + # the stream ends. The caller re-reads committed state and + # reconciles rather than trusting the stream. + durability="exit", + ): + if not isinstance(chunk, tuple) or len(chunk) != 3: # noqa: PLR2004 # (namespace, mode, data) + logger.debug( + "Ignoring unrecognized /offload stream chunk: " + "type=%s len=%s", + type(chunk).__name__, + len(chunk) if isinstance(chunk, tuple) else "n/a", + ) + continue + _namespace, mode, data = chunk + if mode != "updates" or not isinstance(data, dict): + logger.debug( + "Ignoring /offload stream chunk: mode=%r data=%s", + mode, + type(data).__name__, + ) + continue + if OFFLOAD_GRAPH_NODE in data: + saw_node_update = True + for interrupt_obj in data.get("__interrupt__") or []: + iid = getattr(interrupt_obj, "id", None) + value = getattr(interrupt_obj, "value", None) + if iid and is_hook_interrupt_payload(value): + pending[iid] = await self._hooks.fulfill_interrupt(value) + else: + # Two distinct conditions share this branch and the + # user-facing message below only fits one, so record + # which it was: a hook payload with no id is a + # transport bug, while a non-hook payload is the + # `ask`-permission `HITLRequest` case. + unresolvable.append( + f"id={iid!r} " + f"hook_payload={is_hook_interrupt_payload(value)} " + f"type={type(value).__name__}" + ) + if unresolvable: + # Not every interrupt this graph can raise is a hook + # invocation this client can answer: a `PreToolUse` hook + # returning an `ask` permission makes the hook middleware + # itself raise a plain `HITLRequest`, and the operation graph + # has no HITL middleware to route it. Report rather than + # break silently -- the run stays paused either way, and the + # caller would otherwise see an unchanged event and tell the + # user the conversation is already compact. + logger.warning( + "Offload received %d interrupt(s) this client cannot " + "answer; leaving the run paused. Unanswerable: %s. " + "Discarding %d already-executed hook fulfillment(s): %s", + len(unresolvable), + "; ".join(unresolvable), + len(pending), + ", ".join(pending) or "none", + ) + drain_error = ( + "Offload could not complete: the run requested an " + "approval that /offload cannot answer. Check your " + "PreToolUse/PreCompact hook configuration for " + "`compact_conversation`. Any hooks that already ran for " + "this attempt will run again on the next one. Send a " + "new message to continue; the thread may need to reset." + ) + break + if not pending: + break + rounds += 1 + if rounds > _OFFLOAD_MAX_RESUME_ROUNDS: + # Nothing has committed yet -- the hook boundary runs before + # compaction -- so unlike the seeded driver's equivalent + # this is a failure, not a completed offload with a messy + # tail. Report it rather than letting the caller read the + # unchanged event as "nothing to offload". + logger.warning( + "Offload exceeded %d resume rounds; leaving %d " + "interrupt(s) unresolved and discarding their " + "already-executed fulfillment(s): %s", + _OFFLOAD_MAX_RESUME_ROUNDS, + len(pending), + ", ".join(pending), + ) + drain_error = ( + "Offload could not complete: a configured hook kept " + f"interrupting after {_OFFLOAD_MAX_RESUME_ROUNDS} " + "rounds. Any hooks that already ran for this attempt " + "will run again on the next one. Send a new message to " + "continue; the thread may need to reset." + ) + break + resume_input = Command(resume=pending) + if drain_error is None and not saw_node_update: + # The stream ended cleanly, raised nothing, and never carried an + # update from the graph's only node. Either the node did not run + # or the SDK's chunk shape drifted out from under the filters + # above; both leave the run paused server-side, and reporting + # success would surface as the actively-wrong "already compact". + logger.warning( + "Offload stream completed without any %r node update; " + "treating the run as incomplete", + OFFLOAD_GRAPH_NODE, + ) + drain_error = ( + "Offload could not complete: the server finished the run " + "without reporting a result. Your conversation is " + "unchanged. Run /context to confirm, then try again." + ) + except NotFoundError as exc: + # The named graph is optional: custom `graph_ref` configurations + # intentionally register only `agent`. Keep the HTTP detail out of + # the driver-selection policy in `_handle_offload` and fall back to + # the seeded path there. + raise _MissingOffloadGraphError from exc + finally: + # A run on the named `offload` graph rebinds the server thread's + # `graph_id`, so a later out-of-run `aupdate_state(as_node="model")` + # through the main client (e.g. `_aupdate_thread_state` for + # goal/rubric commands) would resolve against a graph with no + # `model` node and fail. Restore the binding through the main + # client's own graph so those writes persist again. This must + # update the thread metadata; an empty state update still resolves + # against the current (offload) graph. + # + # Never fatal: the offload result must not be misreported when only + # this bookkeeping fails. But it is not silent either -- the failure + # leaves an unrelated later command (`/goal`, `/rubric`) to fail + # confusingly, so `_OFFLOAD_REBIND_WARNING` is mounted alongside the + # offload's own result rather than replacing it. + # + # Skipped entirely when `drain_error` is set: the run is still + # suspended server-side in that case, and rebinding would re-point + # the thread away from the graph the paused run belongs to. The + # drain-error text already tells the user the thread may need to + # reset, and the next agent turn rebinds anyway. + # + # Recorded here and mounted by `_handle_offload`, not mounted here: + # this message says the offload finished, and only the caller knows + # whether it is about to report that. A stream error that still + # committed the compaction is reconciled into a success, and a + # `drain_error` is reported as a failure -- neither is visible from + # inside this `finally`. + if drain_error is not None: + logger.info( + "Leaving the thread bound to the offload graph: the run is " + "still suspended server-side" + ) + else: + try: + await remote.arebind_thread(config) + except Exception: + logger.warning( + "Failed to restore the thread's main graph association " + "after /offload; the next agent turn will rebind it", + exc_info=True, + ) + self._offload_rebind_failed = True + else: + self._offload_rebind_failed = False + return drain_error + + async def _drive_local_seeded_compaction( self, config: RunnableConfig, seed_tool_call_id: str | None = None ) -> str | None: - """Trigger the server-side `compact_conversation` tool with `force=True`. + """Drive a local agent's `compact_conversation` tool with `force=True`. + + This is `/offload`'s sole implementation for local in-process `Pregel` + agents, including ACP mode — not a deprecated path. Server-backed agents + take `_drive_offload_operation_graph` instead. Seeds an assistant `compact_conversation` tool call attributed to the model node, then advances the graph so the agent's own `ToolNode` - executes the tool. The tool is HITL-gated, so `astream(None)` surfaces - an approval interrupt; only the first forced `compact_conversation` - request is approved here (this is an explicit user-initiated - `/offload`). The runtime context carries the seeded call ID so the - compaction middleware can reject every other tool independently of - HITL configuration, including tools requested by the trailing model - turn. + executes the tool. The tool is HITL-gated and the seed is not an + approval bypass, so `astream(None)` surfaces an approval interrupt; only + the first forced `compact_conversation` request is approved here (this is + an explicit user-initiated `/offload`). The runtime context carries the + seeded call ID so the compaction middleware can reject every other tool + independently of HITL configuration, including tools requested by the + trailing model turn. A first-turn `Command(update=..., goto=...)` is intentionally avoided: the LangGraph API server rebuilds it with `goto=None` and crashes @@ -14991,6 +15521,9 @@ async def _drive_server_side_compaction( # Remote dev servers separate checkpoint persistence from HTTP thread # registration; register before mutating state so the write lands. + # `_handle_offload` only routes local `Pregel` agents here, so this is + # inert on that path; it is kept for direct and test callers, which do + # drive this method against a remote agent. if remote := self._remote_agent(): await remote.aensure_thread( {"configurable": {"thread_id": self._lc_thread_id}} @@ -15163,16 +15696,15 @@ async def _drain(stream_input: Any) -> list[tuple[str, dict[str, Any]]]: # noqa # rejected gated call could prompt another. The middleware blocks # execution even when HITL is disabled; this bound handles HITL retries. try: - max_resume_rounds = 10 pending = await _drain(None) rounds = 0 while pending: rounds += 1 - if rounds > max_resume_rounds: + if rounds > _OFFLOAD_MAX_RESUME_ROUNDS: logger.warning( "Offload exceeded %d resume rounds; leaving %d interrupt(s) " "unresolved", - max_resume_rounds, + _OFFLOAD_MAX_RESUME_ROUNDS, len(pending), ) # Compaction itself already committed in round 1, so the caller @@ -15209,7 +15741,13 @@ async def _remove_offload_artifacts( messages: list[Any], prior_event: object, ) -> None: - """Restore state changed by a no-op `/offload` graph run. + """Restore state changed by a no-op seeded `/offload` run. + + Applies only to the seeded (local-agent) path. The operation graph + commits no synthetic messages and, since it stops before the model call + when the cutoff would not advance, no replacement event either — so it + has nothing for this to undo. + Best-effort: a failed restoration is logged and swallowed rather than raised. The no-op path answers the seed with a valid tool result, so the diff --git a/libs/code/deepagents_code/client/launch/server.py b/libs/code/deepagents_code/client/launch/server.py index 05f6195979..dbead8d3e3 100644 --- a/libs/code/deepagents_code/client/launch/server.py +++ b/libs/code/deepagents_code/client/launch/server.py @@ -167,10 +167,17 @@ def generate_langgraph_json( ) -> Path: """Generate a `langgraph.json` config file for `langgraph dev`. + Registers the interactive `agent` graph, and — for the default `graph_ref` + only — its paired `offload` operation graph for the `/offload` command. + Args: output_dir: Directory to write the config file. graph_ref: Python "module:attribute" reference to the graph, where the attribute is a graph factory (e.g. `make_graph`) or a graph object. + Only the default ref registers a paired `offload` graph; a custom + ref serves `agent` alone. Because `_handle_offload` routes every + server-backed agent to `remote.for_graph("offload")`, `/offload` + fails against a custom ref's missing graph. env_file: Optional path to an env file. checkpointer_path: Import path to an async context manager that yields a `BaseCheckpointSaver`. When set, the server persists checkpoint data @@ -183,6 +190,12 @@ def generate_langgraph_json( "dependencies": ["."], "graphs": {"agent": graph_ref}, } + # The built-in pair shares one `_build_graph_factories` closure and is the + # only documented offload factory. A custom `graph_ref` is not required to + # expose a matching `make_offload_graph`, so do not generate an unresolved + # reference for it. + if graph_ref == "deepagents_code.server_graph:make_graph": + config["graphs"]["offload"] = "deepagents_code.server_graph:make_offload_graph" if env_file: config["env"] = env_file if checkpointer_path: diff --git a/libs/code/deepagents_code/client/remote_client.py b/libs/code/deepagents_code/client/remote_client.py index 85e441e08d..47b3d0ef8a 100644 --- a/libs/code/deepagents_code/client/remote_client.py +++ b/libs/code/deepagents_code/client/remote_client.py @@ -136,6 +136,7 @@ def __init__( self._api_key = api_key self._headers = headers self._graph: Any = None + self._sibling_clients: dict[str, RemoteAgent] = {} def _get_graph(self) -> Any: # noqa: ANN401 """Lazily create the `RemoteGraph` instance. @@ -154,6 +155,38 @@ def _get_graph(self) -> Any: # noqa: ANN401 ) return self._graph + def for_graph(self, graph_name: str) -> RemoteAgent: + """Return a client for another graph served by the same runtime. + + Cached per graph name, mirroring the `_graph` cache: each `RemoteGraph` + builds its own `httpx` async *and* sync client, neither of which is + pooled by the SDK or closed here, so constructing one per call would + leak two connection pools on every use. + + Args: + graph_name: Registered LangGraph graph name. + + Returns: + A client that preserves this connection's URL and credentials. The + same instance is returned for repeated calls with one name, and + `self` when asked for the graph this client already serves. + """ + if graph_name == self._graph_name: + # Otherwise this builds a second client for the graph the receiver + # already serves -- the exact duplicate-connection-pool leak the + # cache exists to prevent. + return self + cached = self._sibling_clients.get(graph_name) + if cached is None: + cached = RemoteAgent( + self._url, + graph_name=graph_name, + api_key=self._api_key, + headers=self._headers, + ) + self._sibling_clients[graph_name] = cached + return cached + async def astream( self, input: dict | Any, # noqa: A002, ANN401 @@ -469,6 +502,56 @@ async def aensure_thread(self, config: dict[str, Any]) -> None: ) raise + async def arebind_thread(self, config: Mapping[str, Any]) -> None: + """Associate an existing remote thread with this client's graph. + + A run through a sibling graph changes the server-side thread's + `graph_id`. Restore it before a caller makes an out-of-run state + update, whose `as_node` is resolved against that association. + + The metadata write merges server-side, so it does not clobber other + thread metadata. + + Args: + config: Config with `configurable.thread_id`. + + Raises: + ValueError: If `thread_id` is not present in `config`. + AttributeError: If the installed LangGraph SDK no longer exposes + `RemoteGraph._validate_client`. + Exception: Any transport or API failure from the metadata update, + logged at WARNING and re-raised for the caller to classify. + """ # noqa: DOC502 — `ValueError` raised by `_require_thread_id` + thread_id = _require_thread_id(config) + graph = self._get_graph() + + # Private SDK accessor. Checked explicitly because the alternative is an + # `AttributeError` that the caller's blanket rebind handler downgrades + # to a warning -- `/offload` would keep "working" while every later + # `/goal` and `/rubric` failed on a mis-bound thread, permanently and + # with nothing connecting the two. + validate_client = getattr(graph, "_validate_client", None) + if validate_client is None: + msg = ( + "RemoteGraph._validate_client is unavailable; the LangGraph SDK " + "changed and threads can no longer be rebound to their graph." + ) + raise AttributeError(msg) + + try: + client = validate_client() + await client.threads.update( + thread_id, metadata={"graph_id": self._graph_name} + ) + except Exception: + logger.warning( + "Failed to associate thread %s with graph %s", + thread_id, + self._graph_name, + exc_info=True, + ) + raise + def with_config(self, config: dict[str, Any]) -> RemoteAgent: # noqa: ARG002 """Return self (config is passed per-call, not stored). diff --git a/libs/code/deepagents_code/configurable_model.py b/libs/code/deepagents_code/configurable_model.py index 0227faf992..6e197431d6 100644 --- a/libs/code/deepagents_code/configurable_model.py +++ b/libs/code/deepagents_code/configurable_model.py @@ -221,7 +221,7 @@ def _resolve_openai_prompt_cache_key_enabled() -> bool: Called once when `ConfigurableModelMiddleware` is constructed. The read is kept off the blockbuster-guarded server loop by the caller: on the server path `create_cli_agent` runs inside `asyncio.to_thread` (see - `server_graph._make_graph`), so the synchronous `config.toml` read happens + `server_graph._make_graphs`), so the synchronous `config.toml` read happens on a worker thread. On an unexpected failure this defaults to enabled: breaking agent diff --git a/libs/code/deepagents_code/offload_middleware.py b/libs/code/deepagents_code/offload_middleware.py index c9279c09bc..9f92af0b1f 100644 --- a/libs/code/deepagents_code/offload_middleware.py +++ b/libs/code/deepagents_code/offload_middleware.py @@ -7,10 +7,13 @@ import logging from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, cast +from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, Protocol, cast +from uuid import NAMESPACE_URL, uuid4, uuid5 from deepagents.backends.protocol import FILE_NOT_FOUND from deepagents.middleware.summarization import ( + SummarizationEvent, + SummarizationState, SummarizationToolMiddleware, create_summarization_middleware, create_summarization_tool_middleware, @@ -19,11 +22,16 @@ ToolRuntime, # noqa: TC002 # inspected for runtime injection ) from langchain_core.exceptions import ContextOverflowError -from langchain_core.messages import ToolMessage +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage from langchain_core.tools import InjectedToolArg, StructuredTool +from langgraph.config import get_config +from langgraph.errors import GraphBubbleUp +from langgraph.graph.message import add_messages from langgraph.types import Command +from typing_extensions import TypedDict from deepagents_code._cli_context import CLIContextSchema +from deepagents_code.cost_tracking import CostState, CostTrackingMiddleware from deepagents_code.hooks.models.domain import ( CompactTrigger, HookEvent, @@ -32,6 +40,7 @@ ) from deepagents_code.hooks.server_middleware import ( _DEFAULT_DEADLINE, + _PRE_TOOL_STATE_KEY, _event_enabled, _hook_context, _invoke_hook, @@ -42,6 +51,7 @@ if TYPE_CHECKING: from collections.abc import Awaitable, Callable + from deepagents.backends.composite import CompositeBackend from deepagents.backends.protocol import ( BackendProtocol, EditResult, @@ -55,20 +65,181 @@ ModelResponse, ) from langchain.chat_models import BaseChatModel + from langgraph.graph.state import CompiledStateGraph from langgraph.prebuilt.tool_node import ToolCallRequest + from langgraph.runtime import Runtime + + from deepagents_code.hooks.server_middleware import ServerHooksMiddleware logger = logging.getLogger(__name__) +class _OffloadState(CostState, SummarizationState, total=False): + """Checkpoint channels the explicit forced-compaction graph reads and writes. + + `_summarization_event` is inherited from `SummarizationState` rather than + re-declared so it keeps the SDK's own annotation (`NotRequired`, `| None`, + and the `PrivateStateAttr` marker) instead of a divergent copy that claimed + the value is always present. + + Note this schema does *not* by itself keep any channel out of the graph's + input: `PrivateStateAttr` / `OmitFromInput` are honored by `create_agent`, + not by a raw `StateGraph`, so every declared channel would otherwise be + client-writable. `_OffloadInput` is what actually restricts the surface. + + `total=False` describes the inherited shape only -- this body declares no + keys of its own, so the modifier is deliberate documentation rather than a + constraint on anything written here. + """ + + +class _OffloadInput(TypedDict): + """The only channel a caller may write as *run input* to an `/offload` run. + + Passed as `StateGraph(input_schema=...)` so the restriction is enforced by + the graph rather than by the client's discipline in what it sends. Replaying + any other channel is actively harmful: `_session_cost_usd` reduces with + `operator.add`, so echoing back the checkpointed total would double the + thread's recorded spend on every `/offload`, and `_session_cost_transfers` + (`operator.or_`) would resurrect settled subagent transfers. Writable + `_summarization_event` would additionally let a local caller set the + compaction cutoff directly (see THREAT_MODEL TB10). + + Scope: this narrows *run input* only. An out-of-run state update + (`POST /threads/{id}/state`, i.e. `aupdate_state`) resolves its writes + against the full state schema, so it is unaffected by this type — which is + also why `RemoteAgent.arebind_thread` exists. + + `total=True` on purpose: `messages` is mandatory, and the one construction + site (`app._drive_offload_operation_graph`) annotates its literal with this + type so omission is a type error there. That check is worth having but is + not the real guarantee: `.get("messages", [])` at the call site satisfies + the type with `[]`, and against a real server an empty list is *destructive* + rather than merely useless — run input replaces this channel outright (see + the replay comment at that call site). The enforcement that matters is the + node's own empty-`messages` guard, which raises instead of silently + reporting "already compact"; `test_offload_server_side.py` pins the + server-side behavior end to end. + + The node still reads the full `_OffloadState` from the checkpoint; only the + run *input* is narrowed. + """ + + messages: Annotated[list[AnyMessage], add_messages] + + +class OffloadServerResources(NamedTuple): + """Middleware the `/offload` operation graph shares with the agent graph.""" + + compaction: CLICompactionMiddleware + """Compaction implementation bound to the agent's composite backend.""" + + hooks: ServerHooksMiddleware + """Lifecycle middleware carrying the `PreCompact`/`PreToolUse` boundary. + + Not optional: `create_cli_agent` mounts this middleware unconditionally, and + an absent instance would make the operation graph skip the only hook gate + `/offload` still crosses. + """ + + +_OFFLOAD_RESOURCES_ATTR = "_cli_offload_resources" +"""Attribute carrying `OffloadServerResources` on the composite backend. + +`create_cli_agent` builds this middleware but returns only the agent and the +backend, and every call site unpacks that pair positionally (3 in +`deepagents_code`, 61 in `tests`), so widening the return would churn the whole +test suite for one server-only consumer. The backend is the one object both the +agent graph and the separately-resolved `offload` graph already hold, which +makes it the carrier. Reaching it through `attach_offload_resources` / +`offload_resources_from` instead of a bare `getattr` keeps the attribute name in +one place and the unchecked write behind a single type assertion. +""" + + +def attach_offload_resources( + backend: CompositeBackend, resources: OffloadServerResources +) -> None: + """Publish the operation graph's middleware on the shared backend. + + Args: + backend: Composite backend returned alongside the agent graph. + resources: Middleware instances the `offload` graph must reuse. + + Raises: + ValueError: If the compaction middleware is bound to a different + backend than the one it is being attached to. The whole point of + this carrier is that `/offload` archives into the *agent's* backend; + a mis-wired pair would instead write history somewhere the agent + cannot read it, and the symptom (missing history) would surface only + much later. Checked here so the failure lands at construction. + """ + bound = getattr(resources.compaction._summarization, "_backend", None) + if bound is not None and bound is not backend: + msg = ( + "Offload compaction middleware is bound to a different backend than " + "the one it is being published on; /offload would archive into " + "storage the agent cannot read." + ) + raise ValueError(msg) + if getattr(backend, _OFFLOAD_RESOURCES_ATTR, None) is not None: + # Last-write-wins is the existing behavior and safe for the one + # production caller, but a second `create_cli_agent` sharing a backend + # would silently re-point `/offload` at whichever ran last. + logger.warning( + "Replacing offload resources already published on this backend; " + "/offload will use the most recently attached middleware" + ) + setattr(backend, _OFFLOAD_RESOURCES_ATTR, resources) + + +def offload_resources_from(backend: CompositeBackend) -> OffloadServerResources | None: + """Read back the middleware published by `attach_offload_resources`. + + Args: + backend: Composite backend returned alongside the agent graph. + + Returns: + The published resources, or `None` when the backend carries none (so + the caller can fail with its own message rather than an + `AttributeError`). + """ + resources = getattr(backend, _OFFLOAD_RESOURCES_ATTR, None) + return resources if isinstance(resources, OffloadServerResources) else None + + +def _event_cutoff(event: object) -> int: + """Return the absolute cutoff index carried by a `_summarization_event`. + + Args: + event: A `_summarization_event` mapping (as persisted in state), or + `None`. + + Returns: + The `cutoff_index`, or `0` when the event is missing or malformed. + """ + if isinstance(event, dict): + cutoff = event.get("cutoff_index") + if isinstance(cutoff, int): + return cutoff + return 0 + + COMPACTION_FAILURE_PREFIX = "Compaction failed" """Stable prefix for forced-compaction failure tool messages. -`/offload` drives the tool server-side and can only observe the resulting -`ToolMessage` text across the LangGraph server boundary, so it keys failure -detection on this prefix. Owning the literal here means the producer -(`_forced_compact_error`) and both consumers (`app._drive_server_side_compaction` -live-stream detection and `app._find_compaction_failure` committed-state scan) -reference one constant instead of re-hardcoding the wording independently. +The seeded driver drives the tool rather than calling it, so the only failure +signal it gets back is the resulting `ToolMessage` text; it therefore keys +failure detection on this prefix. (The prefix predates the split of `/offload` +into two paths, when that message always crossed the LangGraph server boundary. +The driver is now used for local in-process agents, which cross no such +boundary, but it still only sees message text.) Owning the literal here means +the producers +(`_forced_compact_error` and the operation graph's node, which reuses the prefix +in the `RuntimeError` it raises) and both consumers +(`app._drive_local_seeded_compaction` live-stream detection and +`app._find_compaction_failure` committed-state scan) reference one constant +instead of re-hardcoding the wording independently. Note: this value is deliberately identical to the leading text of the SDK's own model-initiated compaction-failure message, so a failure emitted by either path @@ -142,11 +313,33 @@ class RuntimeModelConfig(NamedTuple): context_limit: int | None -def _runtime_model_config(runtime: ToolRuntime) -> RuntimeModelConfig: - """Read the active model configuration from a tool runtime. +class _HasRunContext(Protocol): + """Anything carrying a per-run context object. + + The compaction helpers read `context` and nothing else, so they accept both + the `ToolRuntime` injected into the tool and the plain LangGraph `Runtime` + the operation graph's node receives -- which is not a `ToolRuntime`. Stating + the dependency this narrowly means a helper that starts touching, say, + `tool_call_id` fails to type-check instead of breaking the operation graph + at runtime. + """ + + @property + def context(self) -> object: + """The run's context object. + + Typed as `object` rather than `Any`: every consumer narrows the shape + with `isinstance` before touching it, so `object` type-checks the same + code while still rejecting an unnarrowed attribute access. + """ + ... + + +def _runtime_model_config(runtime: _HasRunContext) -> RuntimeModelConfig: + """Read the active model configuration from a run context carrier. Args: - runtime: Runtime injected into the compaction tool. + runtime: Runtime carrying the current `CLIContext`. Returns: The active model specification, invocation parameters, profile @@ -161,10 +354,13 @@ def _runtime_model_config(runtime: ToolRuntime) -> RuntimeModelConfig: context_limit=context.model_context_limit, ) if isinstance(context, dict): - model = context.get("model") - params = context.get("model_params") - profile_overrides = context.get("profile_overrides") - context_limit = context.get("model_context_limit") + # The remote boundary delivers the context as JSON, so the keys are + # strings; the values stay unknown and are narrowed individually below. + fields = cast("dict[str, Any]", context) + model = fields.get("model") + params = fields.get("model_params") + profile_overrides = fields.get("profile_overrides") + context_limit = fields.get("model_context_limit") return RuntimeModelConfig( model_spec=model if isinstance(model, str) else None, model_params=dict(params) if isinstance(params, dict) else {}, @@ -349,6 +545,17 @@ class CLICompactionMiddleware(SummarizationToolMiddleware): private `force` input is used only by the user-initiated `/offload` path, which must compact whenever messages exceed the retention window even when the conversation has not reached the SDK's proactive eligibility gate. + + Three entry points, with different error semantics: + + - automatic compaction, on the SDK's own gated path; + - `_run_forced_compact` / `_arun_forced_compact`, the tool-node paths used + by the seeded `/offload` driver, which report failure by *returning* a + `ToolMessage` because a tool node must always answer its call; + - `arun_forced_compaction_update`, the non-tool entry point the `/offload` + operation graph calls, which *raises* instead — that graph has no tool + node to carry a message, so its node turns the exception into a + client-visible error. """ @property @@ -624,7 +831,7 @@ def _guarded_backend(self) -> BackendProtocol: return cast("BackendProtocol", _ArchiveReadGuard(self._summarization._backend)) def _summarization_for_runtime( - self, runtime: ToolRuntime + self, runtime: _HasRunContext ) -> SummarizationMiddleware: """Build a summarizer for the active runtime model when overridden. @@ -667,10 +874,22 @@ def _summarization_for_runtime( context_limit, exc_info=True, ) - backend = self._summarization._backend - summarization = create_summarization_middleware(model, backend) - summarization._backend = self._guarded_backend() - return summarization + # Never pass the `_ArchiveReadGuard` wrapper to the constructor: the SDK + # resolves the archive prefix once in `__init__` via + # `backend.artifacts_root if isinstance(backend, CompositeBackend)`, and + # the guard is not a `CompositeBackend`, so that check would fall back + # to a `/` prefix. The archive write would then miss the + # `conversation_history` route and land in the default backend -- + # silently writing into the user's project tree. An `artifacts_root` + # passthrough on the guard would not help; the `isinstance` is what + # fails. + # + # This is a forward-looking constraint on the *constructor argument*, + # not a bug being fixed: the previous code also passed the real + # composite backend here and only swapped `_backend` for the guard + # afterwards, so the prefix was correct then too. The offload call sites + # apply the guard separately when writing (see `_guarded_backend`). + return create_summarization_middleware(model, self._summarization._backend) def _run_forced_compact(self, runtime: ToolRuntime) -> Command: """Synchronously compact without the SDK eligibility gate. @@ -694,32 +913,26 @@ def _run_forced_compact(self, runtime: ToolRuntime) -> Command: Returns: The compaction state update or an error tool message. """ - tool_call_id = runtime.tool_call_id or "" try: summarization = self._summarization_for_runtime(runtime) messages = runtime.state.get("messages", []) event = runtime.state.get("_summarization_event") effective = summarization._apply_event_to_messages(messages, event) - effective = _without_offload_seed(effective, tool_call_id) + effective = _without_offload_seed(effective, runtime.tool_call_id or "") cutoff = summarization._determine_cutoff_index(effective) if cutoff == 0: - return self._nothing_to_compact(tool_call_id) - + return self._nothing_to_compact(runtime.tool_call_id or "") to_summarize, _ = summarization._partition_messages(effective, cutoff) summary = summarization._create_summary(to_summarize) - backend = self._guarded_backend() - file_path = summarization._offload_to_backend(backend, to_summarize) - # The inherited `_build_compact_result` produces the same event and - # tool message as the SDK's gated path via model-independent helpers - # (string formatting + a staticmethod), so the runtime-selected - # summarizer is not needed to build it. Kept inside the `try` so a - # failure here still returns a ToolMessage rather than raising. + file_path = summarization._offload_to_backend( + self._guarded_backend(), to_summarize + ) return self._build_compact_result( runtime, to_summarize, summary, file_path, event, cutoff ) except Exception as exc: # tool errors must surface as ToolMessages logger.exception("forced compact_conversation failed") - return self._forced_compact_error(tool_call_id, exc) + return self._forced_compact_error(runtime.tool_call_id or "", exc) async def _arun_forced_compact(self, runtime: ToolRuntime) -> Command: """Asynchronously compact without the SDK eligibility gate. @@ -727,7 +940,6 @@ async def _arun_forced_compact(self, runtime: ToolRuntime) -> Command: Returns: The compaction state update or an error tool message. """ - tool_call_id = runtime.tool_call_id or "" try: summarization = await asyncio.to_thread( self._summarization_for_runtime, runtime @@ -735,23 +947,162 @@ async def _arun_forced_compact(self, runtime: ToolRuntime) -> Command: messages = runtime.state.get("messages", []) event = runtime.state.get("_summarization_event") effective = summarization._apply_event_to_messages(messages, event) - effective = _without_offload_seed(effective, tool_call_id) + effective = _without_offload_seed(effective, runtime.tool_call_id or "") cutoff = summarization._determine_cutoff_index(effective) if cutoff == 0: - return self._nothing_to_compact(tool_call_id) - + return self._nothing_to_compact(runtime.tool_call_id or "") to_summarize, _ = summarization._partition_messages(effective, cutoff) summary = await summarization._acreate_summary(to_summarize) - backend = self._guarded_backend() - file_path = await summarization._aoffload_to_backend(backend, to_summarize) - # See `_run_forced_compact` for why the inherited builder is reused - # and why it stays inside the `try`. + file_path = await summarization._aoffload_to_backend( + self._guarded_backend(), to_summarize + ) return self._build_compact_result( runtime, to_summarize, summary, file_path, event, cutoff ) except Exception as exc: # tool errors must surface as ToolMessages logger.exception("forced compact_conversation failed") - return self._forced_compact_error(tool_call_id, exc) + return self._forced_compact_error(runtime.tool_call_id or "", exc) + + async def arun_forced_compaction_update( + self, state: _OffloadState, runtime: _HasRunContext + ) -> dict[str, SummarizationEvent] | None: + """Run forced compaction as a server operation without a tool message. + + Unlike the tool paths, this raises on failure instead of returning a + `ToolMessage`: the operation graph has no tool node to carry one, so its + node converts the exception into a client-visible error. Any summarizer, + backend, or archive failure therefore propagates out of this method + rather than being folded into the return value. + + Args: + state: Checkpointed conversation and prior summarization event. + runtime: Run context carrier used to select the summarizer model. + + Returns: + The state update, or `None` when nothing can be compacted -- either + nothing is old enough to summarize, or the absolute cutoff would + not advance past the prior event. + + Raises: + ValueError: If the run observed no messages at all. On a real server + the run input *replaces* the `messages` channel, so this means + the caller replayed an empty list and has already truncated the + conversation; reporting "nothing to compact" would render that + as success. + """ + summarization = await asyncio.to_thread( + self._summarization_for_runtime, runtime + ) + messages = state.get("messages", []) + event = state.get("_summarization_event") + if not messages: + msg = ( + "Offload ran against an empty conversation. The run input " + "replaces the thread's messages, so this indicates the client " + "replayed an empty list rather than the checkpointed " + "conversation." + ) + raise ValueError(msg) + effective = summarization._apply_event_to_messages(messages, event) + cutoff = summarization._determine_cutoff_index(effective) + if cutoff == 0: + return None + # Resolved once and threaded into the update below: the SDK call is the + # relative-to-absolute conversion, and computing it twice would let the + # value checked here drift from the value committed. + state_cutoff = summarization._compute_state_cutoff(event, cutoff) + if state_cutoff <= _event_cutoff(event): + # Degenerate chained compaction: everything eligible is already + # behind the prior event's cutoff, so only the previous summary + # would be re-summarized. Committing would spend a model call to + # replace the in-context summary with a lossier summary-of-a-summary + # and drop the prior `file_path` from the event -- while the client, + # which keys its report on the *absolute* cutoff advancing, still + # reported "nothing to offload". Stop before the model call so the + # report and the state agree. + return None + to_summarize, _ = summarization._partition_messages(effective, cutoff) + summary = await summarization._acreate_summary(to_summarize) + file_path = await summarization._aoffload_to_backend( + self._guarded_backend(), to_summarize + ) + if file_path is None: + # `_aoffload_to_backend` catches every write failure and returns + # `None`, which also swallows `_ArchiveReadGuard`'s deliberate + # "refusing to overwrite existing history" `RuntimeError`. Its own + # log names neither the thread nor this call site, so record one + # here that does. + # + # Not raised: the compaction is still useful (the summary is + # in-context and the raw messages remain in the checkpoint), and the + # client reports the missing archive to the user as an error rather + # than a success. Escalating here would change that policy, not just + # its observability. + logger.error( + "/offload compacted %d messages but the archive write failed; " + "those messages are not recoverable from storage", + len(to_summarize), + ) + return self._forced_compaction_update( + summarization, summary, file_path, state_cutoff + ) + + @staticmethod + def _forced_compaction_update( + summarization: SummarizationMiddleware, + summary: str, + file_path: str | None, + state_cutoff: int, + ) -> dict[str, SummarizationEvent]: + """Build the state-only result used by the dedicated `/offload` graph. + + Annotated with the SDK's own `SummarizationEvent` rather than a loose + `dict[str, object]` so the hand-built payload is checked against the + shape the channel actually stores -- in particular that + `_build_new_messages_with_path(...)[0]` really is the `HumanMessage` + the event expects. + + Args: + summarization: SDK summarization middleware building the message. + summary: Generated summary text. + file_path: Archive path, or `None` when the write failed. + state_cutoff: **Absolute** cutoff index, already converted from the + relative one by `_compute_state_cutoff`. Taken pre-resolved + rather than converted here so the caller's no-advance check and + the committed value cannot disagree. + + Returns: + The summarization-event state update. + + Raises: + TypeError: If the summarizer's first message is not the + `HumanMessage` the event schema declares. + """ + summary_message = summarization._build_new_messages_with_path( + summary, file_path + )[0] + if not isinstance(summary_message, HumanMessage): + # `_build_new_messages_with_path` is annotated `list[AnyMessage]` + # but documents (and the SDK's own call site assumes, with a type + # suppression) that element 0 is the summary `HumanMessage`. Check + # rather than suppress: the node turns this into a visible + # "Compaction failed" instead of checkpointing an event whose + # `summary_message` violates its own schema. + msg = ( + "Summarizer returned a " + f"{type(summary_message).__name__} summary message; expected " + "HumanMessage." + ) + raise TypeError(msg) + return { + "_summarization_event": { + # Absolute, not relative: a second `/offload` on the same thread + # reads this back as its base. + "cutoff_index": state_cutoff, + "summary_message": summary_message, + "file_path": file_path, + } + } @staticmethod def _forced_compact_error(tool_call_id: str, exc: Exception) -> Command: @@ -812,3 +1163,248 @@ def _create_cli_compaction_middleware( sdk_middleware._summarization, system_prompt=sdk_middleware.system_prompt, ) + + +OFFLOAD_GRAPH_NODE = "force_compact" +"""Name of the `/offload` operation graph's only node. + +Shared with the client so `app._drive_offload_operation_graph` can recognize +this node's `updates` payload as positive evidence that the run reached the +node. A literal on each side would let a rename turn that check into a silent +"the stream produced nothing", which is exactly the condition it detects. +""" + +_OFFLOAD_CALL_NAMESPACE = uuid5(NAMESPACE_URL, "https://deepagents/offload/forced-call") +"""Namespace for deriving the `/offload` hook dispatch's forced tool-call id.""" + + +def _forced_offload_call_id() -> str: + """Return the tool-call id the `/offload` hook dispatch runs against. + + Two requirements pull in opposite directions, and both are load-bearing: + + *Stable across resumes.* `ServerHooksMiddleware` folds this id into its hook + `invocation_id`, and answering a hook interrupt re-executes the node **from + the top** -- LangGraph replays the task rather than resuming mid-coroutine. + A `uuid4()` minted here would therefore differ between the request and the + resume, and `parse_hook_resume_value` rejects a mismatched invocation id as + fatal ("the client answered a different request"). That made `/offload` fail + outright for anyone with a `PreCompact`/`PreToolUse` hook configured, and + made the client's whole fulfill/resume loop unreachable. + + *Distinct across runs.* The client memoizes fulfillments by + `(snapshot_id, invocation_id)` for the session, and the hook `prompt_id` + only rotates on user-prompt submit. A constant would make two `/offload`s + within one turn collide and replay the first run's decision -- including a + denial -- instead of re-running the user's hook. + + The task-scoped `checkpoint_ns` satisfies both: LangGraph reuses the task id + when it replays an interrupted task, and mints a fresh one for every run on + the thread (including a re-run after an abandoned or failed one). + + Returns: + An id stable across this run's resumes and distinct from every other + run's. + """ + try: + config = get_config() + except RuntimeError: + # No runnable context at all -- a direct call outside a graph run. + # Nothing can interrupt or resume such a call, so uniqueness is the only + # property left to preserve, and the `uuid4()` fallback is correct. + return f"offload-precompact-{uuid4()}" + configurable = config.get("configurable") + namespace = ( + configurable.get("checkpoint_ns") if isinstance(configurable, dict) else None + ) + if not namespace: + # A runnable context *without* a usable `checkpoint_ns` is a different + # situation entirely, and a silent fallback here is the failure mode + # this function exists to prevent: the id would differ between the + # request and the resume, `parse_hook_resume_value` would reject the + # mismatch as fatal, and `/offload` would die with "the client answered + # a different request" -- but only for users with hooks configured, and + # with nothing in the logs pointing here. Say so loudly. + logger.warning( + "Deriving the /offload hook call id inside a run but " + "`configurable.checkpoint_ns` is %r; falling back to a random id. " + "Configured PreCompact/PreToolUse hooks will fail to resume this " + "run. This usually means LangGraph moved or renamed the key.", + namespace, + ) + return f"offload-precompact-{uuid4()}" + return f"offload-precompact-{uuid5(_OFFLOAD_CALL_NAMESPACE, namespace)}" + + +def create_forced_compaction_graph( + middleware: CLICompactionMiddleware, + *, + hooks_middleware: ServerHooksMiddleware | None, +) -> CompiledStateGraph[Any, CLIContextSchema, Any, Any]: + """Create the dedicated server graph used by the `/offload` command. + + The run persists a summarization event and the run's cost drain. Unlike the + model-facing `compact_conversation` tool, it never writes an assistant tool + call or a tool result into the conversation: the slash command itself is the + explicit user authorization boundary. It does build a forced tool call for + the hook dispatch below, but that message exists only in memory and is never + checkpointed. + + Args: + middleware: Compaction implementation configured with the agent's + composite backend. + hooks_middleware: Server lifecycle middleware shared with the + interactive graph, or an explicit `None` to run with no hook gate. + It is invoked against an in-memory forced tool call so `PreCompact` + (and `PreToolUse`, which the same hook boundary dispatches for the + forced call) retains its normal authorization boundary without + persisting synthetic conversation messages. Required rather than + defaulted so skipping the gate cannot happen by omission. + + Only the *pre* events fire here. The node dispatches + `aafter_model`; it does not run a `ToolNode`, so + `PostToolUse`/`PostToolUseFailure` -- which the agent graph raises + for a model-initiated `compact_conversation` via + `ServerHooksMiddleware.awrap_tool_call` -- do not fire for + `/offload`. A matcher on `compact_conversation` therefore sees + automatic compaction but not the explicit command. This is a + deliberate consequence of `/offload` no longer executing a tool + (see THREAT_MODEL TB2), not an oversight. + + Returns: + A checkpointable graph that performs one forced compaction attempt. + """ + from langgraph.graph import END, START, StateGraph + + cost_tracking: CostTrackingMiddleware[CLIContextSchema] = CostTrackingMiddleware() + + async def force_compact( + state: _OffloadState, runtime: Runtime[CLIContextSchema] + ) -> dict[str, object]: + if hooks_middleware is not None: + try: + # Stable across this run's resumes, fresh for the next run. Both + # halves matter and neither is obvious -- see the helper. Inside + # the `try` so an unexpected failure in the derivation gets the + # message-preserving `RuntimeError` re-wrap below rather than + # reaching the user as "An internal error occurred". + forced_call_id = _forced_offload_call_id() + hook_update = await hooks_middleware.aafter_model( + cast( + "Any", + { + "messages": [ + AIMessage( + content="", + tool_calls=[ + { + "name": "compact_conversation", + "args": {"force": True}, + "id": forced_call_id, + } + ], + ) + ] + }, + ), + cast("Runtime[Any]", runtime), + ) + except GraphBubbleUp: + # Hook approval requests pause the operation graph through this + # control-flow exception. The caller streams and fulfills it. + raise + except Exception as exc: + logger.exception("/offload hook dispatch failed") + # Same `RuntimeError` re-raise rationale as the compaction + # failure below, but worded so a hook-layer failure is not read + # as a compaction failure. Nothing has been written yet here. + msg = ( + f"Offload hooks failed: {type(exc).__name__}: {exc}. " + "Your conversation is unchanged." + ) + raise RuntimeError(msg) from exc + # Read the key from the producer rather than re-spelling it: a + # mismatch here silently yields `{}`, which reads as "no outcome" + # and would compact straight through a hook denial. + outcomes = hook_update.get(_PRE_TOOL_STATE_KEY, {}) + outcome = outcomes.get(forced_call_id, {}) + if outcome.get("behavior") == "deny": + # Returning `{}` here would be indistinguishable from "nothing + # old enough to compact", so the client would report a hook + # veto as "the conversation is already compact". Raise instead + # so the reason reaches the user. Deliberately not re-wrapped by + # the handler below: this message is already user-facing. + reason = outcome.get("reason") or "Blocked by a compaction hook" + raise RuntimeError(str(reason)) + if outcome.get("context"): + # The agent's tool path injects an allowing hook's + # `additionalContext` into the tool result via + # `_append_message_text`. This graph has no tool result and no + # model turn to inject into, so the text is dropped. Log it so a + # hook author whose output silently vanishes has something to + # find. + logger.warning( + "Discarding PreToolUse additionalContext for the /offload " + "forced compact_conversation call; the operation graph has " + "no tool result to carry it" + ) + try: + update = await middleware.arun_forced_compaction_update(state, runtime) + except GraphBubbleUp: + # Same rationale as the hook dispatch above: LangGraph's control-flow + # exceptions (interrupt, retry) are not failures and must not be + # rewritten into a "Compaction failed" message that pauses nothing. + raise + except Exception as exc: + logger.exception("forced /offload compaction failed") + # Re-raise as `RuntimeError`: the server serializes exceptions for + # the client and preserves the message only for an allowlist of + # builtin types, replacing every other one with "An internal error + # occurred" -- which is what an `OSError` from the archive write or + # a provider SDK error would otherwise become. + # + # Deliberately no cost drain on this path: the drain returns an + # update the raise would discard, and it is destructive, so the + # summarizer's spend would be lost outright. Left undrained it is + # charged on the next turn's first step instead. + msg = ( + f"{COMPACTION_FAILURE_PREFIX}: {type(exc).__name__}: {exc}. " + "Your conversation is unchanged." + ) + raise RuntimeError(msg) from exc + # Summary generation invokes a model outside the normal agent loop, so + # drain it here rather than leaving this run's checkpoint incomplete. + # + # A drain failure must not propagate: `update` already reflects an + # archive section written to the backend, so raising here would discard + # it and tell the user their conversation is unchanged while leaving an + # orphaned section no `_summarization_event` references. Undrained spend + # is merely charged on the next turn, which is the same outcome as the + # failure path above. + # + # `after_agent`, not `aafter_agent`: `CostTrackingMiddleware` implements + # only the sync hook, so awaiting the async one would silently resolve + # to `AgentMiddleware`'s empty base method and drain nothing. Dispatched + # through a thread because the pricing lookup reads from disk and this + # runs on the blockbuster-guarded server loop. + try: + cost_update = await asyncio.to_thread( + cost_tracking.after_agent, state, runtime + ) + except Exception: + logger.exception( + "Failed to drain summary cost after /offload; the spend is " + "charged on the next turn instead" + ) + cost_update = None + return {**(update or {}), **(cost_update or {})} + + graph = StateGraph( + cast("Any", _OffloadState), + context_schema=CLIContextSchema, + input_schema=cast("Any", _OffloadInput), + ) + graph.add_node(OFFLOAD_GRAPH_NODE, force_compact) + graph.add_edge(START, OFFLOAD_GRAPH_NODE) + graph.add_edge(OFFLOAD_GRAPH_NODE, END) + return graph.compile() diff --git a/libs/code/deepagents_code/server_graph.py b/libs/code/deepagents_code/server_graph.py index b133cda91e..fd796efcad 100644 --- a/libs/code/deepagents_code/server_graph.py +++ b/libs/code/deepagents_code/server_graph.py @@ -15,7 +15,7 @@ import atexit import logging import sys -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NamedTuple from deepagents_code._server_config import ServerConfig from deepagents_code._startup_error import ( @@ -27,6 +27,8 @@ if TYPE_CHECKING: from collections.abc import Awaitable, Callable + from deepagents.backends.composite import CompositeBackend + logger = logging.getLogger(__name__) _sandbox_cm: Any = None @@ -185,15 +187,34 @@ def _mcp_tool_is_explicitly_read_only(tool: Any) -> bool: # noqa: ANN401 return mcp_tool_is_coherently_read_only(tool) -async def _make_graph() -> Any: # noqa: ANN401 - """Create the agent graph from environment-based configuration. +class ServerRuntime(NamedTuple): + """The one-per-process result of building this server's agent. + + A named tuple rather than a bare pair so the two slots are addressed by name: + both are structurally opaque to the type checker, and a positional + transposition would hand LangGraph the backend as its compiled graph while + making `offload_resources_from` return `None` — surfacing as "/offload has + no implementation", which points at the wrong subsystem entirely. + """ + + agent: Any + """Compiled LangGraph agent graph served as `agent`.""" + + backend: CompositeBackend + """Composite backend the agent was built with, carrying the offload + resources the `offload` graph resolves through + `offload_middleware.offload_resources_from`.""" + + +async def _make_graphs() -> ServerRuntime: + """Create the agent graph and the backend carrying its shared resources. Reads `DEEPAGENTS_CODE_SERVER_*` env vars via `ServerConfig.from_env()` (the inverse of `ServerConfig.to_env()` used by the app process), resolves a model, assembles tools, and compiles the agent graph. Returns: - Compiled LangGraph agent graph. + The normal agent graph and its configured composite backend. """ config = ServerConfig.from_env() @@ -319,7 +340,7 @@ def _cleanup_sandbox() -> None: ) sys.exit(1) - def _create_cli_agent_sync() -> Any: # noqa: ANN401 + def _create_cli_graphs_sync() -> ServerRuntime: async_subagents = load_async_subagents() or None auto_mode_enabled = config.interactive and sandbox_backend is None @@ -332,7 +353,7 @@ def _create_cli_agent_sync() -> Any: # noqa: ANN401 if config.enable_interpreter: settings.enable_interpreter = True - agent, _composite_backend = create_cli_agent( + agent, composite_backend = create_cli_agent( model=result.model, assistant_id=config.assistant_id, tools=tools, @@ -363,60 +384,133 @@ def _create_cli_agent_sync() -> Any: # noqa: ANN401 goal_criteria_tools=read_only_context_tools, rubric_grader_tools=read_only_context_tools, ) - return agent + return ServerRuntime(agent=agent, backend=composite_backend) + + return await asyncio.to_thread(_create_cli_graphs_sync) + + +class GraphFactories(NamedTuple): + """The factory pair `langgraph.json` resolves its two graph refs against. + + Named for the same reason as `ServerRuntime`, one level up and with a worse + failure: both slots are zero-arg async callables returning `Any`, so a + positional transposition type-checks, registers the offload graph as `agent` + and the agent graph as `offload`, and — because `generate_langgraph_json` + derives both refs from one module — starts the server cleanly. The first + user message would then run a graph with no model node. + """ + + agent: Callable[[], Awaitable[Any]] + """Factory for the interactive graph served as `agent`.""" - return await asyncio.to_thread(_create_cli_agent_sync) + offload: Callable[[], Awaitable[Any]] + """Factory for the `/offload` operation graph served as `offload`.""" -def _build_graph_factory( - builder: Callable[[], Awaitable[Any]] | None = None, -) -> Callable[[], Awaitable[Any]]: - """Build the cached async graph factory exposed to `langgraph dev`. +def _build_graph_factories( + builder: Callable[[], Awaitable[ServerRuntime]] | None = None, +) -> GraphFactories: + """Build paired factories that share one server-resource initialization. - The returned coroutine function is what `langgraph.json` references. It keeps - its cache and lock in this closure rather than in module-level globals, so - importing the module (e.g. for import-only checks) introduces no shared + The LangGraph server resolves named graphs independently. Keeping the cache + here ensures the `agent` and `offload` graphs use the same sandbox, + composite backend, and MCP sessions rather than constructing a second + server runtime for the slash command. + + The cache and lock live in this closure rather than in module-level globals, + so importing the module (e.g. for import-only checks) introduces no shared mutable state. Args: - builder: Optional alternate graph builder. + builder: Optional alternate builder for the shared server resources. Returns: - A zero-arg async factory that builds the graph once and returns the - cached instance on every subsequent call. + Named factories for the interactive agent and `/offload` operation + graphs. """ - missing = object() - graph: Any = missing + # `None` is a sound sentinel for both: the builder always returns a + # `ServerRuntime`, and `create_forced_compaction_graph` always a graph. + runtime: ServerRuntime | None = None + offload_graph: Any = None lock = asyncio.Lock() - async def make_graph() -> Any: # noqa: ANN401 - """Create (or return the cached) agent graph for `langgraph dev`. + async def shared_resources() -> ServerRuntime: + """Build (or return the cached) agent graph and composite backend. - LangGraph loads this async factory from the generated `langgraph.json` - and invokes it lazily on its event loop — and again on every run. The - built graph is cached for the process lifetime so MCP discovery, sandbox + LangGraph loads the factories below from the generated `langgraph.json` + and invokes them lazily on its event loop — and again on every run. The + result is cached for the process lifetime so MCP discovery, sandbox creation, and `atexit` registration each happen exactly once; re-running them per request would re-discover MCP servers, leak sandbox sessions, - and stack duplicate `atexit` handlers. Any construction failure is - converted into a startup-error marker (scraped by the parent app - process) before exiting. + and stack duplicate `atexit` handlers. + + Any construction failure is converted into a startup-error marker + (scraped by the parent app process) before **exiting the process**. That + `sys.exit(1)` is the most surprising thing this function does, so it is + stated here rather than left to the `except` block: there is no usable + server without a graph, and failing loudly at startup beats serving + every request with the same error. + + Returns: + The agent graph and the composite backend it was built with. + """ + nonlocal runtime + if runtime is None: + async with lock: + if runtime is None: + try: + runtime = await (builder or _make_graphs)() + except Exception as exc: # noqa: BLE001 # top-level barrier: any construction failure must surface to the parent as a marker + emit_startup_failure(exc) + sys.exit(1) + return runtime + + async def make_graph() -> Any: # noqa: ANN401 + """Return the normal interactive agent graph. Returns: Compiled LangGraph agent graph. """ - nonlocal graph - if graph is not missing: - return graph - async with lock: - if graph is missing: - try: - graph = await (builder or _make_graph)() - except Exception as exc: # noqa: BLE001 # top-level barrier: any construction failure must surface to the parent as a marker - emit_startup_failure(exc) - sys.exit(1) - return graph - - return make_graph - - -make_graph = _build_graph_factory() + return (await shared_resources()).agent + + async def make_offload_graph() -> Any: # noqa: ANN401 + """Return the explicit server-side `/offload` operation graph. + + Returns: + Compiled graph performing one forced compaction attempt. + + Raises: + RuntimeError: If the shared agent backend published no compaction + middleware, which would leave `/offload` with no implementation. + """ + nonlocal offload_graph + # Awaited to completion before taking `lock`, so the two critical + # sections are sequential rather than nested and cannot deadlock. + backend = (await shared_resources()).backend + if offload_graph is None: + async with lock: + if offload_graph is None: + from deepagents_code.offload_middleware import ( + create_forced_compaction_graph, + offload_resources_from, + ) + + offload_resources = offload_resources_from(backend) + if offload_resources is None: + msg = ( + "Agent backend did not publish its offload " + "middleware; /offload has no implementation." + ) + raise RuntimeError(msg) + offload_graph = create_forced_compaction_graph( + offload_resources.compaction, + hooks_middleware=offload_resources.hooks, + ) + return offload_graph + + return GraphFactories(agent=make_graph, offload=make_offload_graph) + + +_factories = _build_graph_factories() +make_graph = _factories.agent +make_offload_graph = _factories.offload diff --git a/libs/code/tests/integration_tests/test_offload_server_side.py b/libs/code/tests/integration_tests/test_offload_server_side.py index a44eb4b67f..63664e1eed 100644 --- a/libs/code/tests/integration_tests/test_offload_server_side.py +++ b/libs/code/tests/integration_tests/test_offload_server_side.py @@ -1,11 +1,11 @@ """Integration coverage for the server-side `/offload` path. -`/offload` drives the agent's own `compact_conversation` tool (with -`force=True`) server-side, so the offloaded archive lands in the agent's -composite backend and is readable via `read_file` in every run mode — not in a -client-local directory the server can never read. These tests construct the app -the PRODUCTION way (`backend=None`) and prove the archive is readable *through -the agent*. +For a server-backed agent `/offload` runs the dedicated `offload` operation +graph, which compacts without a model node or a synthetic tool call. Either way +the offloaded archive lands in the agent's composite backend and is readable via +`read_file` in every run mode — not in a client-local directory the server can +never read. These tests construct the app the PRODUCTION way (`backend=None`) +and prove the archive is readable *through the agent*. """ from __future__ import annotations @@ -20,7 +20,16 @@ def _write_model_config(home_dir: Path) -> None: - """Write a temp config that points the server subprocess at the test model.""" + """Write a temp config that points the server subprocess at the test model. + + The fake model's 8k-token default profile overflows once the system + prompt plus two seeded long turns cross the 85% auto-compaction trigger, + so auto-compaction fires during seeding and leaves `/offload` nothing + genuine to compact. Widening the window past the seeded size keeps the + thread uncompacted until `/offload`, while the fraction-based retention + window (~800 tokens) stays smaller than the seeded ~4.4k, so the forced + compaction still has real work to do. + """ config_dir = home_dir / ".deepagents" config_dir.mkdir(parents=True, exist_ok=True) (config_dir / "config.toml").write_text( @@ -28,6 +37,9 @@ def _write_model_config(home_dir: Path) -> None: [models.providers.itest] class_path = "deepagents_code._testing_models:DeterministicIntegrationChatModel" models = ["fake"] + +[models.providers.itest.profile] +max_input_tokens = 32000 """.strip() + "\n" ) @@ -44,15 +56,19 @@ def _build_long_prompt(turn: int) -> str: async def _run_turn(agent, *, thread_id: str, assistant_id: str, prompt: str) -> None: """Execute one real remote agent turn and drain the stream to completion.""" - from deepagents_code.config import build_stream_config + from deepagents_code.config import build_stream_config, settings config = build_stream_config(thread_id, assistant_id) stream_input = {"messages": [{"role": "user", "content": prompt}]} + # Send the resolved context limit so the server's compaction/summarization + # layers see the same window the model profile was widened to; without it + # the server falls back to its own default and auto-compaction fires early. async for _chunk in agent.astream( stream_input, stream_mode=["messages", "updates"], subgraphs=True, config=config, + context={"model_context_limit": settings.model_context_limit}, durability="exit", ): pass @@ -88,14 +104,20 @@ async def _read_file_through_agent(agent, *, thread_id: str, file_path: str) -> {"name": "read_file", "args": {"file_path": file_path}, "id": tool_call_id} ], ) - await agent.aensure_thread(config) - await agent.aupdate_state(config, {"messages": [seed]}, as_node="model") + # `/offload` restores the thread's main-graph association before returning, + # so this seeds the read straight through the interactive `agent` graph's + # model node. Seeding against that client (rather than the app's default + # graph) also keeps the test pinned to the interactive graph `/offload` + # shares its checkpoint with. + agent_graph = agent.for_graph("agent") + await agent_graph.aensure_thread(config) + await agent_graph.aupdate_state(config, {"messages": [seed]}, as_node="model") interrupt_ids: list[str] = [] tool_contents: list[str] = [] async def _drain(stream_input) -> None: - async for chunk in agent.astream( + async for chunk in agent_graph.astream( stream_input, stream_mode=["messages", "updates"], subgraphs=True, @@ -136,6 +158,7 @@ async def test_offload_runs_server_side_and_is_agent_readable( enough content, runs `/offload`, and asserts: - no `ErrorMessage` and an "Offloaded " success message, + - no HITL interrupt is surfaced, the operation graph having no tool node, - a persisted `_summarization_event` with `cutoff > 0` and `file_path == /conversation_history/.md`, - the archive is readable THROUGH THE AGENT (via its own `read_file` tool), @@ -189,6 +212,19 @@ async def test_offload_runs_server_side_and_is_agent_readable( config = {"configurable": {"thread_id": thread_id}} + # Captured before the run so the replay can be checked against it. + # The `/offload` run input is *authoritative* for the `messages` + # channel against a real server -- it replaces the conversation + # rather than merging into it (streaming `{"messages": []}` here + # empties the thread outright). No unit test can observe that: an + # in-process checkpointer honors the `add_messages` reducer and + # leaves the checkpointed list intact either way. + before_state = await agent.aget_state(config) + messages_before = list( + (getattr(before_state, "values", None) or {}).get("messages", []) + ) + assert messages_before + # Production construction: no client-owned backend. app = DeepAgentsApp( agent=agent, # ty: ignore @@ -198,6 +234,31 @@ async def test_offload_runs_server_side_and_is_agent_readable( thread_id=thread_id, ) + offload_interrupts: list[object] = [] + recorded_chunks = 0 + plain_for_graph = agent.for_graph + + def _recording_for_graph(graph_id: str): # noqa: ANN202 + """Instrument the `offload` client `/offload` actually streams.""" + offload_client = plain_for_graph(graph_id) + plain_astream = offload_client.astream + + async def _recording_astream(*args, **kwargs): # noqa: ANN002, ANN003, ANN202 + """Record every interrupt the server surfaces to the client.""" + nonlocal recorded_chunks + async for chunk in plain_astream(*args, **kwargs): + if isinstance(chunk, tuple) and len(chunk) == 3: + recorded_chunks += 1 + _ns, mode, data = chunk + if mode == "updates" and isinstance(data, dict): + offload_interrupts.extend( + data.get("__interrupt__") or [] + ) + yield chunk + + offload_client.astream = _recording_astream # ty: ignore + return offload_client + async with app.run_test() as pilot: for _ in range(120): await pilot.pause(0.1) @@ -206,15 +267,28 @@ async def test_offload_runs_server_side_and_is_agent_readable( assert app._message_store.total_count > 0 - await app._handle_offload() - - for _ in range(120): - await pilot.pause(0.1) - if any( - "Offloaded " in str(widget._content) - for widget in app.query(AppMessage) - ): - break + agent.for_graph = _recording_for_graph # ty: ignore + try: + await app._handle_offload() + + for _ in range(120): + await pilot.pause(0.1) + if any( + "Offloaded " in str(widget._content) + for widget in app.query(AppMessage) + ): + break + finally: + agent.for_graph = plain_for_graph # ty: ignore + + # The operation graph has no HITL middleware and manufactures no + # tool call, so the slash command is the whole authorization + # boundary: there is nothing left to approve in any approval + # mode (this app runs the default Manual mode). + assert offload_interrupts == [] + # Positive control: the recorder must have seen chunks, so the + # empty-interrupt assertion above cannot pass vacuously. + assert recorded_chunks > 0 app_messages = [ str(widget._content) for widget in app.query(AppMessage) @@ -230,6 +304,18 @@ async def test_offload_runs_server_side_and_is_agent_readable( # The summarization event must be visible through server state. state = await agent.aget_state(config) values = getattr(state, "values", None) or {} + + # `/offload` frees context by advancing the summarization cutoff, not + # by deleting messages: the raw conversation stays in the checkpoint + # so `/context` and resume still see it. Because the replay replaces + # this channel, a stale or empty input would silently truncate it + # here and still report success -- so assert identity, not count. + messages_after = values.get("messages", []) + assert len(messages_after) == len(messages_before) + assert [getattr(m, "id", None) for m in messages_after] == [ + getattr(m, "id", None) for m in messages_before + ] + summarization_event = values.get("_summarization_event") assert summarization_event is not None cutoff = _event_field(summarization_event, "cutoff_index") diff --git a/libs/code/tests/unit_tests/test_agent.py b/libs/code/tests/unit_tests/test_agent.py index 4eba67c07c..953f461631 100644 --- a/libs/code/tests/unit_tests/test_agent.py +++ b/libs/code/tests/unit_tests/test_agent.py @@ -152,6 +152,41 @@ def test_add_interrupt_on_attaches_auto_approve_predicate() -> None: assert config.get("when") is _should_interrupt_tool_call +def test_offload_resources_are_published_on_the_backend(tmp_path: Path) -> None: + """The `/offload` graph resolves its middleware off the returned backend. + + `create_cli_agent` returns only `(agent, backend)`, so the separately-resolved + server `offload` graph reads its compaction and hook middleware from an + attribute on the backend. Nothing else asserts this wiring: if the + `attach_offload_resources` call were dropped or moved above the middleware it + publishes, `/offload` would break for every server-backed agent and only the + slow subprocess integration test would notice. + """ + from deepagents_code.hooks.server_middleware import ServerHooksMiddleware + from deepagents_code.offload_middleware import ( + CLICompactionMiddleware, + offload_resources_from, + ) + + _agent, backend = create_cli_agent( + model=_make_fake_chat_model(), + assistant_id="test-agent", + enable_memory=False, + enable_skills=False, + enable_shell=False, + system_prompt="test prompt", + cwd=tmp_path, + ) + + resources = offload_resources_from(backend) + + assert resources is not None + assert isinstance(resources.compaction, CLICompactionMiddleware) + # Non-optional: an absent instance would silently skip the `PreCompact` / + # `PreToolUse` gate, the only hook boundary `/offload` still crosses. + assert isinstance(resources.hooks, ServerHooksMiddleware) + + def test_local_conversation_history_route_is_persistent(tmp_path: Path) -> None: """Local archives use the stable user data directory across server restarts.""" history_root = tmp_path / ".deepagents" diff --git a/libs/code/tests/unit_tests/test_auto_mode.py b/libs/code/tests/unit_tests/test_auto_mode.py index de6194e80e..7c02598c30 100644 --- a/libs/code/tests/unit_tests/test_auto_mode.py +++ b/libs/code/tests/unit_tests/test_auto_mode.py @@ -826,6 +826,220 @@ async def test_trusted_compaction_is_deterministically_allowed_without_human_rev assert update["messages"] == [ai_message] +def _offload_seed( + tool_call_id: str, + *, + args: dict[str, object] | None = None, + extra_calls: list[ToolCall] | None = None, +) -> AIMessage: + """Build the synthetic assistant message `/offload` injects into state.""" + from deepagents_code.offload_middleware import _offload_seed_message_id + + seed_call: ToolCall = { + "name": "compact_conversation", + "args": {"force": True} if args is None else args, + "id": tool_call_id, + "type": "tool_call", + } + return AIMessage( + content="", + id=_offload_seed_message_id(tool_call_id), + tool_calls=[seed_call, *(extra_calls or [])], + ) + + +async def _after_model_without_plan( + middleware: AutoModeHITLMiddleware, + request: ModelRequest[Any], + ai_message: AIMessage, +) -> dict[str, Any] | None: + """Route a plan-less assistant message, as the `/offload` seed always is.""" + return await middleware.aafter_model( + cast("AgentState[Any]", {"messages": [ai_message]}), + request.runtime, + ) + + +def _offload_request( + tmp_path: Path, + compact_tool: BaseTool, + *, + tool_call_id: str | None = "seed-call", +) -> ModelRequest[Any]: + """Build a request whose run context carries the `/offload` trust signal.""" + request, _store, _key = _request( + tmp_path, + model=_FailIfClassifiedModel(), + tool_name="compact_conversation", + args={}, + tools=[compact_tool], + ) + if tool_call_id is not None: + request.runtime.context["offload_tool_call_id"] = tool_call_id # ty: ignore + return request + + +async def test_seeded_offload_compaction_is_reviewed_without_operation_graph( + tmp_path: Path, +) -> None: + """A synthetic compaction call is not an approval bypass.""" + compact_tool = _tool("compact_conversation") + middleware = _middleware(tmp_path, trusted_compaction_tool=compact_tool) + request = _offload_request(tmp_path, compact_tool) + ai_message = _offload_seed("seed-call") + + with patch( + "deepagents_code.auto_mode.interrupt", + return_value={"decisions": [{"type": "approve"}]}, + ) as review: + await _after_model_without_plan(middleware, request, ai_message) + + assert review.called + + +async def test_seeded_offload_batch_with_extra_gated_call_is_reviewed( + tmp_path: Path, +) -> None: + """An appended gated call keeps the whole batch on the human-review path.""" + compact_tool = _tool("compact_conversation") + middleware = _middleware(tmp_path, trusted_compaction_tool=compact_tool) + request = _offload_request(tmp_path, compact_tool) + ai_message = _offload_seed( + "seed-call", + extra_calls=[ + { + "name": "execute", + "args": {"command": "curl https://example.com"}, + "id": "model-call", + "type": "tool_call", + } + ], + ) + + with patch( + "deepagents_code.auto_mode.interrupt", + return_value={"decisions": [{"type": "approve"}, {"type": "approve"}]}, + ) as review: + await _after_model_without_plan(middleware, request, ai_message) + + reviewed = [ + action["name"] for action in review.call_args.args[0]["action_requests"] + ] + assert reviewed == ["compact_conversation", "execute"] + + +@pytest.mark.parametrize( + ("tool_call_id", "seed_call_id", "seeded_args", "trusted"), + [ + pytest.param( + "seed-call", "seed-call", {"force": True}, False, id="untrusted-tool" + ), + pytest.param(None, "seed-call", {"force": True}, True, id="no-offload-context"), + pytest.param( + "other-call", "seed-call", {"force": True}, True, id="message-id-mismatch" + ), + pytest.param( + "seed-call", "other-call", {"force": True}, True, id="tool-call-id-mismatch" + ), + pytest.param("seed-call", "seed-call", {}, True, id="unforced-args"), + ], +) +async def test_plan_less_forced_compaction_always_reaches_review( + tmp_path: Path, + tool_call_id: str | None, + seed_call_id: str, + seeded_args: dict[str, object], + trusted: bool, +) -> None: + """A plan-less forced `compact_conversation` always reaches human review. + + There is no seed-based approval bypass any more -- `/offload` runs the + server-side operation graph instead -- so `plan is None` routes to manual + review no matter what the run context claims. Every case here therefore + expects the same outcome by design; the parameters are not each individually + load-bearing. They are retained as a regression guard: each varies one of the + signals the removed bypass keyed on (trusted tool, offload context, seed + message ID, tool-call ID, `force` args), so reintroducing a bypass keyed on + any single one of them fails here. + """ + compact_tool = _tool("compact_conversation") + middleware = _middleware( + tmp_path, trusted_compaction_tool=compact_tool if trusted else None + ) + request = _offload_request(tmp_path, compact_tool, tool_call_id=tool_call_id) + ai_message = _offload_seed(seed_call_id, args=seeded_args) + + with patch( + "deepagents_code.auto_mode.interrupt", + return_value={"decisions": [{"type": "approve"}]}, + ) as review: + await _after_model_without_plan(middleware, request, ai_message) + + assert review.call_args.args[0]["action_requests"][0]["name"] == ( + "compact_conversation" + ) + + +async def test_seed_message_carrying_other_gated_tool_is_reviewed( + tmp_path: Path, +) -> None: + """A seed-ID'd message does not bypass review for a non-compaction tool.""" + compact_tool = _tool("compact_conversation") + middleware = _middleware(tmp_path, trusted_compaction_tool=compact_tool) + request = _offload_request(tmp_path, compact_tool) + ai_message = _offload_seed( + "seed-call", + extra_calls=[ + { + "name": "execute", + "args": {"command": "curl https://example.com", "force": True}, + "id": "seed-call", + "type": "tool_call", + } + ], + ) + + with patch( + "deepagents_code.auto_mode.interrupt", + return_value={"decisions": [{"type": "approve"}, {"type": "approve"}]}, + ) as review: + await _after_model_without_plan(middleware, request, ai_message) + + reviewed = [ + action["name"] for action in review.call_args.args[0]["action_requests"] + ] + assert reviewed == ["compact_conversation", "execute"] + + +async def test_model_generated_compaction_during_offload_is_reviewed( + tmp_path: Path, +) -> None: + """A later model message cannot reuse the authorized ID to skip review.""" + compact_tool = _tool("compact_conversation") + middleware = _middleware(tmp_path, trusted_compaction_tool=compact_tool) + request = _offload_request(tmp_path, compact_tool) + ai_message = AIMessage( + content="", + id="model-generated-message", + tool_calls=[ + { + "name": "compact_conversation", + "args": {"force": True}, + "id": "seed-call", + "type": "tool_call", + } + ], + ) + + with patch( + "deepagents_code.auto_mode.interrupt", + return_value={"decisions": [{"type": "approve"}]}, + ) as review: + await _after_model_without_plan(middleware, request, ai_message) + + assert review.called + + async def test_same_name_custom_compaction_tool_requires_classifier( tmp_path: Path, ) -> None: diff --git a/libs/code/tests/unit_tests/test_compact_tool.py b/libs/code/tests/unit_tests/test_compact_tool.py index 28b2d82c51..9b5b53115f 100644 --- a/libs/code/tests/unit_tests/test_compact_tool.py +++ b/libs/code/tests/unit_tests/test_compact_tool.py @@ -8,7 +8,7 @@ import warnings from types import MethodType, SimpleNamespace -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -219,6 +219,194 @@ async def test_force_bypasses_sdk_eligibility_gate(self) -> None: assert result.update is not None assert result.update["_summarization_event"]["cutoff_index"] == 2 + async def test_forced_compact_writes_through_the_archive_guard(self) -> None: + """The archive write is guarded even though the summarizer is not. + + `_ArchiveReadGuard` deliberately exposes only the read/write methods it + has to intercept, so it cannot stand in for the composite backend the + summarizer needs (see `test_runtime_model_builds_matching_summarizer`). + The two therefore have to be handed out separately, and this pins that + the write half still gets the guard. + """ + summarization = self._summarization() + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + runtime.state = {"messages": [HumanMessage("one"), HumanMessage("two")]} + runtime.tool_call_id = "tool-call" + + await middleware._arun_forced_compact(runtime) + + write_backend = summarization._aoffload_to_backend.await_args.args[0] + assert isinstance(write_backend, _ArchiveReadGuard) + assert write_backend._backend is summarization._backend + assert not hasattr(write_backend, "artifacts_root") + + async def test_operation_path_writes_through_the_archive_guard(self) -> None: + """The `/offload` operation graph's write path has the same invariant. + + The guard is applied per write site rather than by the backend's type, so + the operation graph's entry point does not inherit it from the tool paths + — it has to apply it itself, and nothing but a test says so. + """ + summarization = self._summarization() + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + + await middleware.arun_forced_compaction_update( + {"messages": [HumanMessage("one"), HumanMessage("two")]}, runtime + ) + + write_backend = summarization._aoffload_to_backend.await_args.args[0] + assert isinstance(write_backend, _ArchiveReadGuard) + assert write_backend._backend is summarization._backend + + async def test_operation_path_returns_an_absolute_cutoff(self) -> None: + """The committed event must carry the absolute cutoff, not the relative one. + + `_determine_cutoff_index` is relative to the *effective* conversation + (post-previous-summary), while the persisted `cutoff_index` indexes the + full message list — `_compute_state_cutoff` converts between them. The + two coincide on a thread's first `/offload`, so returning the relative + value passes every other test here and only corrupts the *second* + `/offload`, which reads this back as its base. + """ + summarization = self._summarization() + summarization._determine_cutoff_index.return_value = 2 + summarization._compute_state_cutoff.return_value = 9 + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + prior = {"cutoff_index": 7, "summary_message": None, "file_path": None} + + result = await middleware.arun_forced_compaction_update( + cast( + "Any", + { + "messages": [HumanMessage("one"), HumanMessage("two")], + "_summarization_event": prior, + }, + ), + runtime, + ) + + assert result is not None + event = result["_summarization_event"] + summarization._compute_state_cutoff.assert_called_once_with(prior, 2) + assert event["cutoff_index"] == 9 + assert event["file_path"] == "/conversation_history/thread.md" + assert isinstance(event["summary_message"], HumanMessage) + + async def test_operation_path_refuses_a_chained_no_advance_compaction( + self, + ) -> None: + """A compaction that would not advance the cutoff must not commit. + + The degenerate chained case: everything eligible already sits behind the + prior event, so the only thing left to summarize is the previous summary + itself. `_compute_state_cutoff` returns the prior absolute cutoff + unchanged, and the client — which keys its report on that value moving — + reports "nothing to offload". Committing anyway would spend a model + call, replace the in-context summary with a summary-of-a-summary, and + drop the prior archive's `file_path`, all while telling the user nothing + happened. Stop before the model call so the report and the state agree. + """ + summarization = self._summarization() + summarization._determine_cutoff_index.return_value = 1 + summarization._compute_state_cutoff.return_value = 7 + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + prior = { + "cutoff_index": 7, + "summary_message": None, + "file_path": "/conversation_history/thread.md", + } + + result = await middleware.arun_forced_compaction_update( + cast( + "Any", + { + "messages": [HumanMessage("summary"), HumanMessage("recent")], + "_summarization_event": prior, + }, + ), + runtime, + ) + + assert result is None + # Neither the billable step nor the archive write may happen. + summarization._acreate_summary.assert_not_awaited() + summarization._aoffload_to_backend.assert_not_awaited() + + async def test_operation_path_rejects_an_empty_conversation(self) -> None: + """An empty `messages` must raise rather than report a clean no-op. + + On a real server the run input *replaces* this channel, so the node + seeing no messages means the client already truncated the thread. + Returning `None` would render that wipe to the user as "your + conversation is already compact". + """ + summarization = self._summarization() + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + + with pytest.raises(ValueError, match="empty conversation"): + await middleware.arun_forced_compaction_update( + cast("Any", {"messages": [], "_summarization_event": None}), runtime + ) + + summarization._acreate_summary.assert_not_awaited() + + async def test_operation_path_logs_a_failed_archive_write( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A `None` archive path must leave a trace naming this call site. + + `_aoffload_to_backend` catches every write failure and returns `None` — + including `_ArchiveReadGuard`'s deliberate "refusing to overwrite + existing history" `RuntimeError`. The compaction still commits (the + client reports the missing archive to the user), but without this the + only server-side record is a warning inside the SDK that names neither + the thread nor `/offload`. + """ + summarization = self._summarization() + summarization._aoffload_to_backend = AsyncMock(return_value=None) + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + + with caplog.at_level("ERROR"): + result = await middleware.arun_forced_compaction_update( + {"messages": [HumanMessage("one"), HumanMessage("two")]}, runtime + ) + + assert result is not None + assert result["_summarization_event"]["file_path"] is None + assert "archive write failed" in caplog.text + + async def test_operation_path_returns_none_when_nothing_to_compact(self) -> None: + """A cutoff of 0 must be `None`, not an event pinning cutoff 0. + + The caller distinguishes "nothing old enough" from a real compaction by + this return value; an empty-but-present event would advance nothing while + still reading as success. + """ + summarization = self._summarization() + summarization._determine_cutoff_index = MagicMock(return_value=0) + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + + result = await middleware.arun_forced_compaction_update( + {"messages": [HumanMessage("one")]}, runtime + ) + + assert result is None + summarization._aoffload_to_backend.assert_not_awaited() + def test_runtime_model_builds_matching_summarizer(self) -> None: """A `/model` override selects the summarizer used by `/offload`.""" startup = self._summarization() @@ -249,8 +437,13 @@ def test_runtime_model_builds_matching_summarizer(self) -> None: extra_kwargs={"temperature": 0}, profile_overrides=None, ) - create_summarization.assert_called_once_with(active_model, startup._backend) - assert actual._backend._backend is startup._backend + create_summarization.assert_called_once() + assert create_summarization.call_args.args[0] is active_model + # The summarizer gets the composite backend itself, not the + # `_ArchiveReadGuard` wrapper: it reads `artifacts_root` to prefix the + # archive path, and the guard exposes no such attribute. The write path + # applies the guard separately (see `test_forced_compact_writes_guarded`). + assert create_summarization.call_args.args[1] is startup._backend def test_runtime_profile_overrides_and_context_limit_are_applied(self) -> None: """Server-side offload uses the CLI's effective model profile.""" @@ -285,7 +478,9 @@ def test_runtime_profile_overrides_and_context_limit_are_applied(self) -> None: profile_overrides={"max_input_tokens": 32_000}, ) assert active_model.profile["max_input_tokens"] == 24_000 - create_summarization.assert_called_once_with(active_model, startup._backend) + create_summarization.assert_called_once() + assert create_summarization.call_args.args[0] is active_model + assert create_summarization.call_args.args[1] is startup._backend async def test_force_noops_when_nothing_old_enough(self) -> None: """Forced compaction still no-ops at cutoff 0 (bypasses only the gate).""" diff --git a/libs/code/tests/unit_tests/test_offload.py b/libs/code/tests/unit_tests/test_offload.py index ad39f016d3..a7e867f8fe 100644 --- a/libs/code/tests/unit_tests/test_offload.py +++ b/libs/code/tests/unit_tests/test_offload.py @@ -5,13 +5,16 @@ import os import stat import tempfile +from collections.abc import Callable # noqa: TC003 from contextlib import nullcontext from pathlib import Path, PureWindowsPath -from typing import Any +from types import SimpleNamespace +from typing import Annotated, Any, TypedDict, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from deepagents.backends.utils import validate_path +from langgraph.graph.message import add_messages from deepagents_code import offload from deepagents_code._session_stats import format_token_count @@ -46,6 +49,20 @@ def _make_dict_messages(n: int) -> list[dict[str, Any]]: return messages +def _make_dict_message( + content: str, *, message_id: str | None = None +) -> dict[str, Any]: + """Create one serialized human-message payload with a stable id.""" + return { + "content": content, + "additional_kwargs": {}, + "response_metadata": {}, + "type": "human", + "name": None, + "id": message_id or f"msg-{content}", + } + + def _make_dict_summary_message() -> dict[str, Any]: """Create a serialized summary message payload from remote state.""" return { @@ -80,11 +97,30 @@ def _state_values( def _setup_server_offload_app(app: DeepAgentsApp) -> MagicMock: - """Configure a `DeepAgentsApp` for server-side offload unit tests. + """Configure a `DeepAgentsApp` as a server-backed agent for offload tests. + + The operation-graph path reads state via `_get_thread_state_values` and + drives the graph via `_drive_offload_operation_graph`; tests patch those + seams directly, so only the remote identity/flags are set here. The agent + is specced as a `RemoteAgent` so `_remote_agent()` narrows to it. + """ + from deepagents_code.client.remote_client import RemoteAgent + + agent = MagicMock(spec=RemoteAgent) + agent.aupdate_state = AsyncMock() + app._agent = agent + app._backend = None + app._lc_thread_id = "test-thread" + app._agent_running = False + return agent + - The server-side path reads state via `_get_thread_state_values` and drives - the tool via `_drive_server_side_compaction`; tests patch those seams - directly, so only the plain identity/flags are set here. +def _setup_local_offload_app(app: DeepAgentsApp) -> MagicMock: + """Configure a `DeepAgentsApp` as a local in-process agent for offload tests. + + A plain `MagicMock` agent is *not* a `RemoteAgent`, so `_remote_agent()` + returns `None` and `_handle_offload` takes the seeded in-process path + (`_drive_local_seeded_compaction`) instead of the operation graph. """ agent = MagicMock() agent.aupdate_state = AsyncMock() @@ -173,7 +209,7 @@ async def test_nothing_to_compact_noop(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -256,7 +292,7 @@ async def test_successful_offload_drives_server_tool(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ) as mock_drive, @@ -272,6 +308,125 @@ async def test_successful_offload_drives_server_tool(self) -> None: # Offloaded count is the new cutoff of six minus a prior cutoff of zero. assert any("Offloaded 6 older messages" in str(w._content) for w in msgs) + async def test_rebind_warning_accompanies_a_reported_success(self) -> None: + """A rebind failure is surfaced where the success is reported. + + The driver cannot mount it itself: it does not know whether the caller + is about to report success. Reporting the offload as finished while + staying silent about the mis-bound thread leaves a later `/goal` or + `/rubric` to fail with nothing connecting it to `/offload`. + """ + from deepagents_code.app import _OFFLOAD_REBIND_WARNING + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + + before = _state_values(_make_dict_messages(10)) + after = _state_values(_make_dict_messages(12), _summary_event(6)) + + async def drive(*_args: object, **_kwargs: object) -> None: # noqa: RUF029 + app._offload_rebind_failed = True + + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, after], + ), + patch.object(app, "_drive_offload_operation_graph", side_effect=drive), + ): + await app._handle_offload() + await pilot.pause() + + assert any( + "Offloaded 6 older messages" in str(w._content) + for w in app.query(AppMessage) + ) + assert any( + _OFFLOAD_REBIND_WARNING in str(w._content) + for w in app.query(ErrorMessage) + ) + + async def test_rebind_warning_is_withheld_from_a_reported_failure(self) -> None: + """A drain error must not be paired with "Offload finished, but...". + + Both `drain_error` breaks exit the stream loop normally, so a flag that + only tracked "the stream did not raise" would mount the rebind warning + alongside the failure — two messages asserting opposite outcomes. + """ + from deepagents_code.app import _OFFLOAD_REBIND_WARNING + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + + before = _state_values(_make_dict_messages(10)) + + async def drive(*_args: object, **_kwargs: object) -> str: # noqa: RUF029 + app._offload_rebind_failed = True + return "Offload could not complete: a configured hook kept ..." + + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before], + ), + patch.object(app, "_drive_offload_operation_graph", side_effect=drive), + ): + await app._handle_offload() + await pilot.pause() + + errors = [str(w._content) for w in app.query(ErrorMessage)] + + assert any("could not complete" in text for text in errors) + assert not any(_OFFLOAD_REBIND_WARNING in text for text in errors) + + async def test_unreadable_state_is_not_reported_as_already_compact(self) -> None: + """An empty state read must not render as a benign no-op. + + On the operation-graph path the state re-read is the only evidence of + the outcome — there is no `ToolMessage` to fall back on — and + `_get_thread_state_values` collapses a missing snapshot (a 404 after the + run rebound the thread, a server restart) to `{}`. Reporting that as + "already compact" tells the user nothing happened when the conversation + may well have been compacted. + """ + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + + before = _state_values(_make_dict_messages(10)) + + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, {}], + ), + patch.object( + app, + "_drive_offload_operation_graph", + new_callable=AsyncMock, + return_value=None, + ), + ): + await app._handle_offload() + await pilot.pause() + + errors = [str(w._content) for w in app.query(ErrorMessage)] + infos = [str(w._content) for w in app.query(AppMessage)] + + assert any("could not be confirmed" in text for text in errors) + assert not any("already compact" in text for text in infos) + async def test_committed_offload_survives_stream_failure(self) -> None: """A checkpointed tool update wins over a later stream failure.""" app = DeepAgentsApp() @@ -291,7 +446,7 @@ async def test_committed_offload_survives_stream_failure(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, side_effect=RuntimeError("stream unavailable"), ), @@ -327,7 +482,7 @@ async def test_offload_shows_feedback_message(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -378,7 +533,7 @@ async def test_offload_updates_context_tokens(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -407,7 +562,7 @@ async def test_no_ui_clear_reload(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -429,11 +584,23 @@ class TestOffloadEdgeCases: """Test edge cases in the offload logic.""" async def test_noop_does_not_report_offloaded(self) -> None: - """A no-op restores history and shows the no-op message, not success.""" + """A no-op restores history and shows the no-op message, not success. + + The local seeded driver commits its synthetic seed, tool result, and + trailing turn on a no-op, so `_handle_offload` removes those artifacts + via `aupdate_state`. + """ app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - agent = _setup_server_offload_app(app) + # A plain (non-remote) mock agent drives the seeded in-process + # path, whose no-op branch restores state via `aupdate_state`. + agent = MagicMock() + agent.aupdate_state = AsyncMock() + app._agent = agent + app._backend = None + app._lc_thread_id = "test-thread" + app._agent_running = False # Prior event present; after-state cutoff unchanged -> nothing moved. event = _summary_event(6) @@ -476,7 +643,7 @@ async def test_noop_does_not_report_offloaded(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_local_seeded_compaction", new_callable=AsyncMock, return_value=None, ), @@ -489,13 +656,46 @@ async def test_noop_does_not_report_offloaded(self) -> None: "the conversation is already compact" in str(w._content) for w in msgs ) assert not any("Offloaded " in str(w._content) for w in msgs) - agent.aupdate_state.assert_awaited_once() - update = agent.aupdate_state.call_args.args[1] - assert [message.id for message in update["messages"]] == [ - "offload-seed-test", - "offload-result-test", - "offload-trailing-test", - ] + # The seeded no-op restores the pre-run conversation by removing + # the committed artifacts. + agent.aupdate_state.assert_awaited() + + async def test_noop_operation_graph_writes_nothing(self) -> None: + """The operation graph commits no synthetic artifacts on a no-op.""" + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + agent = _setup_server_offload_app(app) + + # Prior event present; after-state cutoff unchanged -> nothing moved. + event = _summary_event(6) + messages = _make_dict_messages(8) + before = _state_values(messages, event) + after = _state_values(messages, event) + + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, after], + ), + patch.object( + app, + "_drive_offload_operation_graph", + new_callable=AsyncMock, + return_value=None, + ), + ): + await app._handle_offload() + await pilot.pause() + + msgs = app.query(AppMessage) + assert any( + "the conversation is already compact" in str(w._content) for w in msgs + ) + assert not any("Offloaded " in str(w._content) for w in msgs) + agent.aupdate_state.assert_not_awaited() async def test_cutoff_one_offloads_single_message(self) -> None: """A cutoff of 1 reports a single offloaded message.""" @@ -516,7 +716,7 @@ async def test_cutoff_one_offloads_single_message(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -555,7 +755,7 @@ async def test_reoffload_uses_absolute_cutoff_delta(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -593,7 +793,7 @@ async def test_reoffload_noop_restores_prior_summary(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -601,13 +801,7 @@ async def test_reoffload_noop_restores_prior_summary(self) -> None: await app._handle_offload() await pilot.pause() - agent.aupdate_state.assert_awaited_once() - update = agent.aupdate_state.call_args.args[1] - assert update["_summarization_event"] is prior_event - assert [message.id for message in update["messages"]] == [ - "offload-seed", - "offload-result", - ] + agent.aupdate_state.assert_not_awaited() assert any( "Nothing to offload" in str(widget._content) for widget in app.query(AppMessage) @@ -643,7 +837,7 @@ def capture_running(_config: object, _seed_id: object = None) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, side_effect=capture_running, ), @@ -676,7 +870,7 @@ async def test_agent_running_reset_after_failure(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, side_effect=RuntimeError("stream down"), ), @@ -718,7 +912,7 @@ async def test_missing_archive_path_warns_about_unrecoverable_history( ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -760,7 +954,7 @@ async def test_tool_reported_compaction_failure_shows_error(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=tool_error, ), @@ -802,7 +996,7 @@ async def test_stale_compaction_failure_is_not_reported(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -848,7 +1042,7 @@ async def test_current_durable_compaction_failure_is_reported(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -861,18 +1055,110 @@ async def test_current_durable_compaction_failure_is_reported(self) -> None: for widget in app.query(ErrorMessage) ) - async def test_failed_run_removes_dangling_seed(self) -> None: - """A raising run cleans up the committed seed before surfacing failure. + async def test_failed_operation_graph_run_has_no_seed_to_clean_up(self) -> None: + """A failed operation-graph run needs no seed cleanup. + + The operation graph commits no synthetic tool call, so there is nothing + that could wedge the next turn with a dangling `tool_use` and the cleanup + must not run. The failure is still surfaced to the user. The seeded + driver's counterpart is + `test_failed_seeded_run_removes_dangling_seed`. + """ + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + + before = _state_values(_make_dict_messages(6)) + reconciled = _state_values(_make_dict_messages(6)) # cutoff unchanged + + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, reconciled], + ), + patch.object( + app, + "_drive_offload_operation_graph", + new_callable=AsyncMock, + side_effect=RuntimeError("stream boom"), + ), + patch.object( + app, + "_remove_unanswered_offload_seed", + new_callable=AsyncMock, + ) as cleanup, + ): + await app._handle_offload() + await pilot.pause() + + cleanup.assert_not_awaited() + assert any( + "Offload failed" in str(widget._content) + for widget in app.query(ErrorMessage) + ) + + async def test_operation_graph_double_failure_surfaces_one_error(self) -> None: + """Stream failure plus a failed reconcile still reports exactly once. - When the drive raises and the committed cutoff has not advanced, the - seeded (and now unanswered) tool call must be removed so it does not - wedge the next turn; the failure is still surfaced to the user. + The seeded driver additionally warns that the thread may be inconsistent + when it cannot confirm seed removal (see + `test_seeded_double_failure_warns_thread_may_be_inconsistent`). The + operation graph has no seed, so no cleanup runs and no wedge warning is + appropriate -- the user should see the "Offload failed" error alone + rather than a second, inapplicable warning. """ app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() _setup_server_offload_app(app) + before = _state_values(_make_dict_messages(6)) + + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, RuntimeError("reconcile read boom")], + ), + patch.object( + app, + "_drive_offload_operation_graph", + new_callable=AsyncMock, + side_effect=RuntimeError("stream boom"), + ), + patch.object( + app, + "_remove_unanswered_offload_seed", + new_callable=AsyncMock, + return_value=False, + ) as cleanup, + ): + await app._handle_offload() + await pilot.pause() + + cleanup.assert_not_awaited() + error_text = " ".join( + str(widget._content) for widget in app.query(ErrorMessage) + ) + assert "Offload failed" in error_text + assert "inconsistent state" not in error_text + + async def test_failed_seeded_run_removes_dangling_seed(self) -> None: + """A failed local seeded run must not leave an unanswered tool call. + + The seeded driver commits a synthetic assistant `tool_use` before the + tool runs, so a run that fails without compacting has to remove it or the + model API rejects the next turn. + """ + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_local_offload_app(app) + before = _state_values(_make_dict_messages(6)) reconciled = _state_values(_make_dict_messages(6)) # cutoff unchanged @@ -885,7 +1171,7 @@ async def test_failed_run_removes_dangling_seed(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_local_seeded_compaction", new_callable=AsyncMock, side_effect=RuntimeError("stream boom"), ), @@ -893,6 +1179,7 @@ async def test_failed_run_removes_dangling_seed(self) -> None: app, "_remove_unanswered_offload_seed", new_callable=AsyncMock, + return_value=True, ) as cleanup, ): await app._handle_offload() @@ -904,19 +1191,18 @@ async def test_failed_run_removes_dangling_seed(self) -> None: for widget in app.query(ErrorMessage) ) - async def test_double_failure_warns_thread_may_be_inconsistent(self) -> None: - """Stream failure + failed reconcile + failed cleanup warns the user. + async def test_seeded_double_failure_warns_thread_may_be_inconsistent(self) -> None: + """Unconfirmed seed removal warns the user the thread may be wedged. When the drive raises, the reconcile state-read also fails, and the - best-effort seed cleanup cannot confirm removal (returns False), the - user is warned the thread may be inconsistent -- in addition to the - surfaced "Offload failed" error -- so a later cryptic `tool_use` - rejection is not their only signal. + best-effort seed cleanup cannot confirm removal (returns `False`), the + user is warned -- in addition to the surfaced "Offload failed" error -- + so a later cryptic `tool_use` rejection is not their only signal. """ app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) + _setup_local_offload_app(app) before = _state_values(_make_dict_messages(6)) @@ -929,7 +1215,7 @@ async def test_double_failure_warns_thread_may_be_inconsistent(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_local_seeded_compaction", new_callable=AsyncMock, side_effect=RuntimeError("stream boom"), ), @@ -947,8 +1233,8 @@ async def test_double_failure_warns_thread_may_be_inconsistent(self) -> None: error_text = " ".join( str(widget._content) for widget in app.query(ErrorMessage) ) - assert "inconsistent state" in error_text assert "Offload failed" in error_text + assert "inconsistent state" in error_text async def test_compaction_run_failure_shows_error(self) -> None: """Should show error and leave state untouched when the run raises.""" @@ -968,7 +1254,7 @@ async def test_compaction_run_failure_shows_error(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, side_effect=RuntimeError("stream unavailable"), ), @@ -979,6 +1265,51 @@ async def test_compaction_run_failure_shows_error(self) -> None: error_msgs = app.query(ErrorMessage) assert any("Offload failed" in str(w._content) for w in error_msgs) + async def test_remote_exception_is_unwrapped_not_shown_as_a_dict(self) -> None: + """A server-side failure must read as prose, not a Python dict repr. + + The operation graph reports failure by raising, so a server-backed + `/offload` surfaces a `RemoteException` whose sole arg is the server's + error payload dict. `str()` on that renders `{'error': ..., 'message': + ...}`; only `format_agent_exception` unwraps it. Every other test on this + path raises a plain `RuntimeError`, for which the two are identical. + """ + from langgraph.pregel.remote import RemoteException + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + + before = _state_values(_make_dict_messages(10)) + remote_exc = RemoteException( + { + "error": "RuntimeError", + "message": "Compaction failed: OSError: disk full.", + } + ) + + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, before], + ), + patch.object( + app, + "_drive_offload_operation_graph", + new_callable=AsyncMock, + side_effect=remote_exc, + ), + ): + await app._handle_offload() + await pilot.pause() + + error_text = " ".join(str(w._content) for w in app.query(ErrorMessage)) + assert "disk full" in error_text + assert "{'error'" not in error_text + async def test_spinner_hidden_after_failure(self) -> None: """Should hide spinner even when offload fails.""" app = DeepAgentsApp() @@ -997,7 +1328,7 @@ async def test_spinner_hidden_after_failure(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, side_effect=RuntimeError("backend down"), ), @@ -1417,7 +1748,7 @@ async def test_ephemeral_storage_appends_caveat_to_success(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -1452,7 +1783,7 @@ async def test_persistent_storage_omits_caveat(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_offload_operation_graph", new_callable=AsyncMock, return_value=None, ), @@ -1477,9 +1808,15 @@ async def test_cleanup_failure_keeps_noop_report(self) -> None: app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - agent = _setup_server_offload_app(app) - # The no-op branch restores state via aupdate_state; make it fail. + # A plain (non-remote) mock agent drives the seeded in-process + # path, whose no-op branch restores state via `aupdate_state`; + # make that write fail. + agent = MagicMock() agent.aupdate_state = AsyncMock(side_effect=RuntimeError("write failed")) + app._agent = agent + app._backend = None + app._lc_thread_id = "test-thread" + app._agent_running = False before = _state_values(_make_dict_messages(4)) after = _state_values(_make_dict_messages(6)) @@ -1493,7 +1830,7 @@ async def test_cleanup_failure_keeps_noop_report(self) -> None: ), patch.object( app, - "_drive_server_side_compaction", + "_drive_local_seeded_compaction", new_callable=AsyncMock, return_value=None, ), @@ -1621,241 +1958,754 @@ async def test_ordinary_runs_are_unchanged(self) -> None: handler.assert_awaited_once_with(request) -class TestDriveServerSideCompaction: - """Unit-test the server-side `compact_conversation` trigger mechanism.""" +class TestOffloadSessionStartHook: + """`SessionStart(COMPACT)` fires once, from whichever driver ran.""" @staticmethod - def _fake_remote_agent( - tool_content: str, - ) -> tuple[Any, list[Any], list[object]]: - """Build a fake `RemoteAgent` that interrupts then returns a ToolMessage. + def _offloaded_state() -> tuple[dict[str, Any], dict[str, Any]]: + """Build before/after state where compaction advanced the cutoff. - First `astream(None)` surfaces a HITL approval interrupt; the resume - stream (`Command(resume=...)`) yields a `ToolMessage` with the supplied - content so callers can exercise both the success and failure branches. + Returns: + The pre-run and post-run thread state. """ - from langchain_core.messages import ToolMessage - - from deepagents_code.client.remote_client import RemoteAgent + before = _state_values(_make_dict_messages(6)) + after = _state_values(_make_dict_messages(6)) + after["_summarization_event"] = { + "cutoff_index": 4, + "summary_message": {"type": "ai", "content": "summary"}, + "file_path": "/conversation_history/t.md", + } + return before, after - astream_inputs: list[Any] = [] - astream_contexts: list[object] = [] + async def test_operation_graph_path_fires_the_compact_boundary(self) -> None: + """A configured `SessionStart` hook must see `/offload` on this path too. - class _Interrupt: - id = "interrupt-1" - value = { # noqa: RUF012 # test stub; immutability irrelevant - "action_requests": [ - {"name": "compact_conversation", "args": {"force": True}} - ] - } + The seeded driver fires this from inside its drain, keyed on the + compaction tool result. The operation graph produces no tool result, so + without an explicit call here a configured hook would silently never run + for server-backed `/offload` — and every success test would still pass. + """ + from deepagents_code.hooks.models.domain import SessionStartCause - async def _astream(stream_input: object, **kwargs: object): # noqa: RUF029, ANN202 - astream_inputs.append(stream_input) - astream_contexts.append(kwargs.get("context")) - if stream_input is None: - yield ((), "updates", {"__interrupt__": [_Interrupt()]}) - else: - yield ( - (), - "messages", - (ToolMessage(content=tool_content, tool_call_id="x"), {}), - ) + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + before, after = self._offloaded_state() - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream - return agent, astream_inputs, astream_contexts + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, after], + ), + patch.object( + app, + "_drive_offload_operation_graph", + new_callable=AsyncMock, + return_value=None, + ), + patch.object( + app, + "_run_session_start_hook", + new_callable=AsyncMock, + return_value=True, + ) as hook, + ): + await app._handle_offload() + await pilot.pause() - async def test_seeds_tool_call_and_resumes_interrupt(self) -> None: - """Seeds a forced `compact_conversation` call and approves the interrupt.""" - from langgraph.types import Command + hook.assert_awaited_once_with(SessionStartCause.COMPACT) - from deepagents_code.config import settings + async def test_seeded_path_does_not_fire_it_twice(self) -> None: + """The seeded driver owns the boundary, so `_handle_offload` must not. + Both firing would run a user's `SessionStart` hook twice for one + `/offload`. + """ app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - agent, astream_inputs, astream_contexts = self._fake_remote_agent( - "Conversation compacted. Summarized 2 messages into a concise summary." - ) - app._agent = agent - app._lc_thread_id = "test-thread" - app._model_override = "provider:active-model" - app._model_params_override = {"temperature": 0} - app._profile_override = {"max_input_tokens": 4096} + _setup_local_offload_app(app) + before, after = self._offloaded_state() - config = {"configurable": {"thread_id": "test-thread"}} - with patch.object(settings, "model_context_limit", 4096): - result = await app._drive_server_side_compaction(config) # ty: ignore + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, after], + ), + patch.object( + app, + "_drive_local_seeded_compaction", + new_callable=AsyncMock, + return_value=None, + ), + patch.object( + app, + "_run_session_start_hook", + new_callable=AsyncMock, + return_value=True, + ) as hook, + ): + await app._handle_offload() + await pilot.pause() + + hook.assert_not_awaited() + + async def test_hook_stop_still_reports_the_committed_offload(self) -> None: + """A stopping hook must not hide an offload that already committed. + + Compaction is durable by the time this hook runs, so returning early + would leave the user with only "stopped by a hook" while their + conversation *was* compacted and the status bar kept pre-offload counts. + """ + app = DeepAgentsApp() + async with app.run_test() as pilot: await pilot.pause() + _setup_server_offload_app(app) + before, after = self._offloaded_state() + tokens: list[int] = [] - assert result is None + with ( + patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=[before, after], + ), + patch.object( + app, + "_drive_offload_operation_graph", + new_callable=AsyncMock, + return_value=None, + ), + patch.object( + app, + "_run_session_start_hook", + new_callable=AsyncMock, + return_value=False, + ), + patch.object(app, "_on_tokens_update", side_effect=tokens.append), + ): + await app._handle_offload() + await pilot.pause() - # Seed is attributed to the model node so the tool-call routing - # reaches the ToolNode. - agent.aupdate_state.assert_awaited_once() - seed_values = agent.aupdate_state.call_args.args[1] - (seed_msg,) = seed_values["messages"] - (tool_call,) = seed_msg.tool_calls - assert tool_call["name"] == "compact_conversation" - assert tool_call["args"] == {"force": True} - assert agent.aupdate_state.call_args.kwargs["as_node"] == "model" + text = " ".join(str(w._content) for w in app.query(AppMessage)) + assert "Offloaded" in text + # The status bar is refreshed rather than left pre-offload. + assert tokens - # Stream is advanced with None, then resumed after the interrupt. - assert astream_inputs[0] is None - assert isinstance(astream_inputs[1], Command) - resume = astream_inputs[1].resume - assert "interrupt-1" in resume - expected = { - "model": "provider:active-model", - "model_params": {"temperature": 0}, - "profile_overrides": {"max_input_tokens": 4096}, - "model_context_limit": 4096, - "thread_id": "test-thread", - "offload_tool_call_id": tool_call["id"], - } - assert len(astream_contexts) == 2 - for context in astream_contexts: - assert isinstance(context, dict) - normalized = {str(key): value for key, value in context.items()} - assert {key: normalized[key] for key in expected} == expected - async def test_records_summary_and_trailing_usage_in_cost_breakdown(self) -> None: - """Manual offload usage reconciles by type and serving model.""" - from langchain_core.messages import AIMessage, ToolMessage +class TestServerOperationOffload: + """The slash command uses the explicit server operation graph.""" - from deepagents_code.client.remote_client import RemoteAgent + async def test_aborts_when_the_state_refresh_fails(self) -> None: + """A failed re-read aborts the offload instead of replaying the stale snapshot. - class _Interrupt: - id = "interrupt-1" - value = { # noqa: RUF012 # test stub; immutability irrelevant - "action_requests": [ - {"name": "compact_conversation", "args": {"force": True}} - ] - } + The run input replaces the thread's `messages` channel, so falling back + to the caller's pre-run snapshot — taken before `_set_agent_running` + blocked new turns — would delete any message committed in the gap + (shared/external threads, a concurrent completion). No stream may + start when the fresh state cannot be obtained. + """ + app = DeepAgentsApp() + operation = MagicMock() + streamed = False - summary = AIMessage( - content="summary", - id="summary-request", - usage_metadata={ - "input_tokens": 200, - "output_tokens": 20, - "total_tokens": 220, - }, - response_metadata={ - "model_name": "summary-model", - "model_provider": "anthropic", - }, - ) - trailing = AIMessage( - content="done", - id="trailing-request", - usage_metadata={ - "input_tokens": 100, - "output_tokens": 10, - "total_tokens": 110, - }, - response_metadata={ - "model_name": "active-model", - "model_provider": "openai", - }, - ) + async def stream(*_args: object, **_kwargs: object): # noqa: ANN202, RUF029 + nonlocal streamed + streamed = True + yield (), "updates", {"force_compact": {}} - async def _astream( # noqa: ANN202, RUF029 - stream_input: object, **_kwargs: object - ): - if stream_input is None: - yield ((), "updates", {"__interrupt__": [_Interrupt()]}) - return - yield ( - (), - "messages", - (summary, {"lc_source": "summarization"}), - ) - yield ( - (), - "messages", - ( - ToolMessage( - content="Conversation compacted. Summarized 2 messages.", - tool_call_id="compact-call", - ), - {}, - ), + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + side_effect=RuntimeError("checkpoint store unreachable") ) - yield ((), "messages", (trailing, {})) + app._lc_thread_id = "test-thread" + with ( + patch.object(app, "_remote_agent", return_value=remote), + pytest.raises(RuntimeError, match="Could not refresh thread state"), + ): + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream + assert streamed is False + + async def test_aborts_when_the_state_refresh_comes_back_empty(self) -> None: + """An empty re-read aborts the offload rather than replaying stale state. + + `_get_thread_state_values` collapses a missing snapshot (a 404 after a + rebind, a server restart, an un-flushed checkpoint) to `{}`, which the + old `or state_values` fallback silently replaced with the pre-run + snapshot — replaying that would truncate the live conversation. + """ app = DeepAgentsApp() - app._model_override = "openai:active-model" - app._set_session_cost(0.50) - for stats in (app._thread_stats, app._session_stats): - stats.record_request( - "active-model", - 1_000, - 100, - provider="openai", - cost_usd=0.50, - ) + operation = MagicMock() + streamed = False - def _cost( - _usage: object, - model_name: str, - _provider: str = "", - ) -> float: - return {"summary-model": 0.20, "active-model": 0.05}[model_name] + async def stream(*_args: object, **_kwargs: object): # noqa: ANN202, RUF029 + nonlocal streamed + streamed = True + yield (), "updates", {"force_compact": {}} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation async with app.run_test() as pilot: await pilot.pause() - app._agent = agent + app._agent = MagicMock() + app._agent.aget_state = AsyncMock(return_value=SimpleNamespace(values={})) app._lc_thread_id = "test-thread" - with patch( - "deepagents_code.cost_tracking.estimate_cost", side_effect=_cost + with ( + patch.object(app, "_remote_agent", return_value=remote), + pytest.raises(RuntimeError, match="came back empty"), ): - result = await app._drive_server_side_compaction( - {"configurable": {"thread_id": "test-thread"}} + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, ) - await pilot.pause() - assert result is None - assert app._thread_stats.request_count == 3 - assert app._session_stats.request_count == 3 - assert app._thread_stats.per_kind["assistant"].cost_usd == pytest.approx(0.50) - assert app._thread_stats.per_kind["offload"].request_count == 2 - assert app._thread_stats.per_kind["offload"].input_tokens == 300 - assert app._thread_stats.per_kind["offload"].output_tokens == 30 - assert app._thread_stats.per_kind["offload"].cost_usd == pytest.approx(0.25) - assert app._thread_stats.per_model[ - "anthropic", "summary-model" - ].cost_usd == pytest.approx(0.20) - assert app._thread_stats.per_model[ - "openai", "active-model" - ].cost_usd == pytest.approx(0.55) + assert streamed is False - # The estimates keep the running total aligned until the graph-owned - # checkpoint total arrives and clears the provisional amount. - assert app._session_cost_usd == pytest.approx(0.50) - assert app._displayed_cost_usd == pytest.approx(0.75) - app._set_session_cost(0.75) - assert app._displayed_cost_usd == pytest.approx(0.75) - summary_text = app._format_cost_summary() - assert "Estimated thread cost: $0.75" in summary_text - assert "Assistant: $0.50" in summary_text - assert "Offload: $0.25" in summary_text - assert "anthropic:summary-model: $0.20" in summary_text - assert "openai:active-model: $0.55" in summary_text - assert "detailed usage metadata was unavailable" not in summary_text + async def test_streams_named_offload_graph_without_seed_context(self) -> None: + """The operation has no synthetic model call or HITL resume loop.""" + app = DeepAgentsApp() + operation = MagicMock() + stream_args: list[object] = [] + stream_kwargs: dict[str, object] = {} + + async def stream(*args: object, **kwargs: object): # noqa: ANN202, RUF029 + stream_args.extend(args) + stream_kwargs.update(kwargs) + yield (), "updates", {"force_compact": {}} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + messages = [{"type": "human", "content": "hi"}] + # A realistic thread carries far more than `messages`. Every extra + # channel here must be withheld from the run input. + state_values = { + "messages": messages, + "_summarization_event": {"cutoff_index": 2, "summary_message": {}}, + "_session_cost_usd": 0.42, + "_session_cost_transfers": {"scope": {"total": 1.0}}, + "_goal_objective": "ship it", + "todos": [{"content": "t", "status": "pending"}], + } + async with app.run_test() as pilot: + await pilot.pause() + # The driver re-reads thread state before streaming (a stale replay + # would overwrite the live conversation), and that read is what + # registers the thread server-side. + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace(values=state_values) + ) + app._lc_thread_id = "test-thread" + with patch.object(app, "_remote_agent", return_value=remote): + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, state_values + ) - async def test_resume_replay_records_usage_once(self) -> None: - """A usage message replayed after an interrupt is not double-counted.""" - from langchain_core.messages import AIMessage, ToolMessage + remote.aensure_thread.assert_awaited_once_with( + {"configurable": {"thread_id": "test-thread"}} + ) + remote.for_graph.assert_called_once_with("offload") + # `messages` is replayed so the node sees the real conversation rather + # than an emptied list -- and *nothing else* is. Replaying + # `_session_cost_usd` would double the thread's persisted spend on every + # `/offload` (it reduces with `operator.add`), and a writable + # `_summarization_event` would let the caller set the compaction cutoff. + # Asserted by exact key set so a newly-replayed channel fails here. + assert stream_args == [{"messages": messages}] + context = stream_kwargs["context"] + assert isinstance(context, dict) + assert "offload_tool_call_id" not in context + + async def test_forwards_the_active_model_selection_to_the_offload_run( + self, + ) -> None: + """The run context must carry the model the summarizer should use. + + `_runtime_model_config` reads exactly these fields off the context to + build the summarizer, so dropping one silently summarizes with the + startup default under a `/model` or profile override. The seeded driver + has an equivalent test; this path had none, and the integration test's + own `model_context_limit` fix shows the failure mode is live. + """ + app = DeepAgentsApp() + operation = MagicMock() + stream_kwargs: dict[str, object] = {} + + async def stream(*_args: object, **kwargs: object): # noqa: ANN202, RUF029 + stream_kwargs.update(kwargs) + yield (), "updates", {"force_compact": {}} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + app._model_params_override = {"temperature": 0.1} + app._profile_override = {"reasoning": "high"} + with ( + patch.object(app, "_remote_agent", return_value=remote), + patch.object( + app, "_effective_model_spec", return_value="anthropic:some-model" + ), + ): + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) + + context = cast("dict[str, Any]", stream_kwargs["context"]) + assert context["model"] == "anthropic:some-model" + assert context["model_params"] == {"temperature": 0.1} + assert context["profile_overrides"] == {"reasoning": "high"} + assert context["thread_id"] == "test-thread" + from deepagents_code.config import settings + + assert context["model_context_limit"] == settings.model_context_limit + + async def test_an_unanswerable_interrupt_wins_over_fulfillable_ones(self) -> None: + """A mixed round must report, not resume with a partial answer. + + One round can yield both a hook interrupt this client can fulfill and a + `HITLRequest` it cannot. Resuming with only the fulfillable half leaves + the other pending forever, so the unanswerable one has to decide the + round. Plausibly correct but unpinned before this test — a reordering + would flip it silently. + """ + app = DeepAgentsApp() + operation = MagicMock() + hook_interrupt = MagicMock() + hook_interrupt.id = "interrupt-1" + hook_interrupt.value = {"type": "hook_invocation", "invocation_id": "inv-1"} + alien = MagicMock() + alien.id = "interrupt-2" + alien.value = {"type": "hitl_request"} + streams: list[object] = [] + + async def stream(stream_input: object, **_kwargs: object): # noqa: ANN202, RUF029 + streams.append(stream_input) + yield (), "updates", {"__interrupt__": [hook_interrupt, alien]} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + manager = MagicMock(spec=HooksManager) + manager.apply_graph_context = MagicMock() + manager.fulfill_interrupt = AsyncMock(return_value={"decision": "ok"}) + with ( + patch.object(app, "_remote_agent", return_value=remote), + self._hooks_patch(manager), + ): + result = await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) + + assert result is not None + assert "cannot answer" in result + # No resume round: the fulfillable half is deliberately discarded. + assert len(streams) == 1 + + async def test_restores_main_graph_association(self) -> None: + """A `graph_id` rebind from the named-graph run is reset before returning. + + The server records the last-run graph on the thread; leaving it at + `offload` would send later out-of-run `as_node="model"` state writes to + a graph with no model node. The driver rebinds the thread's `graph_id` + metadata through the main client. An empty `aupdate_state` cannot do + this -- it resolves against the thread's *current* graph association -- + so the rebind must go through `arebind_thread`. + """ + app = DeepAgentsApp() + operation = MagicMock() + + async def stream(*_args: object, **_kwargs: object): # noqa: ANN202, RUF029 + yield (), "updates", {"force_compact": {}} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + with patch.object(app, "_remote_agent", return_value=remote): + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) + + remote.arebind_thread.assert_awaited_once_with( + {"configurable": {"thread_id": "test-thread"}} + ) + + async def test_graph_restore_failure_warns_without_failing_the_offload( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A failed rebind is recorded for the caller, not mounted here. + + The offload itself succeeded, so this must not surface as a failure. It + must not be silent either: the thread stays bound to the `offload` + graph, so an unrelated later `/goal` or `/rubric` would fail with no + explanation. The driver only records it -- `_OFFLOAD_REBIND_WARNING` + says the offload finished, and only `_handle_offload` knows whether it + is about to report that. + """ + app = DeepAgentsApp() + operation = MagicMock() + + async def stream(*_args: object, **_kwargs: object): # noqa: ANN202, RUF029 + yield (), "updates", {"force_compact": {}} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock(side_effect=RuntimeError("server gone")) + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + with ( + patch.object(app, "_remote_agent", return_value=remote), + caplog.at_level("WARNING"), + ): + result = await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) + await pilot.pause() + + # Not reported as a drain failure -- the offload still succeeded. + assert result is None + assert app._offload_rebind_failed is True + error_text = " ".join( + str(widget._content) for widget in app.query(ErrorMessage) + ) + + assert "Failed to restore the thread's main graph association" in caplog.text + # Recorded, not rendered: the caller decides when this is safe to say. + assert error_text == "" + + async def test_stream_error_is_not_masked_by_a_failing_rebind(self) -> None: + """A rebind failure in the `finally` must not replace the stream error. + + The caller distinguishes a committed-but-interrupted offload from a + failed one by the exception it sees; swallowing the real error in favor + of the bookkeeping failure would lose that. The driver mounts nothing + either way — it records the rebind failure and lets `_handle_offload` + decide, which is what lets a stream error that still committed the + compaction be reconciled into a success *and* carry the warning. + """ + from deepagents_code.app import _OFFLOAD_REBIND_WARNING + + app = DeepAgentsApp() + operation = MagicMock() + + async def stream(*_args: object, **_kwargs: object): # noqa: ANN202, RUF029 + yield (), "updates", {"force_compact": {}} + msg = "stream died" + raise RuntimeError(msg) + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock(side_effect=RuntimeError("server gone")) + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + with ( + patch.object(app, "_remote_agent", return_value=remote), + pytest.raises(RuntimeError, match="stream died"), + ): + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) + await pilot.pause() + error_text = " ".join( + str(widget._content) for widget in app.query(ErrorMessage) + ) + + assert _OFFLOAD_REBIND_WARNING not in error_text + # Still recorded, so a caller that reconciles this into a success can + # surface it there. + assert app._offload_rebind_failed is True + + @staticmethod + def _hooks_patch(manager: MagicMock): # noqa: ANN205 + """Swap the read-only `_hooks` property for a stubbed manager.""" + return patch.object( + DeepAgentsApp, + "_hooks", + new_callable=lambda: property(lambda _app: manager), + ) + + async def test_fulfills_hook_interrupts_and_resumes(self) -> None: + """A hook interrupt is answered through the hook engine and the graph resumed. + + With a configured `PreCompact`/`PreToolUse` hook the node interrupts at + the hook boundary instead of returning; without this loop `/offload` + would park there and report "Nothing to offload". + """ + from deepagents_code.hooks.manager import HooksManager + + app = DeepAgentsApp() + operation = MagicMock() + streams: list[object] = [] + hook_payload = {"type": "hook_invocation", "invocation_id": "inv-1"} + hook_interrupt = MagicMock() + hook_interrupt.id = "interrupt-1" + hook_interrupt.value = hook_payload + + async def stream(input_: object, **_kwargs: object): # noqa: ANN202, RUF029 + streams.append(input_) + if len(streams) == 1: + yield (), "updates", {"__interrupt__": [hook_interrupt]} + else: + yield (), "updates", {"force_compact": {}} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + manager = MagicMock(spec=HooksManager) + manager.apply_graph_context = MagicMock() + manager.fulfill_interrupt = AsyncMock(return_value={"decision": "ok"}) + with ( + patch.object(app, "_remote_agent", return_value=remote), + self._hooks_patch(manager), + ): + result = await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) + + assert result is None + manager.fulfill_interrupt.assert_awaited_once_with(hook_payload) + assert len(streams) == 2 + resume = streams[1] + # The resume command carries the fulfilled hook decision keyed by the + # server-supplied interrupt id. + assert getattr(resume, "resume", None) == {"interrupt-1": {"decision": "ok"}} + + async def test_unbounded_hook_interrupts_stop_at_resume_cap(self) -> None: + """A hook interrupting every round is bounded *and* reported. + + Reporting is the load-bearing half. The driver returns `None` for a + successful run, and the caller reads an unchanged `_summarization_event` + as "nothing to offload" -- so a silent give-up would tell the user their + conversation is already compact while the run sits paused mid-interrupt + and nothing was compacted at all. + """ + app = DeepAgentsApp() + operation = MagicMock() + hook_interrupt = MagicMock() + hook_interrupt.id = "interrupt-1" + hook_interrupt.value = {"type": "hook_invocation", "invocation_id": "inv-1"} + rounds = 0 + + async def stream(_input: object, **_kwargs: object): # noqa: ANN202, RUF029 + nonlocal rounds + rounds += 1 + yield (), "updates", {"__interrupt__": [hook_interrupt]} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + manager = MagicMock(spec=HooksManager) + manager.apply_graph_context = MagicMock() + manager.fulfill_interrupt = AsyncMock(return_value={}) + with ( + patch.object(app, "_remote_agent", return_value=remote), + self._hooks_patch(manager), + ): + result = await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) + + from deepagents_code.app import _OFFLOAD_MAX_RESUME_ROUNDS + + assert rounds == _OFFLOAD_MAX_RESUME_ROUNDS + 1 + assert result is not None + assert "could not complete" in result + assert str(_OFFLOAD_MAX_RESUME_ROUNDS) in result + + async def test_unanswerable_interrupt_is_reported_not_silently_dropped( + self, + ) -> None: + """An approval-shaped interrupt this client cannot answer is an error. + + A `PreToolUse` hook returning an `ask` permission makes the hook + middleware raise a plain `HITLRequest`, which is not a hook-invocation + payload. The operation graph has no HITL middleware to route it, so the + run stays paused; dropping it silently would surface as "the + conversation is already compact". + """ + app = DeepAgentsApp() + operation = MagicMock() + approval = MagicMock() + approval.id = "interrupt-1" + approval.value = {"action_request": {"action": "compact_conversation"}} + rounds = 0 + + async def stream(_input: object, **_kwargs: object): # noqa: ANN202, RUF029 + nonlocal rounds + rounds += 1 + yield (), "updates", {"__interrupt__": [approval]} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + manager = MagicMock(spec=HooksManager) + manager.apply_graph_context = MagicMock() + manager.fulfill_interrupt = AsyncMock() + with ( + patch.object(app, "_remote_agent", return_value=remote), + self._hooks_patch(manager), + ): + result = await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + {"messages": [{"type": "human", "content": "hi"}]}, + ) + + # Bails on the first round rather than burning the whole resume budget, + # and never tries to fulfill a payload the hook engine cannot parse. + assert rounds == 1 + assert result is not None + assert "approval" in result + manager.fulfill_interrupt.assert_not_awaited() + + +class TestDriveLegacySeededCompaction: + """Unit-test the seeded in-process `compact_conversation` trigger. + + This driver serves local `Pregel` agents, which have no server operation + graph; server-backed agents use the dedicated `offload` operation instead. + """ + + @staticmethod + def _fake_remote_agent( + tool_content: str, + ) -> tuple[Any, list[Any], list[object]]: + """Build a fake `RemoteAgent` that interrupts then returns a ToolMessage. + + First `astream(None)` surfaces a HITL approval interrupt; the resume + stream (`Command(resume=...)`) yields a `ToolMessage` with the supplied + content so callers can exercise both the success and failure branches. + """ + from langchain_core.messages import ToolMessage from deepagents_code.client.remote_client import RemoteAgent + astream_inputs: list[Any] = [] + astream_contexts: list[object] = [] + class _Interrupt: id = "interrupt-1" value = { # noqa: RUF012 # test stub; immutability irrelevant @@ -1864,53 +2714,268 @@ class _Interrupt: ] } - usage_message = AIMessage( - content="summary", - id="replayed-request", - usage_metadata={ - "input_tokens": 200, - "output_tokens": 20, - "total_tokens": 220, - }, - response_metadata={"model_name": "summary-model"}, - ) - - async def _astream( # noqa: ANN202, RUF029 - stream_input: object, **_kwargs: object - ): - yield ( - (), - "messages", - (usage_message, {"lc_source": "summarization"}), - ) + async def _astream(stream_input: object, **kwargs: object): # noqa: RUF029, ANN202 + astream_inputs.append(stream_input) + astream_contexts.append(kwargs.get("context")) if stream_input is None: yield ((), "updates", {"__interrupt__": [_Interrupt()]}) else: yield ( (), "messages", - (ToolMessage(content="Nothing to compact", tool_call_id="x"), {}), + (ToolMessage(content=tool_content, tool_call_id="x"), {}), ) agent = MagicMock(spec=RemoteAgent) agent.aensure_thread = AsyncMock() agent.aupdate_state = AsyncMock() agent.astream = _astream - app = DeepAgentsApp() + return agent, astream_inputs, astream_contexts + + async def test_seeds_tool_call_and_resumes_interrupt(self) -> None: + """Seeds a forced `compact_conversation` call and approves the interrupt.""" + from langgraph.types import Command + + from deepagents_code.config import settings + app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() + agent, astream_inputs, astream_contexts = self._fake_remote_agent( + "Conversation compacted. Summarized 2 messages into a concise summary." + ) app._agent = agent app._lc_thread_id = "test-thread" - with patch( - "deepagents_code.cost_tracking.estimate_cost", return_value=0.20 - ): - result = await app._drive_server_side_compaction( - {"configurable": {"thread_id": "test-thread"}} - ) - await pilot.pause() - - assert result is None + app._model_override = "provider:active-model" + app._model_params_override = {"temperature": 0} + app._profile_override = {"max_input_tokens": 4096} + + config = {"configurable": {"thread_id": "test-thread"}} + with patch.object(settings, "model_context_limit", 4096): + result = await app._drive_local_seeded_compaction(config) # ty: ignore + await pilot.pause() + + assert result is None + + # Seed is attributed to the model node so the tool-call routing + # reaches the ToolNode. + agent.aupdate_state.assert_awaited_once() + seed_values = agent.aupdate_state.call_args.args[1] + (seed_msg,) = seed_values["messages"] + (tool_call,) = seed_msg.tool_calls + assert tool_call["name"] == "compact_conversation" + assert tool_call["args"] == {"force": True} + assert agent.aupdate_state.call_args.kwargs["as_node"] == "model" + + # Stream is advanced with None, then resumed after the interrupt. + assert astream_inputs[0] is None + assert isinstance(astream_inputs[1], Command) + resume = astream_inputs[1].resume + assert "interrupt-1" in resume + expected = { + "model": "provider:active-model", + "model_params": {"temperature": 0}, + "profile_overrides": {"max_input_tokens": 4096}, + "model_context_limit": 4096, + "thread_id": "test-thread", + "offload_tool_call_id": tool_call["id"], + } + assert len(astream_contexts) == 2 + for context in astream_contexts: + assert isinstance(context, dict) + normalized = {str(key): value for key, value in context.items()} + assert {key: normalized[key] for key in expected} == expected + + async def test_records_summary_and_trailing_usage_in_cost_breakdown(self) -> None: + """Manual offload usage reconciles by type and serving model.""" + from langchain_core.messages import AIMessage, ToolMessage + + from deepagents_code.client.remote_client import RemoteAgent + + class _Interrupt: + id = "interrupt-1" + value = { # noqa: RUF012 # test stub; immutability irrelevant + "action_requests": [ + {"name": "compact_conversation", "args": {"force": True}} + ] + } + + summary = AIMessage( + content="summary", + id="summary-request", + usage_metadata={ + "input_tokens": 200, + "output_tokens": 20, + "total_tokens": 220, + }, + response_metadata={ + "model_name": "summary-model", + "model_provider": "anthropic", + }, + ) + trailing = AIMessage( + content="done", + id="trailing-request", + usage_metadata={ + "input_tokens": 100, + "output_tokens": 10, + "total_tokens": 110, + }, + response_metadata={ + "model_name": "active-model", + "model_provider": "openai", + }, + ) + + async def _astream( # noqa: ANN202, RUF029 + stream_input: object, **_kwargs: object + ): + if stream_input is None: + yield ((), "updates", {"__interrupt__": [_Interrupt()]}) + return + yield ( + (), + "messages", + (summary, {"lc_source": "summarization"}), + ) + yield ( + (), + "messages", + ( + ToolMessage( + content="Conversation compacted. Summarized 2 messages.", + tool_call_id="compact-call", + ), + {}, + ), + ) + yield ((), "messages", (trailing, {})) + + agent = MagicMock(spec=RemoteAgent) + agent.aensure_thread = AsyncMock() + agent.aupdate_state = AsyncMock() + agent.astream = _astream + app = DeepAgentsApp() + app._model_override = "openai:active-model" + app._set_session_cost(0.50) + for stats in (app._thread_stats, app._session_stats): + stats.record_request( + "active-model", + 1_000, + 100, + provider="openai", + cost_usd=0.50, + ) + + def _cost( + _usage: object, + model_name: str, + _provider: str = "", + ) -> float: + return {"summary-model": 0.20, "active-model": 0.05}[model_name] + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = agent + app._lc_thread_id = "test-thread" + with patch( + "deepagents_code.cost_tracking.estimate_cost", side_effect=_cost + ): + result = await app._drive_local_seeded_compaction( + {"configurable": {"thread_id": "test-thread"}} + ) + await pilot.pause() + + assert result is None + assert app._thread_stats.request_count == 3 + assert app._session_stats.request_count == 3 + assert app._thread_stats.per_kind["assistant"].cost_usd == pytest.approx(0.50) + assert app._thread_stats.per_kind["offload"].request_count == 2 + assert app._thread_stats.per_kind["offload"].input_tokens == 300 + assert app._thread_stats.per_kind["offload"].output_tokens == 30 + assert app._thread_stats.per_kind["offload"].cost_usd == pytest.approx(0.25) + assert app._thread_stats.per_model[ + "anthropic", "summary-model" + ].cost_usd == pytest.approx(0.20) + assert app._thread_stats.per_model[ + "openai", "active-model" + ].cost_usd == pytest.approx(0.55) + + # The estimates keep the running total aligned until the graph-owned + # checkpoint total arrives and clears the provisional amount. + assert app._session_cost_usd == pytest.approx(0.50) + assert app._displayed_cost_usd == pytest.approx(0.75) + app._set_session_cost(0.75) + assert app._displayed_cost_usd == pytest.approx(0.75) + summary_text = app._format_cost_summary() + assert "Estimated thread cost: $0.75" in summary_text + assert "Assistant: $0.50" in summary_text + assert "Offload: $0.25" in summary_text + assert "anthropic:summary-model: $0.20" in summary_text + assert "openai:active-model: $0.55" in summary_text + assert "detailed usage metadata was unavailable" not in summary_text + + async def test_resume_replay_records_usage_once(self) -> None: + """A usage message replayed after an interrupt is not double-counted.""" + from langchain_core.messages import AIMessage, ToolMessage + + from deepagents_code.client.remote_client import RemoteAgent + + class _Interrupt: + id = "interrupt-1" + value = { # noqa: RUF012 # test stub; immutability irrelevant + "action_requests": [ + {"name": "compact_conversation", "args": {"force": True}} + ] + } + + usage_message = AIMessage( + content="summary", + id="replayed-request", + usage_metadata={ + "input_tokens": 200, + "output_tokens": 20, + "total_tokens": 220, + }, + response_metadata={"model_name": "summary-model"}, + ) + + async def _astream( # noqa: ANN202, RUF029 + stream_input: object, **_kwargs: object + ): + yield ( + (), + "messages", + (usage_message, {"lc_source": "summarization"}), + ) + if stream_input is None: + yield ((), "updates", {"__interrupt__": [_Interrupt()]}) + else: + yield ( + (), + "messages", + (ToolMessage(content="Nothing to compact", tool_call_id="x"), {}), + ) + + agent = MagicMock(spec=RemoteAgent) + agent.aensure_thread = AsyncMock() + agent.aupdate_state = AsyncMock() + agent.astream = _astream + app = DeepAgentsApp() + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = agent + app._lc_thread_id = "test-thread" + with patch( + "deepagents_code.cost_tracking.estimate_cost", return_value=0.20 + ): + result = await app._drive_local_seeded_compaction( + {"configurable": {"thread_id": "test-thread"}} + ) + await pilot.pause() + + assert result is None assert app._thread_stats.request_count == 1 assert app._session_stats.request_count == 1 assert app._thread_stats.per_kind["offload"].cost_usd == pytest.approx(0.20) @@ -1937,669 +3002,1939 @@ async def test_stream_failure_keeps_usage_recorded_once(self) -> None: response_metadata={"model_name": "summary-model"}, ) - async def _astream( # noqa: ANN202, RUF029 - _stream_input: object, **_kwargs: object - ): - for _ in range(2): - yield ( - (), - "messages", - (usage_message, {"lc_source": "summarization"}), - ) - msg = "stream failed" - raise RuntimeError(msg) + async def _astream( # noqa: ANN202, RUF029 + _stream_input: object, **_kwargs: object + ): + for _ in range(2): + yield ( + (), + "messages", + (usage_message, {"lc_source": "summarization"}), + ) + msg = "stream failed" + raise RuntimeError(msg) + + agent = MagicMock(spec=RemoteAgent) + agent.aensure_thread = AsyncMock() + agent.aupdate_state = AsyncMock() + agent.astream = _astream + app = DeepAgentsApp() + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = agent + app._lc_thread_id = "test-thread" + with ( + patch("deepagents_code.cost_tracking.estimate_cost", return_value=0.20), + pytest.raises(RuntimeError, match="stream failed"), + ): + await app._drive_local_seeded_compaction( + {"configurable": {"thread_id": "test-thread"}} + ) + await pilot.pause() + + assert app._thread_stats.request_count == 1 + assert app._session_stats.request_count == 1 + assert app._thread_stats.per_kind["offload"].cost_usd == pytest.approx(0.20) + assert app._session_cost_usd == pytest.approx(0.0) + assert app._displayed_cost_usd == pytest.approx(0.20) + summary = app._format_cost_summary() + assert "Estimated thread cost: $0.20" in summary + assert "Offload: $0.20" in summary + + async def test_fulfills_precompact_before_manual_approval(self) -> None: + """A precompact hook is fulfilled before the compaction approval.""" + from types import SimpleNamespace + + from langchain_core.messages import ToolMessage + from langgraph.types import Command + + from deepagents_code.client.remote_client import RemoteAgent + from deepagents_code.hooks.interrupt import HOOK_INVOCATION_INTERRUPT_TYPE + + streams: list[object] = [] + + async def _astream( # noqa: ANN202, RUF029 + value: object, **_kwargs: object + ): + index = len(streams) + streams.append(value) + if index == 0: + interrupt = SimpleNamespace( + id="hook-interrupt", + value={"type": HOOK_INVOCATION_INTERRUPT_TYPE}, + ) + elif index == 1: + interrupt = SimpleNamespace( + id="approval-interrupt", + value={ + "action_requests": [ + { + "name": "compact_conversation", + "args": {"force": True}, + } + ] + }, + ) + else: + yield ( + (), + "messages", + (ToolMessage(content="compacted", tool_call_id="compact-call"), {}), + ) + return + yield ((), "updates", {"__interrupt__": [interrupt]}) + + agent = MagicMock( + spec=RemoteAgent, + aensure_thread=AsyncMock(), + aupdate_state=AsyncMock(), + astream=_astream, + ) + app = DeepAgentsApp() + + async with app.run_test() as pilot: + await pilot.pause() + runtime = MagicMock(snapshot_id="snapshot") + runtime.configured_server_events.return_value = ("PreCompact",) + assert app._session_state is not None + app._session_state.hooks = HooksManager.adopting( + runtime, + identity=app._session_state.hook_identity, + ) + app._agent = agent + app._lc_thread_id = "test-thread" + fulfill = AsyncMock(return_value={"hook": "approved"}) + with patch("deepagents_code.hooks.client.fulfill_hook_interrupt", fulfill): + result = await app._drive_local_seeded_compaction( + {"configurable": {"thread_id": "test-thread"}} + ) + + assert result is None + fulfill.assert_awaited_once() + assert len(streams) == 3 + assert isinstance(streams[1], Command) + assert streams[1].resume == {"hook-interrupt": {"hook": "approved"}} + assert isinstance(streams[2], Command) + approval = streams[2].resume + assert isinstance(approval, dict) + assert "approval-interrupt" in approval + + async def test_reports_tool_failure(self) -> None: + """Returns the tool's error text when compaction fails.""" + from deepagents_code.offload_middleware import COMPACTION_FAILURE_PREFIX + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + agent, _inputs, _contexts = self._fake_remote_agent( + f"{COMPACTION_FAILURE_PREFIX}: an error occurred during compaction." + ) + app._agent = agent + app._lc_thread_id = "test-thread" + + config = {"configurable": {"thread_id": "test-thread"}} + result = await app._drive_local_seeded_compaction(config) # ty: ignore + await pilot.pause() + + assert result is not None + assert result.startswith(COMPACTION_FAILURE_PREFIX) + + async def test_forwards_startup_model_profile_to_compaction(self) -> None: + """Profile data is usable even without a session `/model` override.""" + from deepagents_code.config import settings + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + agent, _inputs, contexts = self._fake_remote_agent( + "Conversation compacted. Summarized 2 messages." + ) + app._agent = agent + app._lc_thread_id = "test-thread" + app._model_override = None + app._profile_override = {"max_input_tokens": 4096} + + config = {"configurable": {"thread_id": "test-thread"}} + with ( + patch.object(settings, "model_provider", "provider"), + patch.object(settings, "model_name", "startup-model"), + patch.object(settings, "model_context_limit", 4096), + ): + await app._drive_local_seeded_compaction(config) # ty: ignore + await pilot.pause() + + assert contexts + seed_values = agent.aupdate_state.call_args.args[1] + (seed_msg,) = seed_values["messages"] + (tool_call,) = seed_msg.tool_calls + expected = { + "model": "provider:startup-model", + "model_params": {}, + "profile_overrides": {"max_input_tokens": 4096}, + "model_context_limit": 4096, + "thread_id": "test-thread", + "offload_tool_call_id": tool_call["id"], + } + for context in contexts: + assert isinstance(context, dict) + normalized = {str(key): value for key, value in context.items()} + assert {key: normalized[key] for key in expected} == expected + + async def test_rejects_interrupt_without_identifiable_action(self) -> None: + """Malformed interrupt payloads fail closed instead of being approved.""" + from langgraph.types import Command + + from deepagents_code.client.remote_client import RemoteAgent + + astream_inputs: list[Any] = [] + + class _Interrupt: + id = "interrupt-unknown" + value: dict[str, Any] = {} # noqa: RUF012 # test stub + + async def _astream( # noqa: RUF029, ANN202 + stream_input: object, **_kwargs: object + ): + astream_inputs.append(stream_input) + if stream_input is None: + yield ((), "updates", {"__interrupt__": [_Interrupt()]}) + + agent = MagicMock(spec=RemoteAgent) + agent.aensure_thread = AsyncMock() + agent.aupdate_state = AsyncMock() + agent.astream = _astream + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + app._agent = agent + app._lc_thread_id = "test-thread" + + config = {"configurable": {"thread_id": "test-thread"}} + result = await app._drive_local_seeded_compaction(config) # ty: ignore + await pilot.pause() + + assert result is None + assert len(astream_inputs) == 2 + assert isinstance(astream_inputs[1], Command) + decision = astream_inputs[1].resume["interrupt-unknown"]["decisions"][0] + assert decision["type"] == "reject" + + async def test_approves_only_first_forced_compaction(self) -> None: + """A repeated forced compaction request is rejected, not approved.""" + from langchain_core.messages import ToolMessage + from langgraph.types import Command + + from deepagents_code.client.remote_client import RemoteAgent + + astream_inputs: list[Any] = [] + guard_ids: list[object] = [] + + class _Interrupt: + def __init__(self, iid: str, tool_name: str, args: dict[str, Any]) -> None: + self.id = iid + self.value = {"action_requests": [{"name": tool_name, "args": args}]} + + async def _astream(stream_input: object, **kwargs: object): # noqa: RUF029, ANN202 + idx = len(astream_inputs) + astream_inputs.append(stream_input) + context = kwargs.get("context") + guard_ids.append( + context.get("offload_tool_call_id") + if isinstance(context, dict) + else None + ) + if idx == 0: + compact = _Interrupt( + "i-compact", "compact_conversation", {"force": True} + ) + yield ((), "updates", {"__interrupt__": [compact]}) + elif idx == 1: + # Model a trailing turn that asks to compact again. + repeated = _Interrupt( + "i-repeated", "compact_conversation", {"force": True} + ) + yield ((), "updates", {"__interrupt__": [repeated]}) + else: + yield ( + (), + "messages", + ( + ToolMessage( + content="Conversation compacted. Summarized 2 messages " + "into a concise summary.", + tool_call_id="x", + ), + {}, + ), + ) + + agent = MagicMock(spec=RemoteAgent) + agent.aensure_thread = AsyncMock() + agent.aupdate_state = AsyncMock() + agent.astream = _astream + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + app._agent = agent + app._lc_thread_id = "test-thread" + + config = {"configurable": {"thread_id": "test-thread"}} + result = await app._drive_local_seeded_compaction(config) # ty: ignore + await pilot.pause() + + assert result is None + # Initial drain + two resumes (compaction, then trailing tool). + assert len(astream_inputs) == 3 + assert isinstance(astream_inputs[1], Command) + assert isinstance(astream_inputs[2], Command) + assert len(set(guard_ids)) == 1 + assert isinstance(guard_ids[0], str) + # Compaction was approved. + compact_decision = astream_inputs[1].resume["i-compact"]["decisions"][0] + assert compact_decision["type"] == "approve" + # A second compaction request is not the seeded call and is rejected. + repeated_decision = astream_inputs[2].resume["i-repeated"]["decisions"][0] + assert repeated_decision["type"] == "reject" + + async def test_sets_tool_guard_context_without_hitl(self) -> None: + """The per-run tool guard is set even when no HITL interrupt exists.""" + from langchain_core.messages import ToolMessage + + from deepagents_code.client.remote_client import RemoteAgent + + guard_ids: list[object] = [] + + async def _astream(_stream_input: object, **kwargs: object): # noqa: RUF029, ANN202 + context = kwargs.get("context") + guard_ids.append( + context.get("offload_tool_call_id") + if isinstance(context, dict) + else None + ) + yield ( + (), + "messages", + ( + ToolMessage( + content="Conversation compacted. Summarized 2 messages.", + tool_call_id="x", + ), + {}, + ), + ) + + agent = MagicMock(spec=RemoteAgent) + agent.aensure_thread = AsyncMock() + agent.aupdate_state = AsyncMock() + agent.astream = _astream + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + app._agent = agent + app._lc_thread_id = "test-thread" + + config = {"configurable": {"thread_id": "test-thread"}} + result = await app._drive_local_seeded_compaction(config) # ty: ignore + await pilot.pause() + + assert result is None + seed_values = agent.aupdate_state.call_args.args[1] + (seed_msg,) = seed_values["messages"] + (tool_call,) = seed_msg.tool_calls + assert guard_ids == [tool_call["id"]] + + async def test_bounds_resume_loop_and_reports_abandoned_drain(self) -> None: + """A model that keeps requesting tools cannot spin `/offload` forever. + + Every stream yields a fresh gated interrupt, so the resume loop never + drains cleanly. It must stop at the `max_resume_rounds` cap (initial + drain + 10 resumes = 11 streams) and surface a user-visible notice that + the run was left paused, rather than looping indefinitely. + """ + from deepagents_code.client.remote_client import RemoteAgent + + astream_inputs: list[Any] = [] + + class _Interrupt: + def __init__(self, iid: str) -> None: + self.id = iid + self.value = {"action_requests": [{"name": "write_file", "args": {}}]} + + async def _astream(stream_input: object, **_kwargs: object): # noqa: RUF029, ANN202 + idx = len(astream_inputs) + astream_inputs.append(stream_input) + # Never terminate: each round surfaces another gated interrupt. + yield ((), "updates", {"__interrupt__": [_Interrupt(f"i-{idx}")]}) + + agent = MagicMock(spec=RemoteAgent) + agent.aensure_thread = AsyncMock() + agent.aupdate_state = AsyncMock() + agent.astream = _astream + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + app._agent = agent + app._lc_thread_id = "test-thread" + + config = {"configurable": {"thread_id": "test-thread"}} + result = await app._drive_local_seeded_compaction(config) # ty: ignore + await pilot.pause() + + # No compaction failure was reported, so the run returns cleanly. + assert result is None + # Initial drain + exactly 10 resume rounds, then the cap breaks. + assert len(astream_inputs) == 11 + assert any( + "could not be fully drained" in str(widget._content) + for widget in app.query(ErrorMessage) + ) + + +class TestRemoveUnansweredOffloadSeed: + """Cleanup of a committed-but-unanswered `/offload` seed after a failure.""" + + @staticmethod + def _seed_message(tool_call_id: str) -> dict[str, Any]: + """Serialized seed AIMessage carrying the forced compaction tool call.""" + return { + "type": "ai", + "content": "", + "id": f"offload-seed-{tool_call_id}", + "tool_calls": [ + { + "name": "compact_conversation", + "args": {"force": True}, + "id": tool_call_id, + } + ], + } + + async def test_removes_dangling_seed(self) -> None: + """An unanswered seed is removed so it cannot wedge the next turn.""" + from langchain_core.messages import RemoveMessage + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + agent = MagicMock() + agent.aupdate_state = AsyncMock() + app._agent = agent + state = _state_values( + [*_make_dict_messages(2), self._seed_message("seed-call")] + ) + with patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + return_value=state, + ): + await app._remove_unanswered_offload_seed( + {"configurable": {"thread_id": "test-thread"}}, "seed-call" + ) + + agent.aupdate_state.assert_awaited_once() + update = agent.aupdate_state.call_args.args[1] + (removal,) = update["messages"] + assert isinstance(removal, RemoveMessage) + assert removal.id == "offload-seed-seed-call" + + async def test_keeps_answered_seed(self) -> None: + """A seed answered by a ToolMessage is a valid pair and is left intact.""" + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + agent = MagicMock() + agent.aupdate_state = AsyncMock() + app._agent = agent + answered = { + "type": "tool", + "content": "Nothing to compact yet.", + "tool_call_id": "seed-call", + } + state = _state_values( + [*_make_dict_messages(2), self._seed_message("seed-call"), answered] + ) + with patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + return_value=state, + ): + await app._remove_unanswered_offload_seed( + {"configurable": {"thread_id": "test-thread"}}, "seed-call" + ) + + agent.aupdate_state.assert_not_awaited() + + async def test_noop_when_seed_absent(self) -> None: + """Nothing is removed when no seed with the id is present.""" + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + agent = MagicMock() + agent.aupdate_state = AsyncMock() + app._agent = agent + state = _state_values(_make_dict_messages(2)) + with patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + return_value=state, + ): + await app._remove_unanswered_offload_seed( + {"configurable": {"thread_id": "test-thread"}}, "seed-call" + ) + + agent.aupdate_state.assert_not_awaited() + + async def test_returns_true_when_seed_removed(self) -> None: + """Successful removal reports the thread is clean.""" + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + agent = MagicMock() + agent.aupdate_state = AsyncMock() + app._agent = agent + state = _state_values( + [*_make_dict_messages(2), self._seed_message("seed-call")] + ) + with patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + return_value=state, + ): + cleaned = await app._remove_unanswered_offload_seed( + {"configurable": {"thread_id": "test-thread"}}, "seed-call" + ) + + assert cleaned is True + + async def test_returns_false_when_state_read_fails(self) -> None: + """A failed state read cannot confirm cleanup, so it reports unclean.""" + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + agent = MagicMock() + agent.aupdate_state = AsyncMock() + app._agent = agent + with patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + side_effect=RuntimeError("state read boom"), + ): + cleaned = await app._remove_unanswered_offload_seed( + {"configurable": {"thread_id": "test-thread"}}, "seed-call" + ) + + assert cleaned is False + # The dangling seed could not be removed, so nothing was written. + agent.aupdate_state.assert_not_awaited() + + async def test_returns_false_when_removal_write_fails(self) -> None: + """A failed removal write leaves the seed and reports unclean.""" + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + agent = MagicMock() + agent.aupdate_state = AsyncMock(side_effect=RuntimeError("write boom")) + app._agent = agent + state = _state_values( + [*_make_dict_messages(2), self._seed_message("seed-call")] + ) + with patch.object( + app, + "_get_thread_state_values", + new_callable=AsyncMock, + return_value=state, + ): + cleaned = await app._remove_unanswered_offload_seed( + {"configurable": {"thread_id": "test-thread"}}, "seed-call" + ) + + assert cleaned is False + + +class TestFormatTokenCount: + """Test the format_token_count helper function.""" + + def test_zero(self) -> None: + assert format_token_count(0) == "0" + + def test_below_threshold(self) -> None: + assert format_token_count(999) == "999" + + def test_at_threshold(self) -> None: + assert format_token_count(1000) == "1.0K" + + def test_above_threshold(self) -> None: + assert format_token_count(1500) == "1.5K" + + def test_large_value(self) -> None: + assert format_token_count(200000) == "200.0K" + + def test_millions(self) -> None: + assert format_token_count(1_000_000) == "1.0M" + + def test_above_million(self) -> None: + assert format_token_count(2_500_000) == "2.5M" + + +class TestOffloadHelpers: + """Pure helpers backing `/offload` accounting and failure detection.""" + + def test_summarization_cutoff_reads_int(self) -> None: + from deepagents_code.app import _summarization_cutoff + + assert _summarization_cutoff({"cutoff_index": 4}) == 4 + + def test_summarization_cutoff_defaults_zero_on_malformed(self) -> None: + from deepagents_code.app import _summarization_cutoff + + assert _summarization_cutoff(None) == 0 + assert _summarization_cutoff({"cutoff_index": "x"}) == 0 + assert _summarization_cutoff({}) == 0 + assert _summarization_cutoff("not-a-dict") == 0 + + def test_effective_conversation_applies_event(self) -> None: + from deepagents_code.app import _effective_conversation + + messages = [f"m{i}" for i in range(5)] + event = {"summary_message": "S", "cutoff_index": 2} + assert _effective_conversation(messages, event) == ["S", "m2", "m3", "m4"] + + def test_effective_conversation_degrades_on_malformed(self) -> None: + from deepagents_code.app import _effective_conversation + + messages = ["m0", "m1"] + # No event, non-dict event, missing summary, and non-int cutoff all + # return the messages unchanged rather than raising or emitting a None. + assert _effective_conversation(messages, None) == messages + assert _effective_conversation(messages, "x") == messages + assert _effective_conversation(messages, {"cutoff_index": 1}) == messages + assert _effective_conversation(messages, {"summary_message": "S"}) == messages + + def test_effective_conversation_cutoff_past_end(self) -> None: + from deepagents_code.app import _effective_conversation + + event = {"summary_message": "S", "cutoff_index": 9} + assert _effective_conversation(["m0"], event) == ["S"] + + def test_message_text_handles_str_and_block_list(self) -> None: + from deepagents_code.app import _message_text + + assert _message_text(MagicMock(content="hello")) == "hello" + # A block-list content is concatenated, not stringified to "[{...}]". + blocks = [ + {"type": "text", "text": "Compaction "}, + {"type": "text", "text": "failed"}, + ] + assert _message_text({"content": blocks}) == "Compaction failed" + assert _message_text({"content": None}) == "" + + def test_find_compaction_failure_scans_durable_state(self) -> None: + from langchain_core.messages import HumanMessage, ToolMessage + + from deepagents_code.app import _find_compaction_failure + from deepagents_code.offload_middleware import COMPACTION_FAILURE_PREFIX + + failing = ToolMessage( + content=f"{COMPACTION_FAILURE_PREFIX}: boom", + tool_call_id="tc", + ) + messages = [HumanMessage("hi"), failing] + assert ( + _find_compaction_failure(messages) == f"{COMPACTION_FAILURE_PREFIX}: boom" + ) + + def test_find_compaction_failure_ignores_success(self) -> None: + from langchain_core.messages import ToolMessage + + from deepagents_code.app import _find_compaction_failure + + ok = ToolMessage(content="Conversation compacted.", tool_call_id="tc") + assert _find_compaction_failure([ok]) is None + # Serialized-dict tool message form is handled too. + assert _find_compaction_failure([{"type": "tool", "content": "ok"}]) is None + + +def _deny_dispatched_call( + reason: str | None, +) -> Callable[[Any, Any], dict[str, Any]]: + """Build an `aafter_model` stub that denies whichever call was dispatched. + + Keys the outcome on the tool-call id the node actually generated rather than + a fixed literal. The node derives that id per run (so a hook fulfillment + cannot be memoized across two `/offload`s in one turn), so a hardcoded key + here would silently stop matching and the denial would be read as "no + outcome" instead of failing loudly. + + Also asserts the dispatched call's `name`/`args`, which are the values + `ServerHooksMiddleware._after_model` gates on: without `compact_conversation` + it never raises `PreCompact` at all, and without `force: True` it raises the + event as `CompactTrigger.AUTO`, silently exempting `/offload` from a hook + scoped to manual compaction. Uses the middleware's own state key for the + same reason -- a re-spelling would make the node read `{}` and compact + straight through this denial. + + Args: + reason: Denial reason, or `None` to omit it. + + Returns: + A side-effect callable for an `AsyncMock`. + """ + from deepagents_code.hooks.server_middleware import _PRE_TOOL_STATE_KEY + + def deny(state: Any, _runtime: Any) -> dict[str, Any]: # noqa: ANN401 + call = state["messages"][0].tool_calls[0] + assert call["name"] == "compact_conversation" + assert call["args"] == {"force": True} + outcome: dict[str, Any] = {"behavior": "deny"} + if reason is not None: + outcome["reason"] = reason + return {_PRE_TOOL_STATE_KEY: {call["id"]: outcome}} + + return deny + + +class TestForcedCompactionGraph: + """Lifecycle and cost guarantees of the dedicated `/offload` graph.""" + + async def test_precompact_denial_skips_forced_compaction(self) -> None: + """A configured `PreCompact` hook can deny manual `/offload`. + + The denial has to raise rather than return an empty update: an empty + update is indistinguishable from "nothing old enough to compact", so the + client would render a hook veto as "the conversation is already compact" + and the reason would never reach the user. + """ + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph + + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock() + hooks = MagicMock() + hooks.aafter_model = AsyncMock( + side_effect=_deny_dispatched_call("policy forbids compaction") + ) + + graph = create_forced_compaction_graph(middleware, hooks_middleware=hooks) + with pytest.raises(RuntimeError, match="policy forbids compaction"): + await cast("Any", graph).ainvoke({"messages": []}, context=CLIContext()) + + hooks.aafter_model.assert_awaited_once() + middleware.arun_forced_compaction_update.assert_not_awaited() + + async def test_precompact_denial_without_reason_still_surfaces(self) -> None: + """A denial carrying no reason still reports something actionable.""" + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph + + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock() + hooks = MagicMock() + hooks.aafter_model = AsyncMock(side_effect=_deny_dispatched_call(None)) + + graph = create_forced_compaction_graph(middleware, hooks_middleware=hooks) + with pytest.raises(RuntimeError, match="Blocked by a compaction hook"): + await cast("Any", graph).ainvoke({"messages": []}, context=CLIContext()) + + middleware.arun_forced_compaction_update.assert_not_awaited() + + async def test_compaction_failure_is_raised_with_a_preserved_message(self) -> None: + """A node failure raises `RuntimeError` so its text survives the server. + + The LangGraph server preserves an exception's message only for an + allowlist of builtin types and replaces every other one with "An + internal error occurred", so an `OSError` from the archive write has to + be re-raised as an allowlisted type to stay diagnosable. + """ + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph + + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock( + side_effect=OSError("disk is full") + ) + + graph = create_forced_compaction_graph(middleware, hooks_middleware=None) + with pytest.raises(RuntimeError, match="OSError: disk is full") as exc_info: + await cast("Any", graph).ainvoke({"messages": []}, context=CLIContext()) + + assert "Your conversation is unchanged." in str(exc_info.value) + + async def test_failed_compaction_leaves_summary_spend_undrained( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A failure must not drain cost it is about to discard. + + The drain is destructive and the raise discards any update built on this + path, so draining here would lose the summarizer's spend outright. + Undrained records are charged on the next turn's first step instead. + """ + from deepagents_code import offload_middleware + from deepagents_code._cli_context import CLIContext + + cost_tracking = MagicMock() + cost_tracking.after_agent = MagicMock(return_value={"_session_cost_usd": 0.25}) + monkeypatch.setattr( + offload_middleware, "CostTrackingMiddleware", lambda: cost_tracking + ) + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock( + side_effect=OSError("disk is full") + ) + + graph = offload_middleware.create_forced_compaction_graph( + middleware, hooks_middleware=None + ) + with pytest.raises(RuntimeError): + await cast("Any", graph).ainvoke({"messages": []}, context=CLIContext()) + + cost_tracking.after_agent.assert_not_called() + + def test_only_messages_is_a_writable_input_channel(self) -> None: + """The run input surface is restricted by the graph, not by the client. + + `PrivateStateAttr` / `OmitFromInput` on the state schema are honored by + `create_agent`, *not* by a raw `StateGraph`, so without an explicit + `input_schema` every channel `_OffloadState` declares would be writable + by any local caller on the same `noop`-auth port (see THREAT_MODEL TB10): + `_summarization_event` would let them set the compaction cutoff, and + `_session_cost_usd` (an `operator.add` channel) would let them inflate + the thread's recorded spend. + """ + from deepagents_code.offload_middleware import create_forced_compaction_graph + + graph = create_forced_compaction_graph(MagicMock(), hooks_middleware=None) + + assert set(graph.get_input_jsonschema()["properties"]) == {"messages"} + + async def test_private_channel_in_input_is_dropped_not_applied(self) -> None: + """An injected private channel must not reach the node or the checkpoint. + + The schema check above is the mechanism; this is the behavior. A replayed + `_session_cost_usd` would *add* to the thread's total, so echoing the + checkpointed value back would double the recorded spend on every + `/offload` and compound across runs. + """ + from langchain_core.messages import HumanMessage + + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph + + seen: dict[str, Any] = {} + middleware = MagicMock() + + async def capture(state: Any, _runtime: Any) -> dict[str, Any]: # noqa: ANN401, RUF029 + seen.update(state) + return {} + + middleware.arun_forced_compaction_update = AsyncMock(side_effect=capture) + graph = create_forced_compaction_graph(middleware, hooks_middleware=None) + + out = await cast("Any", graph).ainvoke( + { + "messages": [HumanMessage("hi", id="m1")], + "_session_cost_usd": 0.42, + "_summarization_event": {"cutoff_index": 99}, + }, + context=CLIContext(), + ) + + assert seen.get("_session_cost_usd") == pytest.approx(0.0) + assert seen.get("_summarization_event") is None + assert [m.content for m in seen["messages"]] == ["hi"] + assert out.get("_session_cost_usd") == pytest.approx(0.0) + + async def test_forced_hook_call_id_is_unique_per_invocation(self) -> None: + """Two `/offload`s in one turn must not share a hook invocation id. + + `ServerHooksMiddleware` derives its `invocation_id` from this id plus the + thread, hook snapshot, and prompt id — and the prompt id only rotates on + user-prompt submit. A constant id would therefore collide across two + `/offload`s in one turn and replay the first decision, so a first-run + *denial* would silently deny the second run too. + """ + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph + + seen_ids: list[str] = [] + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock(return_value={}) + hooks = MagicMock() + + async def record(state: Any, _runtime: Any) -> dict[str, Any]: # noqa: ANN401, RUF029 + seen_ids.append(state["messages"][0].tool_calls[0]["id"]) + return {} + + hooks.aafter_model = AsyncMock(side_effect=record) + graph = create_forced_compaction_graph(middleware, hooks_middleware=hooks) + + await cast("Any", graph).ainvoke({"messages": []}, context=CLIContext()) + await cast("Any", graph).ainvoke({"messages": []}, context=CLIContext()) + + assert len(seen_ids) == 2 + assert seen_ids[0] != seen_ids[1] + + async def test_hook_interrupt_survives_a_real_resume_round_trip( + self, tmp_path: Path + ) -> None: + """A fulfilled hook interrupt must let the offload finish. + + The one test that exercises the whole loop against the real + `ServerHooksMiddleware` rather than a stub. Answering an `interrupt()` + re-executes the node **from the top**, so anything the node derives + before dispatching must survive that replay. It previously minted the + forced tool-call id with `uuid4()`; the middleware folds that id into + its hook `invocation_id`, so the resumed execution computed a different + one and `parse_hook_resume_value` rejected the client's answer as + fatal — making `/offload` fail outright for every user with a + `PreCompact`/`PreToolUse` hook configured, and rendering the client's + entire fulfill/resume loop unreachable. + + Mocking either side hides this: the bug lives in the interaction between + node replay and invocation-id derivation. + """ + from uuid import uuid4 + + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.types import Command + + from deepagents_code._cli_context import CLIContextSchema + from deepagents_code.hooks.interrupt import ( + build_hook_resume_value, + parse_hook_interrupt_payload, + ) + from deepagents_code.hooks.models.domain import HookEvent, PreCompactDecision + from deepagents_code.hooks.models.transport import HookInvocationResponse + from deepagents_code.hooks.server_middleware import ServerHooksMiddleware + from deepagents_code.offload_middleware import create_forced_compaction_graph + + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock( + return_value={"_summarization_event": {"cutoff_index": 4}} + ) + # The graph compiles without a checkpointer (the LangGraph server owns + # durability in production), but an interrupt cannot resume without one. + built = create_forced_compaction_graph( + middleware, hooks_middleware=ServerHooksMiddleware(cwd=tmp_path) + ) + graph = cast("Any", built).builder.compile(checkpointer=InMemorySaver()) + + context = CLIContextSchema( + hooks_snapshot_id="snap", + hooks_server_events=[HookEvent.PRE_COMPACT.value], + thread_id="t1", + ) + config = {"configurable": {"thread_id": str(uuid4())}} + + first = await graph.ainvoke( + {"messages": [HumanMessage("hi", id="m1")]}, + config=config, + context=context, + ) + interrupts = first["__interrupt__"] + assert len(interrupts) == 1 + request = parse_hook_interrupt_payload(interrupts[0].value) + assert request is not None + + resume_value = build_hook_resume_value( + HookInvocationResponse( + protocol_version=1, + invocation_id=request.invocation_id, + snapshot_id=request.snapshot_id, + decision=PreCompactDecision(event=HookEvent.PRE_COMPACT), + ) + ) + result = await graph.ainvoke( + Command(resume={interrupts[0].id: resume_value}), + config=config, + context=context, + ) + + # The resume was accepted and the compaction actually ran. + assert result["_summarization_event"]["cutoff_index"] == 4 + middleware.arun_forced_compaction_update.assert_awaited_once() + + async def test_hook_dispatch_failure_is_reported_as_a_hook_failure(self) -> None: + """A hook-layer crash must not reach the user as "internal error". + + The server replaces the message of any exception outside its builtin + allowlist, so this has to be re-raised as `RuntimeError` like the + compaction failure — but worded so it is not read as a compaction bug. + """ + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream - app = DeepAgentsApp() + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock() + hooks = MagicMock() + hooks.aafter_model = AsyncMock(side_effect=OSError("hook socket died")) + graph = create_forced_compaction_graph(middleware, hooks_middleware=hooks) - async with app.run_test() as pilot: - await pilot.pause() - app._agent = agent - app._lc_thread_id = "test-thread" - with ( - patch("deepagents_code.cost_tracking.estimate_cost", return_value=0.20), - pytest.raises(RuntimeError, match="stream failed"), - ): - await app._drive_server_side_compaction( - {"configurable": {"thread_id": "test-thread"}} - ) - await pilot.pause() + with pytest.raises(RuntimeError, match=r"Offload hooks failed.*socket died"): + await cast("Any", graph).ainvoke({"messages": []}, context=CLIContext()) - assert app._thread_stats.request_count == 1 - assert app._session_stats.request_count == 1 - assert app._thread_stats.per_kind["offload"].cost_usd == pytest.approx(0.20) - assert app._session_cost_usd == pytest.approx(0.0) - assert app._displayed_cost_usd == pytest.approx(0.20) - summary = app._format_cost_summary() - assert "Estimated thread cost: $0.20" in summary - assert "Offload: $0.20" in summary + middleware.arun_forced_compaction_update.assert_not_awaited() - async def test_fulfills_precompact_before_manual_approval(self) -> None: - """A precompact hook is fulfilled before the compaction approval.""" - from types import SimpleNamespace + async def test_hook_interrupt_bubbles_to_the_operation_graph(self) -> None: + """A hook approval pause must reach the server's interrupt stream.""" + from langgraph.errors import GraphInterrupt + from langgraph.types import Interrupt - from langchain_core.messages import ToolMessage - from langgraph.types import Command + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph - from deepagents_code.client.remote_client import RemoteAgent - from deepagents_code.hooks.interrupt import HOOK_INVOCATION_INTERRUPT_TYPE + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock() + hooks = MagicMock() + hooks.aafter_model = AsyncMock( + side_effect=GraphInterrupt((Interrupt(value={"type": "hook_invocation"}),)) + ) + graph = create_forced_compaction_graph(middleware, hooks_middleware=hooks) - streams: list[object] = [] + result = await cast("Any", graph).ainvoke( + {"messages": []}, context=CLIContext() + ) - async def _astream( # noqa: ANN202, RUF029 - value: object, **_kwargs: object - ): - index = len(streams) - streams.append(value) - if index == 0: - interrupt = SimpleNamespace( - id="hook-interrupt", - value={"type": HOOK_INVOCATION_INTERRUPT_TYPE}, - ) - elif index == 1: - interrupt = SimpleNamespace( - id="approval-interrupt", - value={ - "action_requests": [ - { - "name": "compact_conversation", - "args": {"force": True}, - } - ] - }, - ) - else: - yield ( - (), - "messages", - (ToolMessage(content="compacted", tool_call_id="compact-call"), {}), - ) - return - yield ((), "updates", {"__interrupt__": [interrupt]}) + assert result["__interrupt__"] + middleware.arun_forced_compaction_update.assert_not_awaited() - agent = MagicMock( - spec=RemoteAgent, - aensure_thread=AsyncMock(), - aupdate_state=AsyncMock(), - astream=_astream, + async def test_cost_drain_failure_does_not_discard_the_compaction(self) -> None: + """A bookkeeping failure must not throw away a committed archive write. + + By the time the drain runs, the archive section is already written. Raising + here would report "your conversation is unchanged" while leaving an + orphaned section no `_summarization_event` references. + """ + from deepagents_code import offload_middleware + from deepagents_code._cli_context import CLIContext + + event = {"cutoff_index": 3, "summary_message": None, "file_path": "/a.md"} + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock( + return_value={"_summarization_event": event} ) - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - runtime = MagicMock(snapshot_id="snapshot") - runtime.configured_server_events.return_value = ("PreCompact",) - assert app._session_state is not None - app._session_state.hooks = HooksManager.adopting( - runtime, - identity=app._session_state.hook_identity, + with patch.object( + offload_middleware.CostTrackingMiddleware, + "after_agent", + new=MagicMock(side_effect=RuntimeError("pricing down")), + ): + graph = offload_middleware.create_forced_compaction_graph( + middleware, hooks_middleware=None + ) + out = await cast("Any", graph).ainvoke( + {"messages": []}, context=CLIContext() ) - app._agent = agent - app._lc_thread_id = "test-thread" - fulfill = AsyncMock(return_value={"hook": "approved"}) - with patch("deepagents_code.hooks.client.fulfill_hook_interrupt", fulfill): - result = await app._drive_server_side_compaction( - {"configurable": {"thread_id": "test-thread"}} - ) - assert result is None - fulfill.assert_awaited_once() - assert len(streams) == 3 - assert isinstance(streams[1], Command) - assert streams[1].resume == {"hook-interrupt": {"hook": "approved"}} - assert isinstance(streams[2], Command) - approval = streams[2].resume - assert isinstance(approval, dict) - assert "approval-interrupt" in approval + assert out["_summarization_event"] == event - async def test_reports_tool_failure(self) -> None: - """Returns the tool's error text when compaction fails.""" - from deepagents_code.offload_middleware import COMPACTION_FAILURE_PREFIX + async def test_operation_graph_preserves_agent_only_channels(self) -> None: + """The narrow operation schema must not drop the agent's other state. - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - agent, _inputs, _contexts = self._fake_remote_agent( - f"{COMPACTION_FAILURE_PREFIX}: an error occurred during compaction." - ) - app._agent = agent - app._lc_thread_id = "test-thread" + Both graphs run against one thread, but `_OffloadState` declares only + the two channels the operation needs. A schema that dropped or replaced + the rest would silently destroy conversation state the agent owns, so + this pins that a real checkpoint round-trip leaves it intact. + """ + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.graph import END, START, StateGraph + + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "shared-thread"}} + + # Stand in for the agent graph: a wider schema on the same thread. + class _WideState(TypedDict, total=False): + messages: Annotated[list, add_messages] + todos: list[str] + _summarization_event: dict + + def _seed(state: _WideState) -> dict: # noqa: ARG001 # node input unused + return { + "messages": [HumanMessage(content="secret history", id="m0")], + "todos": ["keep me"], + } - config = {"configurable": {"thread_id": "test-thread"}} - result = await app._drive_server_side_compaction(config) # ty: ignore - await pilot.pause() + wide = StateGraph(cast("Any", _WideState)) + wide.add_node("seed", cast("Any", _seed)) + wide.add_edge(START, "seed") + wide.add_edge("seed", END) + await cast("Any", wide.compile(checkpointer=checkpointer)).ainvoke({}, config) - assert result is not None - assert result.startswith(COMPACTION_FAILURE_PREFIX) + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock( + return_value={"_summarization_event": {"cutoff_index": 1}} + ) + offload = create_forced_compaction_graph(middleware, hooks_middleware=None) + # `create_forced_compaction_graph` compiles without a checkpointer + # because the LangGraph server attaches its own to every registered + # graph. Do the same here so both graphs share one thread. + cast("Any", offload).checkpointer = checkpointer + await cast("Any", offload).ainvoke({}, config, context=CLIContext()) + + # Read back through the wide schema, as the agent graph would. + state = await cast("Any", wide.compile(checkpointer=checkpointer)).aget_state( + config + ) + assert state.values["todos"] == ["keep me"] + assert [message.id for message in state.values["messages"]] == ["m0"] + assert state.values["_summarization_event"] == {"cutoff_index": 1} + + async def test_checkpointed_summarization_event_reaches_the_node(self) -> None: + """The prior event must arrive from the checkpoint, not from input. + + `_OffloadInput` deliberately keeps `_summarization_event` out of the run + input, and a separate test pins that. This is the other half: the node + still has to *read* it from `_OffloadState`. Narrowing that schema — an + easy-looking cleanup, given `_OffloadInput` sits right beside it — would + hand the node `event=None`, so a second `/offload` would re-summarize + already-archived messages and compute its cutoff from the wrong base. + Both graph-level tests would still pass. + """ + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.graph import END, START, StateGraph - async def test_forwards_startup_model_profile_to_compaction(self) -> None: - """Profile data is usable even without a session `/model` override.""" - from deepagents_code.config import settings + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - agent, _inputs, contexts = self._fake_remote_agent( - "Conversation compacted. Summarized 2 messages." - ) - app._agent = agent - app._lc_thread_id = "test-thread" - app._model_override = None - app._profile_override = {"max_input_tokens": 4096} + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "resumed-thread"}} + prior_event = {"cutoff_index": 7, "summary_message": None, "file_path": "/a.md"} - config = {"configurable": {"thread_id": "test-thread"}} - with ( - patch.object(settings, "model_provider", "provider"), - patch.object(settings, "model_name", "startup-model"), - patch.object(settings, "model_context_limit", 4096), - ): - await app._drive_server_side_compaction(config) # ty: ignore - await pilot.pause() + class _WideState(TypedDict, total=False): + messages: Annotated[list, add_messages] + _summarization_event: dict - assert contexts - seed_values = agent.aupdate_state.call_args.args[1] - (seed_msg,) = seed_values["messages"] - (tool_call,) = seed_msg.tool_calls - expected = { - "model": "provider:startup-model", - "model_params": {}, - "profile_overrides": {"max_input_tokens": 4096}, - "model_context_limit": 4096, - "thread_id": "test-thread", - "offload_tool_call_id": tool_call["id"], - } - for context in contexts: - assert isinstance(context, dict) - normalized = {str(key): value for key, value in context.items()} - assert {key: normalized[key] for key in expected} == expected + def _seed(state: _WideState) -> dict: # noqa: ARG001 # node input unused + return { + "messages": [HumanMessage(content="older", id="m0")], + "_summarization_event": prior_event, + } - async def test_rejects_interrupt_without_identifiable_action(self) -> None: - """Malformed interrupt payloads fail closed instead of being approved.""" - from langgraph.types import Command + wide = StateGraph(cast("Any", _WideState)) + wide.add_node("seed", cast("Any", _seed)) + wide.add_edge(START, "seed") + wide.add_edge("seed", END) + await cast("Any", wide.compile(checkpointer=checkpointer)).ainvoke({}, config) - from deepagents_code.client.remote_client import RemoteAgent + seen: dict[str, Any] = {} - astream_inputs: list[Any] = [] + async def capture(state: Any, _runtime: Any) -> dict[str, Any]: # noqa: ANN401, RUF029 + seen.update(state) + return {} - class _Interrupt: - id = "interrupt-unknown" - value: dict[str, Any] = {} # noqa: RUF012 # test stub + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock(side_effect=capture) + offload = create_forced_compaction_graph(middleware, hooks_middleware=None) + cast("Any", offload).checkpointer = checkpointer + await cast("Any", offload).ainvoke( + {"messages": []}, config, context=CLIContext() + ) - async def _astream( # noqa: RUF029, ANN202 - stream_input: object, **_kwargs: object - ): - astream_inputs.append(stream_input) - if stream_input is None: - yield ((), "updates", {"__interrupt__": [_Interrupt()]}) + assert seen["_summarization_event"] == prior_event - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream + async def test_replayed_input_reaches_the_node_without_duplicating(self) -> None: + """The driver's state replay must arrive intact and not double up. - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - app._agent = agent - app._lc_thread_id = "test-thread" + `_drive_offload_operation_graph` replays the thread's messages as the run + input because an empty input leaves a *server-backed* run with nothing to + compact. The `add_messages` reducer is what makes that safe: replaying + messages that are already checkpointed merges them by ID rather than + appending a second copy. + """ + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.graph import END, START, StateGraph + + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "shared-thread"}} + seeded = [ + HumanMessage(content="history", id="m0"), + HumanMessage(content="more", id="m1"), + ] - config = {"configurable": {"thread_id": "test-thread"}} - result = await app._drive_server_side_compaction(config) # ty: ignore - await pilot.pause() + class _WideState(TypedDict, total=False): + messages: Annotated[list, add_messages] - assert result is None - assert len(astream_inputs) == 2 - assert isinstance(astream_inputs[1], Command) - decision = astream_inputs[1].resume["interrupt-unknown"]["decisions"][0] - assert decision["type"] == "reject" + def _seed(state: _WideState) -> dict: # noqa: ARG001 # node input unused + return {"messages": seeded} - async def test_approves_only_first_forced_compaction(self) -> None: - """A repeated forced compaction request is rejected, not approved.""" - from langchain_core.messages import ToolMessage - from langgraph.types import Command + wide = StateGraph(cast("Any", _WideState)) + wide.add_node("seed", cast("Any", _seed)) + wide.add_edge(START, "seed") + wide.add_edge("seed", END) + await cast("Any", wide.compile(checkpointer=checkpointer)).ainvoke({}, config) - from deepagents_code.client.remote_client import RemoteAgent + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock(return_value=None) + offload = create_forced_compaction_graph(middleware, hooks_middleware=None) + cast("Any", offload).checkpointer = checkpointer + await cast("Any", offload).ainvoke( + {"messages": seeded}, config, context=CLIContext() + ) - astream_inputs: list[Any] = [] - guard_ids: list[object] = [] + await_args = middleware.arun_forced_compaction_update.await_args + assert await_args is not None + state_arg = await_args.args[0] + assert [message.id for message in state_arg["messages"]] == ["m0", "m1"] - class _Interrupt: - def __init__(self, iid: str, tool_name: str, args: dict[str, Any]) -> None: - self.id = iid - self.value = {"action_requests": [{"name": tool_name, "args": args}]} + async def test_nothing_to_compact_returns_an_empty_update(self) -> None: + """A `None` update must not become a partial write.""" + from deepagents_code._cli_context import CLIContext + from deepagents_code.offload_middleware import create_forced_compaction_graph - async def _astream(stream_input: object, **kwargs: object): # noqa: RUF029, ANN202 - idx = len(astream_inputs) - astream_inputs.append(stream_input) - context = kwargs.get("context") - guard_ids.append( - context.get("offload_tool_call_id") - if isinstance(context, dict) - else None - ) - if idx == 0: - compact = _Interrupt( - "i-compact", "compact_conversation", {"force": True} - ) - yield ((), "updates", {"__interrupt__": [compact]}) - elif idx == 1: - # Model a trailing turn that asks to compact again. - repeated = _Interrupt( - "i-repeated", "compact_conversation", {"force": True} - ) - yield ((), "updates", {"__interrupt__": [repeated]}) - else: - yield ( - (), - "messages", - ( - ToolMessage( - content="Conversation compacted. Summarized 2 messages " - "into a concise summary.", - tool_call_id="x", - ), - {}, - ), - ) + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock(return_value=None) - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream + graph = create_forced_compaction_graph(middleware, hooks_middleware=None) + result = await cast("Any", graph).ainvoke( + {"messages": []}, context=CLIContext() + ) + + assert "_summarization_event" not in result + + async def test_summary_cost_is_drained_into_graph_update( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The graph checkpoints summarizer spend before `/offload` returns. + + The stub implements only `after_agent`, exactly like the real + `CostTrackingMiddleware`, and inherits nothing. Awaiting `aafter_agent` + instead therefore raises `AttributeError` into the drain's `except` and + the spend silently vanishes — which is what a `MagicMock` stub hid: + `AgentMiddleware.aafter_agent` exists as an empty base method, so + `await ...aafter_agent(...)` returns `None` in production while a mock + materializes a working one in the test. + """ + from deepagents_code import offload_middleware + from deepagents_code._cli_context import CLIContext + + calls: list[str] = [] - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - app._agent = agent - app._lc_thread_id = "test-thread" + class _OnlySyncDrain: + """Mirrors `CostTrackingMiddleware`'s real method surface.""" - config = {"configurable": {"thread_id": "test-thread"}} - result = await app._drive_server_side_compaction(config) # ty: ignore - await pilot.pause() + def after_agent(self, _state: object, _runtime: object) -> dict[str, float]: + calls.append("after_agent") + return {"_session_cost_usd": 0.25} - assert result is None - # Initial drain + two resumes (compaction, then trailing tool). - assert len(astream_inputs) == 3 - assert isinstance(astream_inputs[1], Command) - assert isinstance(astream_inputs[2], Command) - assert len(set(guard_ids)) == 1 - assert isinstance(guard_ids[0], str) - # Compaction was approved. - compact_decision = astream_inputs[1].resume["i-compact"]["decisions"][0] - assert compact_decision["type"] == "approve" - # A second compaction request is not the seeded call and is rejected. - repeated_decision = astream_inputs[2].resume["i-repeated"]["decisions"][0] - assert repeated_decision["type"] == "reject" + monkeypatch.setattr( + offload_middleware, "CostTrackingMiddleware", _OnlySyncDrain + ) + middleware = MagicMock() + middleware.arun_forced_compaction_update = AsyncMock( + return_value={"_summarization_event": {"cutoff_index": 2}} + ) - async def test_sets_tool_guard_context_without_hitl(self) -> None: - """The per-run tool guard is set even when no HITL interrupt exists.""" - from langchain_core.messages import ToolMessage + graph = offload_middleware.create_forced_compaction_graph( + middleware, hooks_middleware=None + ) + result = await cast("Any", graph).ainvoke( + {"messages": []}, context=CLIContext() + ) - from deepagents_code.client.remote_client import RemoteAgent + assert calls == ["after_agent"] + assert result["_session_cost_usd"] == pytest.approx(0.25) - guard_ids: list[object] = [] - async def _astream(_stream_input: object, **kwargs: object): # noqa: RUF029, ANN202 - context = kwargs.get("context") - guard_ids.append( - context.get("offload_tool_call_id") - if isinstance(context, dict) - else None - ) - yield ( - (), - "messages", - ( - ToolMessage( - content="Conversation compacted. Summarized 2 messages.", - tool_call_id="x", - ), - {}, +class TestOffloadDriverSelection: + """`_handle_offload` must pick its driver by agent kind, not by accident. + + Every other test in this file asserts routing *indirectly*, by patching the + driver it expects and letting an unpatched one blow up against a mock. That + fails loudly but misleadingly. These pin the predicate itself. + """ + + async def test_server_agent_uses_only_the_operation_graph(self) -> None: + """A `RemoteAgent` must never take the seeded in-process path.""" + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + with ( + patch.object( + app, + "_get_thread_state_values", + AsyncMock(return_value=_state_values([_make_dict_message("hi")])), ), - ) + patch.object( + app, "_drive_offload_operation_graph", AsyncMock(return_value=None) + ) as operation, + patch.object( + app, "_drive_local_seeded_compaction", AsyncMock(return_value=None) + ) as seeded, + ): + await app._handle_offload() - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream + operation.assert_awaited_once() + seeded.assert_not_awaited() + + async def test_custom_server_graph_falls_back_to_seeded_compaction(self) -> None: + """A custom graph without `offload` retains the pre-operation behavior.""" + from deepagents_code.app import _MissingOffloadGraphError app = DeepAgentsApp() + before = _state_values(_make_dict_messages(2)) + after = _state_values(_make_dict_messages(2), _summary_event(1)) async with app.run_test() as pilot: await pilot.pause() - app._agent = agent - app._lc_thread_id = "test-thread" + _setup_server_offload_app(app) + with ( + patch.object( + app, + "_get_thread_state_values", + AsyncMock(side_effect=[before, after]), + ), + patch.object( + app, + "_drive_offload_operation_graph", + AsyncMock(side_effect=_MissingOffloadGraphError()), + ) as operation, + patch.object( + app, "_drive_local_seeded_compaction", AsyncMock(return_value=None) + ) as seeded, + ): + await app._handle_offload() - config = {"configurable": {"thread_id": "test-thread"}} - result = await app._drive_server_side_compaction(config) # ty: ignore + operation.assert_awaited_once() + seeded.assert_awaited_once() + + async def test_local_agent_uses_only_the_seeded_driver(self) -> None: + """A local in-process `Pregel` agent must never stream the named graph.""" + app = DeepAgentsApp() + async with app.run_test() as pilot: await pilot.pause() + _setup_local_offload_app(app) + with ( + patch.object( + app, + "_get_thread_state_values", + AsyncMock(return_value=_state_values([_make_dict_message("hi")])), + ), + patch.object( + app, "_drive_offload_operation_graph", AsyncMock(return_value=None) + ) as operation, + patch.object( + app, "_drive_local_seeded_compaction", AsyncMock(return_value=None) + ) as seeded, + ): + await app._handle_offload() - assert result is None - seed_values = agent.aupdate_state.call_args.args[1] - (seed_msg,) = seed_values["messages"] - (tool_call,) = seed_msg.tool_calls - assert guard_ids == [tool_call["id"]] + seeded.assert_awaited_once() + operation.assert_not_awaited() - async def test_bounds_resume_loop_and_reports_abandoned_drain(self) -> None: - """A model that keeps requesting tools cannot spin `/offload` forever. - Every stream yields a fresh gated interrupt, so the resume loop never - drains cleanly. It must stop at the `max_resume_rounds` cap (initial - drain + 10 resumes = 11 streams) and surface a user-visible notice that - the run was left paused, rather than looping indefinitely. - """ - from deepagents_code.client.remote_client import RemoteAgent +class TestOffloadReplaySafety: + """The run input replaces the `messages` channel, so it must be current. - astream_inputs: list[Any] = [] + Against a real LangGraph server the `/offload` run input is authoritative + for `messages`: streaming `{"messages": []}` empties an eight-message thread + (see `test_offload_server_side.py`). A stale or partial replay is therefore + not a stale read but a destructive write. + """ - class _Interrupt: - def __init__(self, iid: str) -> None: - self.id = iid - self.value = {"action_requests": [{"name": "write_file", "args": {}}]} + @staticmethod + def _remote_with_stream(chunks: list[Any]) -> MagicMock: + """Build a remote whose offload graph yields `chunks`.""" + operation = MagicMock() + + async def stream(*_args: object, **_kwargs: object): # noqa: ANN202, RUF029 + for chunk in chunks: + yield chunk + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + return remote + + async def test_replays_freshly_read_state_not_the_callers_snapshot(self) -> None: + """A turn committed after the caller's snapshot must not be dropped. + + `_handle_offload` reads state *before* `_set_agent_running(True)`, so a + run committed in that window is missing from its snapshot. Replaying + that snapshot would write the shorter list over the live conversation + and delete the newer turn. + """ + app = DeepAgentsApp() + stale = [_make_dict_message("one")] + fresh = [_make_dict_message("one"), _make_dict_message("two")] + captured: list[object] = [] - async def _astream(stream_input: object, **_kwargs: object): # noqa: RUF029, ANN202 - idx = len(astream_inputs) - astream_inputs.append(stream_input) - # Never terminate: each round surfaces another gated interrupt. - yield ((), "updates", {"__interrupt__": [_Interrupt(f"i-{idx}")]}) + operation = MagicMock() - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream + async def stream(*args: object, **_kwargs: object): # noqa: ANN202, RUF029 + captured.extend(args) + yield (), "updates", {"force_compact": {}} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation - app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - app._agent = agent + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace(values=_state_values(fresh)) + ) app._lc_thread_id = "test-thread" + with patch.object(app, "_remote_agent", return_value=remote): + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + _state_values(stale), + ) - config = {"configurable": {"thread_id": "test-thread"}} - result = await app._drive_server_side_compaction(config) # ty: ignore - await pilot.pause() + assert captured == [{"messages": fresh}] - # No compaction failure was reported, so the run returns cleanly. - assert result is None - # Initial drain + exactly 10 resume rounds, then the cap breaks. - assert len(astream_inputs) == 11 - assert any( - "could not be fully drained" in str(widget._content) - for widget in app.query(ErrorMessage) + async def test_empty_replay_is_refused_rather_than_streamed(self) -> None: + """An unreadable state must abort, not wipe the conversation. + + If both the re-read and the caller's snapshot yield no messages, running + anyway would truncate the thread to zero and report "already compact". + """ + app = DeepAgentsApp() + remote = self._remote_with_stream([((), "updates", {"force_compact": {}})]) + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace(values={"messages": []}) ) + app._lc_thread_id = "test-thread" + with ( + patch.object(app, "_remote_agent", return_value=remote), + pytest.raises(RuntimeError, match="could not be read back"), + ): + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, {"messages": []} + ) + remote.for_graph.assert_not_called() -class TestRemoveUnansweredOffloadSeed: - """Cleanup of a committed-but-unanswered `/offload` seed after a failure.""" + +class TestOffloadDrainCompletion: + """A stream that never reports the node's update is not a success.""" @staticmethod - def _seed_message(tool_call_id: str) -> dict[str, Any]: - """Serialized seed AIMessage carrying the forced compaction tool call.""" - return { - "type": "ai", - "content": "", - "id": f"offload-seed-{tool_call_id}", - "tool_calls": [ - { - "name": "compact_conversation", - "args": {"force": True}, - "id": tool_call_id, - } - ], - } + def _app_with_chunks(app: DeepAgentsApp, chunks: list[Any]) -> MagicMock: + """Point `app` at a remote whose offload stream yields `chunks`.""" + operation = MagicMock() + + async def stream(*_args: object, **_kwargs: object): # noqa: ANN202, RUF029 + for chunk in chunks: + yield chunk + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + return remote - async def test_removes_dangling_seed(self) -> None: - """An unanswered seed is removed so it cannot wedge the next turn.""" - from langchain_core.messages import RemoveMessage + async def test_missing_node_update_is_reported_as_a_failure(self) -> None: + """Chunk-shape drift must not read as "already compact". + Interrupt detection is all the drain loop does, so a chunk shape it + cannot parse is indistinguishable from a clean run: no interrupts, no + error, and a run still paused server-side. The caller would then read an + unadvanced event and tell the user their conversation is already + compact. + """ app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - agent = MagicMock() - agent.aupdate_state = AsyncMock() - app._agent = agent - state = _state_values( - [*_make_dict_messages(2), self._seed_message("seed-call")] - ) - with patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - return_value=state, - ): - await app._remove_unanswered_offload_seed( - {"configurable": {"thread_id": "test-thread"}}, "seed-call" + # 2-tuples: the shape the loop's `len(chunk) != 3` filter discards. + remote = self._app_with_chunks(app, [("updates", {"force_compact": {}})]) + with patch.object(app, "_remote_agent", return_value=remote): + result = await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + _state_values([_make_dict_message("hi")]), ) - agent.aupdate_state.assert_awaited_once() - update = agent.aupdate_state.call_args.args[1] - (removal,) = update["messages"] - assert isinstance(removal, RemoveMessage) - assert removal.id == "offload-seed-seed-call" + assert result is not None + assert "without reporting a result" in result + assert "conversation is unchanged" in result - async def test_keeps_answered_seed(self) -> None: - """A seed answered by a ToolMessage is a valid pair and is left intact.""" + async def test_node_update_marks_the_run_complete(self) -> None: + """The positive control: a real node update reports success.""" app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - agent = MagicMock() - agent.aupdate_state = AsyncMock() - app._agent = agent - answered = { - "type": "tool", - "content": "Nothing to compact yet.", - "tool_call_id": "seed-call", - } - state = _state_values( - [*_make_dict_messages(2), self._seed_message("seed-call"), answered] + remote = self._app_with_chunks( + app, [((), "updates", {"force_compact": {}})] ) - with patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - return_value=state, - ): - await app._remove_unanswered_offload_seed( - {"configurable": {"thread_id": "test-thread"}}, "seed-call" + with patch.object(app, "_remote_agent", return_value=remote): + result = await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + _state_values([_make_dict_message("hi")]), ) - agent.aupdate_state.assert_not_awaited() + assert result is None + + async def test_paused_run_keeps_its_offload_graph_binding(self) -> None: + """A drain failure leaves the run suspended, so do not rebind the thread. + + Rebinding would re-point the thread away from the graph the paused run + belongs to, leaving it unaddressable. + """ + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + remote = self._app_with_chunks(app, [("updates", {})]) + with patch.object(app, "_remote_agent", return_value=remote): + result = await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + _state_values([_make_dict_message("hi")]), + ) + + assert result is not None + remote.arebind_thread.assert_not_awaited() + + +class TestOffloadHookStopReporting: + """A hook that stops the client mid-`/offload` must still say something.""" + + async def test_operation_path_reports_a_hook_stop(self) -> None: + """The operation graph's fulfillment mounts nothing on its own.""" + from deepagents_code.hooks.client_lifecycle import ClientHookStopError - async def test_noop_when_seed_absent(self) -> None: - """Nothing is removed when no seed with the id is present.""" app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() _setup_server_offload_app(app) - agent = MagicMock() - agent.aupdate_state = AsyncMock() - app._agent = agent - state = _state_values(_make_dict_messages(2)) - with patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - return_value=state, + with ( + patch.object( + app, + "_get_thread_state_values", + AsyncMock(return_value=_state_values([_make_dict_message("hi")])), + ), + patch.object( + app, + "_drive_offload_operation_graph", + AsyncMock(side_effect=ClientHookStopError("stopped")), + ), ): - await app._remove_unanswered_offload_seed( - {"configurable": {"thread_id": "test-thread"}}, "seed-call" - ) + await app._handle_offload() + await pilot.pause() + + errors = [str(w._content) for w in app.query(ErrorMessage)] + + assert any("Offload stopped by a hook." in text for text in errors) + + async def test_seeded_path_stays_silent(self) -> None: + """The seeded driver mounts its own stop reason before raising. - agent.aupdate_state.assert_not_awaited() + Mounting a second, generic line here would duplicate it. + """ + from deepagents_code.hooks.client_lifecycle import ClientHookStopError - async def test_returns_true_when_seed_removed(self) -> None: - """Successful removal reports the thread is clean.""" app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - agent = MagicMock() - agent.aupdate_state = AsyncMock() - app._agent = agent - state = _state_values( - [*_make_dict_messages(2), self._seed_message("seed-call")] - ) - with patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - return_value=state, + _setup_local_offload_app(app) + with ( + patch.object( + app, + "_get_thread_state_values", + AsyncMock(return_value=_state_values([_make_dict_message("hi")])), + ), + patch.object( + app, + "_drive_local_seeded_compaction", + AsyncMock(side_effect=ClientHookStopError("stopped")), + ), ): - cleaned = await app._remove_unanswered_offload_seed( - {"configurable": {"thread_id": "test-thread"}}, "seed-call" - ) + await app._handle_offload() + await pilot.pause() - assert cleaned is True + errors = [str(w._content) for w in app.query(ErrorMessage)] - async def test_returns_false_when_state_read_fails(self) -> None: - """A failed state read cannot confirm cleanup, so it reports unclean.""" + assert not any("Offload stopped by a hook." in text for text in errors) + + +class TestOffloadRebindWarningCoverage: + """Every path that reports a finished offload must surface a failed rebind. + + A dropped warning leaves an unrelated later `/goal` or `/rubric` to fail + with an opaque `as_node="model"` error and nothing tying it to the offload. + """ + + async def test_unconfirmed_result_still_warns(self) -> None: + """The empty-state branch says "Offload finished" — so it must warn.""" app = DeepAgentsApp() + + async def _fail_rebind(*_args: object, **_kwargs: object) -> None: # noqa: RUF029 + # `_handle_offload` clears the flag before dispatching, so the + # driver has to be the one that sets it -- as the real one does in + # its `finally`. + app._offload_rebind_failed = True + async with app.run_test() as pilot: await pilot.pause() _setup_server_offload_app(app) - agent = MagicMock() - agent.aupdate_state = AsyncMock() - app._agent = agent - with patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=RuntimeError("state read boom"), + states = [_state_values([_make_dict_message("hi")]), {}] + with ( + patch.object( + app, + "_get_thread_state_values", + AsyncMock(side_effect=states), + ), + patch.object( + app, + "_drive_offload_operation_graph", + AsyncMock(side_effect=_fail_rebind), + ), ): - cleaned = await app._remove_unanswered_offload_seed( - {"configurable": {"thread_id": "test-thread"}}, "seed-call" - ) + await app._handle_offload() + await pilot.pause() - assert cleaned is False - # The dangling seed could not be removed, so nothing was written. - agent.aupdate_state.assert_not_awaited() + errors = [str(w._content) for w in app.query(ErrorMessage)] + + assert any("could not be confirmed" in text for text in errors) + assert any("re-associated with the main agent" in text for text in errors) + + async def test_hook_stop_still_warns(self) -> None: + """A hook stop after a failed rebind must not swallow the warning.""" + from deepagents_code.hooks.client_lifecycle import ClientHookStopError - async def test_returns_false_when_removal_write_fails(self) -> None: - """A failed removal write leaves the seed and reports unclean.""" app = DeepAgentsApp() + + async def _stop(*_args: object, **_kwargs: object) -> None: # noqa: RUF029 + app._offload_rebind_failed = True + msg = "stopped" + raise ClientHookStopError(msg) + async with app.run_test() as pilot: await pilot.pause() _setup_server_offload_app(app) - agent = MagicMock() - agent.aupdate_state = AsyncMock(side_effect=RuntimeError("write boom")) - app._agent = agent - state = _state_values( - [*_make_dict_messages(2), self._seed_message("seed-call")] - ) - with patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - return_value=state, + with ( + patch.object( + app, + "_get_thread_state_values", + AsyncMock(return_value=_state_values([_make_dict_message("hi")])), + ), + patch.object( + app, "_drive_offload_operation_graph", AsyncMock(side_effect=_stop) + ), ): - cleaned = await app._remove_unanswered_offload_seed( - {"configurable": {"thread_id": "test-thread"}}, "seed-call" - ) + await app._handle_offload() + await pilot.pause() - assert cleaned is False + errors = [str(w._content) for w in app.query(ErrorMessage)] + assert any("re-associated with the main agent" in text for text in errors) -class TestFormatTokenCount: - """Test the format_token_count helper function.""" - def test_zero(self) -> None: - assert format_token_count(0) == "0" +class TestOffloadResourceWiring: + """`/offload` must archive into the agent's own backend.""" - def test_below_threshold(self) -> None: - assert format_token_count(999) == "999" + def test_mismatched_backend_is_rejected_at_construction(self) -> None: + """A compaction bound elsewhere would archive where the agent cannot read. - def test_at_threshold(self) -> None: - assert format_token_count(1000) == "1.0K" + The symptom — history that silently is not there — surfaces long after + the mis-wiring, so fail where the pair is built instead. + """ + from deepagents_code.offload_middleware import ( + OffloadServerResources, + attach_offload_resources, + ) - def test_above_threshold(self) -> None: - assert format_token_count(1500) == "1.5K" + backend = SimpleNamespace() + compaction = MagicMock() + compaction._summarization._backend = SimpleNamespace() - def test_large_value(self) -> None: - assert format_token_count(200000) == "200.0K" + with pytest.raises(ValueError, match="different backend"): + attach_offload_resources( + cast("Any", backend), + OffloadServerResources( + compaction=cast("Any", compaction), hooks=cast("Any", MagicMock()) + ), + ) - def test_millions(self) -> None: - assert format_token_count(1_000_000) == "1.0M" + def test_matching_backend_is_published(self) -> None: + """The positive control for the check above.""" + from deepagents_code.offload_middleware import ( + OffloadServerResources, + attach_offload_resources, + offload_resources_from, + ) - def test_above_million(self) -> None: - assert format_token_count(2_500_000) == "2.5M" + backend = SimpleNamespace() + compaction = MagicMock() + compaction._summarization._backend = backend + hooks = MagicMock() + attach_offload_resources( + cast("Any", backend), + OffloadServerResources( + compaction=cast("Any", compaction), hooks=cast("Any", hooks) + ), + ) -class TestOffloadHelpers: - """Pure helpers backing `/offload` accounting and failure detection.""" + resources = offload_resources_from(cast("Any", backend)) + assert resources is not None + assert resources.compaction is compaction + assert resources.hooks is hooks - def test_summarization_cutoff_reads_int(self) -> None: - from deepagents_code.app import _summarization_cutoff - assert _summarization_cutoff({"cutoff_index": 4}) == 4 +class TestForcedOffloadCallId: + """The hook dispatch's call id must be stable across a run's resumes.""" - def test_summarization_cutoff_defaults_zero_on_malformed(self) -> None: - from deepagents_code.app import _summarization_cutoff + def test_missing_checkpoint_namespace_is_logged_not_silent( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A run without a usable `checkpoint_ns` breaks hook resumes. - assert _summarization_cutoff(None) == 0 - assert _summarization_cutoff({"cutoff_index": "x"}) == 0 - assert _summarization_cutoff({}) == 0 - assert _summarization_cutoff("not-a-dict") == 0 + The random fallback makes the id differ between the request and the + resume, which `parse_hook_resume_value` rejects as fatal — so `/offload` + dies with "the client answered a different request", but only for users + with hooks configured. Without a log line there is nothing to point at. + """ + from deepagents_code import offload_middleware + + with ( + patch.object( + offload_middleware, + "get_config", + return_value={"configurable": {}}, + ), + caplog.at_level("WARNING"), + ): + call_id = offload_middleware._forced_offload_call_id() - def test_effective_conversation_applies_event(self) -> None: - from deepagents_code.app import _effective_conversation + assert call_id.startswith("offload-precompact-") + assert "checkpoint_ns" in caplog.text - messages = [f"m{i}" for i in range(5)] - event = {"summary_message": "S", "cutoff_index": 2} - assert _effective_conversation(messages, event) == ["S", "m2", "m3", "m4"] + def test_no_runnable_context_is_not_warned_about( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A direct call outside a graph is expected, not a misconfiguration. - def test_effective_conversation_degrades_on_malformed(self) -> None: - from deepagents_code.app import _effective_conversation + Nothing can interrupt or resume such a call, so the random id is + correct there and must not be reported as a problem. + """ + from deepagents_code import offload_middleware - messages = ["m0", "m1"] - # No event, non-dict event, missing summary, and non-int cutoff all - # return the messages unchanged rather than raising or emitting a None. - assert _effective_conversation(messages, None) == messages - assert _effective_conversation(messages, "x") == messages - assert _effective_conversation(messages, {"cutoff_index": 1}) == messages - assert _effective_conversation(messages, {"summary_message": "S"}) == messages + with ( + patch.object( + offload_middleware, "get_config", side_effect=RuntimeError("no context") + ), + caplog.at_level("WARNING"), + ): + call_id = offload_middleware._forced_offload_call_id() - def test_effective_conversation_cutoff_past_end(self) -> None: - from deepagents_code.app import _effective_conversation + assert call_id.startswith("offload-precompact-") + assert "checkpoint_ns" not in caplog.text - event = {"summary_message": "S", "cutoff_index": 9} - assert _effective_conversation(["m0"], event) == ["S"] + def test_same_namespace_yields_the_same_id(self) -> None: + """Answering a hook interrupt replays the node from the top.""" + from deepagents_code import offload_middleware - def test_message_text_handles_str_and_block_list(self) -> None: - from deepagents_code.app import _message_text + config = {"configurable": {"checkpoint_ns": "force_compact:abc123"}} + with patch.object(offload_middleware, "get_config", return_value=config): + first = offload_middleware._forced_offload_call_id() + second = offload_middleware._forced_offload_call_id() - assert _message_text(MagicMock(content="hello")) == "hello" - # A block-list content is concatenated, not stringified to "[{...}]". - blocks = [ - {"type": "text", "text": "Compaction "}, - {"type": "text", "text": "failed"}, - ] - assert _message_text({"content": blocks}) == "Compaction failed" - assert _message_text({"content": None}) == "" + assert first == second - def test_find_compaction_failure_scans_durable_state(self) -> None: - from langchain_core.messages import HumanMessage, ToolMessage - from deepagents_code.app import _find_compaction_failure - from deepagents_code.offload_middleware import COMPACTION_FAILURE_PREFIX +class TestOffloadStreamShape: + """The operation driver must request the chunk shape it parses.""" - failing = ToolMessage( - content=f"{COMPACTION_FAILURE_PREFIX}: boom", - tool_call_id="tc", - ) - messages = [HumanMessage("hi"), failing] - assert ( - _find_compaction_failure(messages) == f"{COMPACTION_FAILURE_PREFIX}: boom" - ) + async def test_streams_with_subgraphs_enabled(self) -> None: + """`subgraphs=True` is what makes the yielded chunk a real 3-tuple. - def test_find_compaction_failure_ignores_success(self) -> None: + `RemoteAgent.astream` only yields the documented + `(namespace, mode, data)` shape when `subgraphs` is set. Without it a + live server yields `("updates", {...}, None)`, so the driver's unpacking + binds `mode` to the payload dict, every chunk falls through the + `mode != "updates"` filter, and interrupts are dropped in silence — the + run stays paused while the user is told their conversation is already + compact. Unit tests feed the loop chunks directly and cannot see that, + so pin the request instead. + """ + app = DeepAgentsApp() + stream_kwargs: dict[str, object] = {} + operation = MagicMock() + + async def stream(*_args: object, **kwargs: object): # noqa: ANN202, RUF029 + stream_kwargs.update(kwargs) + yield (), "updates", {"force_compact": {}} + + operation.astream = stream + remote = MagicMock() + remote.aensure_thread = AsyncMock() + remote.arebind_thread = AsyncMock() + remote.for_graph.return_value = operation + + async with app.run_test() as pilot: + await pilot.pause() + app._agent = MagicMock() + app._agent.aget_state = AsyncMock( + return_value=SimpleNamespace( + values=_state_values([_make_dict_message("hi")]) + ) + ) + app._lc_thread_id = "test-thread" + with patch.object(app, "_remote_agent", return_value=remote): + await app._drive_offload_operation_graph( + {"configurable": {"thread_id": "test-thread"}}, + _state_values([_make_dict_message("hi")]), + ) + + assert stream_kwargs["subgraphs"] is True + + +class TestSeededDriverAgainstALocalAgent: + """The seeded driver's *only* production shape is a local `Pregel` agent. + + Every other test in `TestDriveLocalSeededCompaction` builds its agent with + `MagicMock(spec=RemoteAgent)`, but `_handle_offload` now routes remote + agents to the operation graph — so the driver is exercised exclusively in + the shape it no longer serves. These use a non-`RemoteAgent` double. + """ + + @staticmethod + def _local_agent(tool_content: str) -> tuple[MagicMock, dict[str, object]]: + """Build a local (non-`RemoteAgent`) agent double. + + Returns: + The agent and a dict recording the kwargs it was streamed with. + """ from langchain_core.messages import ToolMessage - from deepagents_code.app import _find_compaction_failure + stream_kwargs: dict[str, object] = {} - ok = ToolMessage(content="Conversation compacted.", tool_call_id="tc") - assert _find_compaction_failure([ok]) is None - # Serialized-dict tool message form is handled too. - assert _find_compaction_failure([{"type": "tool", "content": "ok"}]) is None + async def _astream(*_args: object, **kwargs: object): # noqa: ANN202, RUF029 + stream_kwargs.update(kwargs) + yield ( + (), + "messages", + (ToolMessage(content=tool_content, tool_call_id="x"), {}), + ) + + agent = MagicMock() + agent.aupdate_state = AsyncMock() + agent.astream = _astream + return agent, stream_kwargs + + async def test_local_agent_is_driven_without_thread_registration(self) -> None: + """A local agent has no server thread to register. + + `aensure_thread` exists only on `RemoteAgent`; calling it here would + raise, and the tool result must still be detected from the stream. + """ + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + agent, stream_kwargs = self._local_agent( + "Conversation compacted. Summarized 2 messages into a summary." + ) + app._agent = agent + app._lc_thread_id = "test-thread" + + result = await app._drive_local_seeded_compaction( # ty: ignore + {"configurable": {"thread_id": "test-thread"}} + ) + + assert result is None + assert not hasattr(agent.aensure_thread, "assert_awaited") # not a RemoteAgent + # The driver parses `(namespace, mode, data)`, so it must ask for it. + assert stream_kwargs["subgraphs"] is True + + async def test_local_agent_failure_is_detected_from_the_tool_message(self) -> None: + """The driver's only failure signal is the `ToolMessage` text.""" + from deepagents_code.offload_middleware import COMPACTION_FAILURE_PREFIX + + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + agent, _ = self._local_agent( + f"{COMPACTION_FAILURE_PREFIX}: OSError: disk is full." + ) + app._agent = agent + app._lc_thread_id = "test-thread" + + result = await app._drive_local_seeded_compaction( # ty: ignore + {"configurable": {"thread_id": "test-thread"}} + ) + + assert result is not None + assert "disk is full" in result + + +class TestOffloadSessionCostSync: + """Operation-graph spend reaches `/cost` only via the post-run state read.""" + + async def test_summary_spend_is_synced_from_the_committed_state(self) -> None: + """The summarizer runs outside the agent loop, so nothing else reports it. + + `stream_mode=["updates"]` gives this path no per-kind itemization, so + `_sync_session_cost_from_state` on the state read back after the run is + the only route from the graph's `CostTrackingMiddleware` to the session + total. Dropping or reordering that call silently under-reports every + server-side `/offload` until the next turn. + """ + app = DeepAgentsApp() + after = _state_values([_make_dict_message("hi")], _summary_event(2)) | { + "_session_cost_usd": 1.25 + } + + async with app.run_test() as pilot: + await pilot.pause() + _setup_server_offload_app(app) + with ( + patch.object( + app, + "_get_thread_state_values", + AsyncMock( + side_effect=[ + _state_values([_make_dict_message("hi")]), + after, + ] + ), + ), + patch.object( + app, "_drive_offload_operation_graph", AsyncMock(return_value=None) + ), + patch.object(app, "_sync_session_cost_from_state") as sync, + patch.object(app, "_run_session_start_hook", AsyncMock()), + ): + await app._handle_offload() + await pilot.pause() + + sync.assert_called_once_with(after) diff --git a/libs/code/tests/unit_tests/test_remote_client.py b/libs/code/tests/unit_tests/test_remote_client.py index f7a211302e..458a3dcde4 100644 --- a/libs/code/tests/unit_tests/test_remote_client.py +++ b/libs/code/tests/unit_tests/test_remote_client.py @@ -1074,3 +1074,143 @@ def test_non_string_error_key_uses_class_name(self) -> None: def test_non_dict_payload_uses_class_name(self) -> None: assert agent_error_type(ValueError("boom")) == "ValueError" + + +# --------------------------------------------------------------------------- +# RemoteAgent.for_graph +# --------------------------------------------------------------------------- + + +class TestForGraph: + """Sibling clients for other graphs served by the same runtime.""" + + def _agent(self) -> RemoteAgent: + return RemoteAgent( + "http://localhost:1234", + graph_name="agent", + api_key="secret-key", + headers={"x-proxy": "yes"}, + ) + + def test_preserves_url_and_credentials(self) -> None: + """A dropped key would 401 only on authenticated deployments.""" + sibling = self._agent().for_graph("offload") + + assert sibling._graph_name == "offload" + assert sibling._url == "http://localhost:1234" + assert sibling._api_key == "secret-key" + assert sibling._headers == {"x-proxy": "yes"} + + def test_repeated_calls_return_one_client(self) -> None: + """Each `RemoteGraph` opens an httpx async *and* sync client. + + Neither is pooled by the SDK nor closed here, so building a client per + call would leak two connection pools on every `/offload`. + """ + agent = self._agent() + + assert agent.for_graph("offload") is agent.for_graph("offload") + + def test_own_graph_name_returns_self(self) -> None: + """Asking for the graph this client already serves must not build a second. + + `test_offload_server_side.py` takes exactly this path + (`agent.for_graph("agent")`), and without the short-circuit it creates + the duplicate httpx sync+async pair the cache exists to prevent. + """ + agent = self._agent() + + assert agent.for_graph("agent") is agent + assert agent._sibling_clients == {} + + def test_distinct_names_get_distinct_clients(self) -> None: + agent = self._agent() + + assert agent.for_graph("offload") is not agent.for_graph("agent") + + def test_sibling_does_not_share_the_parents_graph(self) -> None: + """The cached `RemoteGraph` is per-graph-name, not shared.""" + agent = self._agent() + with patch("langgraph.pregel.remote.RemoteGraph") as remote_graph: + remote_graph.side_effect = lambda name, **_kwargs: SimpleNamespace( + name=name + ) + parent_graph = agent._get_graph() + sibling_graph = agent.for_graph("offload")._get_graph() + + assert parent_graph is not sibling_graph + assert parent_graph.name == "agent" + assert sibling_graph.name == "offload" + + +# --------------------------------------------------------------------------- +# RemoteAgent.arebind_thread +# --------------------------------------------------------------------------- + + +class TestArebindThread: + """Re-associating a thread with this client's graph after a sibling run.""" + + def _agent(self) -> RemoteAgent: + return RemoteAgent("http://localhost:1234", graph_name="agent") + + async def test_writes_graph_id_into_thread_metadata(self) -> None: + """The rebind must go through thread *metadata*, not a state update. + + The server resolves an out-of-run `as_node=` write from + `thread["metadata"]["graph_id"]`. `threads.update_state` sends no graph + id at all, so an empty state update resolves against whichever graph ran + last -- i.e. it cannot undo a sibling-graph run. Only a metadata write + rebinds the thread. + """ + agent = self._agent() + client = MagicMock() + client.threads.update = AsyncMock() + + graph = SimpleNamespace(_validate_client=lambda: client) + with patch.object(agent, "_get_graph", return_value=graph): + await agent.arebind_thread({"configurable": {"thread_id": "t-1"}}) + + client.threads.update.assert_awaited_once_with( + "t-1", metadata={"graph_id": "agent"} + ) + + async def test_missing_thread_id_raises(self) -> None: + """A silent no-op would leave the thread bound to the sibling graph.""" + agent = self._agent() + + with pytest.raises(ValueError, match="thread_id"): + await agent.arebind_thread({"configurable": {}}) + + async def test_missing_sdk_accessor_reports_the_real_cause(self) -> None: + """An SDK rename must not degrade into a generic rebind warning. + + `_validate_client` is private. If it disappears, the bare attribute + access would raise inside the caller's blanket rebind handler, which + downgrades everything to a warning -- `/offload` would keep "working" + while every later `/goal` and `/rubric` failed on a mis-bound thread, + permanently and with nothing connecting the two. + """ + agent = self._agent() + + with ( + patch.object(agent, "_get_graph", return_value=SimpleNamespace()), + pytest.raises(AttributeError, match="_validate_client is unavailable"), + ): + await agent.arebind_thread({"configurable": {"thread_id": "t-1"}}) + + async def test_transport_failure_propagates(self) -> None: + """The caller decides whether a failed rebind is fatal, so re-raise.""" + agent = self._agent() + client = MagicMock() + client.threads.update = AsyncMock(side_effect=RuntimeError("server gone")) + + with ( + patch.object( + agent, + "_get_graph", + return_value=SimpleNamespace(_validate_client=lambda: client), + ), + pytest.raises(RuntimeError, match="server gone"), + ): + await agent.arebind_thread({"configurable": {"thread_id": "t-1"}}) diff --git a/libs/code/tests/unit_tests/test_server_graph.py b/libs/code/tests/unit_tests/test_server_graph.py index 413543aa8d..faecf35960 100644 --- a/libs/code/tests/unit_tests/test_server_graph.py +++ b/libs/code/tests/unit_tests/test_server_graph.py @@ -7,6 +7,7 @@ import sys import threading from types import ModuleType, SimpleNamespace +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -21,6 +22,38 @@ def _import_fresh_server_graph() -> ModuleType: return importlib.import_module("deepagents_code.server_graph") +def _attach_offload_resources( + backend: object, + *, + compaction: object | None = None, + hooks: object | None = None, +) -> tuple[object, object]: + """Publish offload resources whose compaction is bound to `backend`. + + `attach_offload_resources` rejects a compaction middleware bound to a + different backend, since that would make `/offload` archive into storage the + agent cannot read. A bare `MagicMock()` auto-creates a distinct + `_summarization._backend`, so wire it explicitly rather than defeating the + check with a looser assertion. + + Returns: + The compaction and hooks doubles that were published. + """ + from deepagents_code.offload_middleware import ( + OffloadServerResources, + attach_offload_resources, + ) + + compaction = cast("Any", compaction if compaction is not None else MagicMock()) + hooks = cast("Any", hooks if hooks is not None else MagicMock()) + compaction._summarization._backend = backend + attach_offload_resources( + cast("Any", backend), + OffloadServerResources(compaction=compaction, hooks=hooks), + ) + return compaction, hooks + + def _module_with_attrs(name: str, **attrs: object) -> ModuleType: """Create a module stub with dynamically assigned attributes.""" module = ModuleType(name) @@ -38,13 +71,179 @@ async def test_make_graph_caches_first_constructed_graph(self) -> None: module = _import_fresh_server_graph() with patch.object( - module, "_make_graph", new=AsyncMock(return_value=graph_obj) + module, + "_make_graphs", + new=AsyncMock(return_value=module.ServerRuntime(graph_obj, object())), ) as make_graph: assert await module.make_graph() is graph_obj assert await module.make_graph() is graph_obj make_graph.assert_awaited_once_with() + async def test_both_graphs_share_one_resource_initialization(self) -> None: + """`agent` and `offload` must not build two server runtimes. + + The LangGraph server resolves named graphs independently, so a second + initialization would re-discover MCP servers, create a second sandbox, + and stack duplicate `atexit` handlers. + """ + graph_obj = object() + module = _import_fresh_server_graph() + backend = SimpleNamespace() + offload_graph = object() + + _attach_offload_resources(backend) + + with ( + patch.object( + module, + "_make_graphs", + new=AsyncMock(return_value=module.ServerRuntime(graph_obj, backend)), + ) as make_graphs, + patch( + "deepagents_code.offload_middleware.create_forced_compaction_graph", + return_value=offload_graph, + ), + ): + assert await module.make_graph() is graph_obj + assert await module.make_offload_graph() is offload_graph + assert await module.make_offload_graph() is offload_graph + assert await module.make_graph() is graph_obj + + make_graphs.assert_awaited_once_with() + + async def test_offload_graph_reuses_the_agents_middleware(self) -> None: + """The operation graph must run the agent's own compaction middleware.""" + module = _import_fresh_server_graph() + backend = SimpleNamespace() + compaction, hooks = _attach_offload_resources(backend) + + with ( + patch.object( + module, + "_make_graphs", + new=AsyncMock(return_value=module.ServerRuntime(object(), backend)), + ), + patch( + "deepagents_code.offload_middleware.create_forced_compaction_graph", + return_value=object(), + ) as create_graph, + ): + await module.make_offload_graph() + + create_graph.assert_called_once_with(compaction, hooks_middleware=hooks) + + def test_factories_are_addressed_by_name(self) -> None: + """The two graph factories must not be a bare positional pair. + + Both are zero-arg async callables returning `Any`, so a transposition + type-checks — and because `generate_langgraph_json` derives both refs + from one module, the server would start cleanly with the offload graph + registered as `agent`. The first user message would then run a graph + with no model node. + """ + module = _import_fresh_server_graph() + + factories = module._build_graph_factories() + + assert factories._fields == ("agent", "offload") + # The closures are named after the graph they build, so this also pins + # that the module-level bindings did not get crossed. + assert factories.agent.__name__ == "make_graph" + assert factories.offload.__name__ == "make_offload_graph" + assert module.make_graph.__name__ == "make_graph" + assert module.make_offload_graph.__name__ == "make_offload_graph" + + def test_resources_of_the_wrong_type_are_rejected(self) -> None: + """A foreign attribute value must read as absent, not be trusted. + + `offload_resources_from` narrows with `isinstance` so a backend carrying + something else under the same attribute name fails closed with the + caller's own message instead of an `AttributeError` deep inside the + graph build. + """ + from deepagents_code.offload_middleware import ( + _OFFLOAD_RESOURCES_ATTR, + offload_resources_from, + ) + + backend = SimpleNamespace() + setattr(backend, _OFFLOAD_RESOURCES_ATTR, ("compaction", "hooks")) + + assert offload_resources_from(cast("Any", backend)) is None + + async def test_offload_graph_fails_closed_without_published_middleware( + self, + ) -> None: + """A backend carrying no offload resources must fail, not half-work.""" + module = _import_fresh_server_graph() + + with ( + patch.object( + module, + "_make_graphs", + new=AsyncMock( + return_value=module.ServerRuntime(object(), SimpleNamespace()) + ), + ), + pytest.raises(RuntimeError, match="did not publish its offload"), + ): + await module.make_offload_graph() + + async def test_concurrent_resolution_builds_one_runtime(self) -> None: + """The server resolves named graphs concurrently, not in sequence. + + The sequential test above would pass even with a broken double-check; + `gather` is what the server actually does when both graphs are requested + at once, and a lost race here means two sandboxes and two MCP discoveries. + """ + import asyncio + + module = _import_fresh_server_graph() + graph_obj = object() + backend = SimpleNamespace() + + _attach_offload_resources(backend) + + calls = 0 + + async def build() -> object: + nonlocal calls + calls += 1 + # Yield so a second waiter can enter before the cache is populated. + await asyncio.sleep(0) + return module.ServerRuntime(graph_obj, backend) + + offload_graph = object() + with ( + patch.object(module, "_make_graphs", new=build), + patch( + "deepagents_code.offload_middleware.create_forced_compaction_graph", + return_value=offload_graph, + ) as create_graph, + ): + results = await asyncio.gather( + module.make_graph(), + module.make_offload_graph(), + module.make_offload_graph(), + module.make_graph(), + ) + + assert calls == 1 + assert results == [graph_obj, offload_graph, offload_graph, graph_obj] + create_graph.assert_called_once() + + def test_server_runtime_slots_are_named(self) -> None: + """Both slots are opaque to the type checker, so name them. + + A positional transposition would hand LangGraph the backend as its + compiled graph and make `/offload` report "no implementation" — pointing + at the wrong subsystem entirely. + """ + module = _import_fresh_server_graph() + + assert module.ServerRuntime._fields == ("agent", "backend") + def test_criteria_context_tools_use_identity_allowlist_in_tool_order(self) -> None: """Criteria tools should be known context objects in main-tool order.""" module = _import_fresh_server_graph() @@ -120,7 +319,7 @@ async def test_make_graph_emits_marker_and_exits_on_failure( with ( patch.object( module, - "_make_graph", + "_make_graphs", new=AsyncMock(side_effect=ValueError("boom: bad model")), ), pytest.raises(SystemExit) as exc_info, @@ -225,7 +424,7 @@ async def cleanup(self) -> None: # Non-default allowlist so the `fs_tools=` assertion below is # load-bearing: it round-trips through `to_env()`/`from_env()` and # must reach `create_cli_agent`. With the `None` default this - # assertion passed whether or not `_make_graph` read + # assertion passed whether or not `_make_graphs` read # `config.allow_fs_tools`, so a dropped read would go unnoticed. allow_fs_tools=["ls", "read_file"], ) diff --git a/libs/code/tests/unit_tests/test_server_manager.py b/libs/code/tests/unit_tests/test_server_manager.py index 974e9609aa..884666ab3c 100644 --- a/libs/code/tests/unit_tests/test_server_manager.py +++ b/libs/code/tests/unit_tests/test_server_manager.py @@ -550,6 +550,39 @@ def test_relative_paths_written_verbatim_to_langgraph_json( assert config["graphs"]["agent"] == "./server_graph.py:make_graph" assert config["checkpointer"]["path"] == "./checkpointer.py:create_checkpointer" + def test_offload_operation_graph_is_registered(self, tmp_path: Path) -> None: + """`/offload` streams this graph by name, so it must always be served. + + Dropping the entry turns every `/offload` into a server 404 that no + other unit test would catch. + """ + import json + + from deepagents_code.client.launch.server import generate_langgraph_json + + # The production default (see `server_manager`) resolves to the real + # installed module. + generate_langgraph_json(tmp_path) + config = json.loads((tmp_path / "langgraph.json").read_text()) + assert config["graphs"]["agent"] == "deepagents_code.server_graph:make_graph" + assert ( + config["graphs"]["offload"] + == "deepagents_code.server_graph:make_offload_graph" + ) + + def test_custom_graph_does_not_require_an_offload_factory( + self, tmp_path: Path + ) -> None: + """Custom graph references remain valid without an undocumented pair.""" + import json + + from deepagents_code.client.launch.server import generate_langgraph_json + + generate_langgraph_json(tmp_path, graph_ref="custom_graph:make_graph") + + config = json.loads((tmp_path / "langgraph.json").read_text()) + assert config["graphs"] == {"agent": "custom_graph:make_graph"} + class TestWritePyproject: """Tests for the generated runtime pyproject."""