diff --git a/libs/code/THREAT_MODEL.md b/libs/code/THREAT_MODEL.md index ac4c040959..5c252b80a7 100644 --- a/libs/code/THREAT_MODEL.md +++ b/libs/code/THREAT_MODEL.md @@ -85,7 +85,13 @@ │ └─►C12: RemoteAgent────┘ (HTTP+SSE on 127.0.0.1) │ │ (client/remote_client.py) │ │ │ │ -│ C3: Agent Engine (server_graph.py) │ +│ ┌──────────┴───────────┐ │ +│ ▼ ▼ │ +│ C18: Offload HTTP Boundary C3: Agent Engine │ +│ (custom route + operation) (server_graph.py) │ +│ │ │ │ +│ └──────────┬───────────┘ │ +│ │ │ │ (create_cli_agent, deepagents SDK) │ │ │ │ │ User Prompt ──────────────┘ │ @@ -141,7 +147,8 @@ | C15 | LocalContext Middleware | Runs a bash detection script via backend; injects git/project/env context into system prompt each turn | framework-controlled | Yes⁶ | `local_context.LocalContextMiddleware.before_agent`, `local_context.build_detect_script` | | C16 | Custom Subagent Loader | Reads `{dir}/{name}/AGENTS.md` YAML frontmatter from `.deepagents/agents/` and project `.agents/` directories | user-controlled | No | `subagents.list_subagents`, `subagents._parse_subagent_file` | | C17 | Model Config Loader | Resolves model providers, enforces the `models.allowed` policy (exact specs and `provider:*` wildcards), and supports `class_path` for arbitrary `BaseChatModel` instantiation via `importlib` | administrator/user-controlled | N/A | `config.create_model`, `config._create_model_from_class`, `model_config.ModelConfig.load` | -| C18 | Goal/Rubric State Notice | Projects persisted goal objectives, active criteria, and status notes into synthetic messages for the primary model | framework-controlled | Yes | `goal_state_notice.build_goal_state_notice`, `goal_tools.GoalToolsMiddleware` | +| C18 | Server Offload Boundary | Custom HTTP route registered with LangGraph's route-auth layer (inert under the shipped `noop` auth, which relies on the loopback bind); reads thread state, runs the agent's shared compaction/hooks/backend, and commits a state-only result plus cost | framework-controlled | Yes for built-in graph⁷ | `offload_api.offload`, `offload_api._execute_offload`, `offload_middleware.OffloadOperation.execute` | +| C19 | Goal/Rubric State Notice | Projects persisted goal objectives, active criteria, and status notes into synthetic messages for the primary model | framework-controlled | Yes | `goal_state_notice.build_goal_state_notice`, `goal_tools.GoalToolsMiddleware` | **Notes:** 1. `http_request` and `fetch_url` enabled by default; `web_search` requires `TAVILY_API_KEY`. @@ -150,6 +157,7 @@ 4. Sandbox mode requires explicit `--sandbox` CLI flag. 5. Both TUI and non-interactive modes now always spawn a local LangGraph dev server and connect via `RemoteAgent`. 6. `LocalContextMiddleware` is added whenever `LocalShellBackend` or an `_AsyncExecutableBackend` is in use (`agent.py:create_cli_agent`). +7. Custom graph references do not receive dcode's HTTP app and do not support `/offload`. --- @@ -160,7 +168,7 @@ | DC1 | API Keys / Credentials | `ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `TAVILY_API_KEY`, `LANGSMITH_API_KEY`, `LANGGRAPH_API_KEY` | Critical | Process environment only; never written to disk by CLI code | N/A (in-memory) | Process lifetime | All — breach trigger | | DC2 | Conversation Messages | User prompts, LLM responses, tool args/results, goal objectives, rubric criteria, and status notes | High | SQLite (`~/.deepagents/*.db`) via LangGraph checkpointer | No (local file, unencrypted) | Unbounded (session files persist) | GDPR if personal data is discussed | | DC3 | System Prompt Content | `DA_SERVER_SYSTEM_PROMPT` env var; custom AGENTS.md contents | Medium | Process environment (transient); `~/.deepagents/{agent}/AGENTS.md` on disk | No | Config lifetime | None direct | -| DC5 | Offloaded Conversation History | Summarized + raw conversation messages written to sandbox backend | High | Sandbox filesystem at `/conversation_history/{thread_id}.md` | Depends on sandbox provider | Sandbox session lifetime | GDPR if personal data is discussed | +| DC5 | Offloaded Conversation History | Summarized + raw conversation messages written to sandbox backend | High | Sandbox filesystem at `/conversation_history/session_{uuid4hex}.md` | Depends on sandbox provider | Sandbox session lifetime | GDPR if personal data is discussed | ### Data Classification Details @@ -186,13 +194,18 @@ #### DC5: Offloaded Conversation History -- **Fields**: Timestamped, formatted conversation messages written by `offload.offload_messages_to_backend`. -- **Storage**: Sandbox backend filesystem at path `/conversation_history/{thread_id}.md` where `thread_id` is a UUID7 (via `sessions.generate_thread_id`). +- **Fields**: Timestamped, formatted conversation messages written by the SDK's `SummarizationMiddleware._aoffload_to_backend`, reached through `offload_middleware.CLICompactionMiddleware`. +- **Producers**: Three paths write this data. + - Automatic trigger-based compaction. + - The model-initiated `compact_conversation` tool (HITL-gated, see TB2). + - The explicit `/offload` command. This one is available only through C18 on a built-in server, which reads checkpoint state and writes the archive without entering the tool-approval path. + - **Read guard**: The server-owned `/offload` path wraps 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 site. +- **Storage**: Sandbox backend filesystem at path `/conversation_history/session_{uuid4hex}.md`. The leaf is the *summarization session* id (`SummarizationMiddleware._get_history_path`), not the thread id: it is minted per summarization session and persisted under `_summarization_session_id` so later compactions append to the same file. One thread can therefore own several archives. - **Access**: Accessible within the sandbox session; depends on provider access controls. - **Encryption**: Depends on sandbox provider storage backend. - **Retention**: Sandbox session lifetime (destroyed when sandbox is deleted). - **Logging exposure**: Contains full message history including tool results. -- **Gaps**: Thread ID is UUID7 (no path injection risk), but offloaded content is unstructured markdown containing raw conversation data. +- **Gaps**: The filename is a framework-minted `session_` with no user-controlled component (no path injection risk), but offloaded content is unstructured markdown containing raw conversation data. --- @@ -228,6 +241,13 @@ - **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. + - **Authorization**: The explicit `/offload` command does *not* cross it. C18 invokes the agent's shared compaction service directly, with no tool node and no synthetic message. The slash command is the authorization. + - **Hook events**: The operation still dispatches `PreCompact` and `PreToolUse` against an in-memory forced call. Hooks may veto or interrupt. The TUI returns opaque hook replies over the operation protocol. + - **`ask` is fail-closed**: A `PreToolUse` `ask` decision cannot prompt on this path. The operation transport carries hook invocations, not HITL review requests, and `interrupt()` requires a Pregel task. `_ask_permission_via_hitl` therefore converts `ask` into a deny that carries the reason. + - **Archive write**: It reaches `backend.awrite()` without traversing tool approval. See DF25 and DF26. + - **Unsupported deployments**: Local in-process `Pregel` agents, including ACP mode, do not support `/offload`. Custom and older servers without C18 fail at the HTTP boundary rather than entering a client-driven tool path. +- **Key note**: Only the *pre* hook events fire for `/offload` through C18. `PostToolUse`/`PostToolUseFailure`, which `ServerHooksMiddleware` records in `awrap_tool_call` and dispatches from `_before_model` via `_maybe_post_tool_use` on the next model turn, do not fire: there is no tool node to record the pending entry, and no following model turn to drain it. 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 @@ -257,7 +277,13 @@ - **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`. +- **Crossing mechanism**: HTTP POST/GET to `http://127.0.0.1:{port}` using `langgraph.pregel.remote.RemoteGraph` for graph operations and the same configured HTTP client for C18. +- **Key note**: Graph registration and the accepted payload are both narrowed here. + - **Registration**: The default built-in `graph_ref` registers one `agent` graph plus the C18 custom HTTP app. A custom `graph_ref` registers only its graph and does not support `/offload`. + - **Accepted**: Operation identity, model/hook context, and opaque hook replies. + - **Rejected**: Messages, checkpoint identifiers, graph names, and state updates. The server refuses any operation update whose channels fall outside `OffloadStateUpdate`. + - **Checkpoint invariant**: The server reads and hydrates checkpoint messages itself and rejects active, pending, or changed threads. Its final thread-state update is state-only and targets the latest checkpoint, so it cannot branch from a stale checkpoint and hide concurrently appended messages. + - **Enforced by**: `offload_api._execute_offload` and `client.remote_client.RemoteAgent.aoffload`. #### TB11: Config File → Code Execution @@ -340,10 +366,12 @@ | DF22 | C14 Async Config | C3 Agent | AsyncSubAgent specs (URL, graph_id, headers) from config.toml | — | None | TOML parse + dict | | DF23 | C9 Config | C17 Model Config | `class_path` string from `config.toml` | — | TB11 | TOML parse → importlib | | DF24 | C5 MCP Config | MCP Subprocess | `env` dict from `.mcp.json` forwarded to stdio subprocess | DC1 | TB4 | subprocess environment | -| DF25 | C3 Agent | C7 Sandbox | Conversation messages for offload | DC5 | TB6 | `backend.awrite()` | -| DF26 | User / Host FS | C18 Goal/Rubric State Notice | Goal objective, criteria, and status notes; `/rubric file` content | DC2 | TB1, TB12 | TUI command + local file read + checkpoint update | -| DF27 | C18 Goal/Rubric State Notice | External LLM | Synthetic user-role message containing actionable objective, active criteria, and status note | DC2 | TB12, TB7 | LangChain model request over configured provider transport | -| DF28 | Administrator | C9 Config | Managed TOML policy | DC1 | TB13 | Fixed local file read | +| DF25 | C3 Agent, C18 Server Offload Boundary | C7 Sandbox | Conversation messages for offload | DC5 | TB6 | `backend.awrite()` | +| DF26 | C12 RemoteAgent | C18 Server Offload Boundary | Thread ID, operation identity, model/hook context, opaque hook replies; typed result or hook request | — | TB10 | HTTP+JSON (localhost) | +| DF27 | C18 Server Offload Boundary | C8 Sessions | Checkpoint message read; summarization event and additive cost update (never a messages write) | DC2 | None | In-process LangGraph SDK | +| DF28 | User / Host FS | C19 Goal/Rubric State Notice | Goal objective, criteria, and status notes; `/rubric file` content | DC2 | TB1, TB12 | TUI command + local file read + checkpoint update | +| DF29 | C19 Goal/Rubric State Notice | External LLM | Synthetic user-role message containing actionable objective, active criteria, and status note | DC2 | TB12, TB7 | LangChain model request over configured provider transport | +| DF30 | Administrator | C9 Config | Managed TOML policy | DC1 | TB13 | Fixed local file read | ### Flow Details @@ -383,7 +411,7 @@ - **Validation**: Type check only — `env` must be a dict (`mcp_tools._validate_server_config`). No filtering of key names or values. Forwarded directly to `StdioConnection` which passes to `subprocess.Popen`. - **Trust assumption**: User authored or approved the MCP config. Project-level configs go through the approval gate (allow-list or interactive prompt) before loading. -#### DF26/DF27: Goal/Rubric State → Primary-Model Context +#### DF28/DF29: Goal/Rubric State → Primary-Model Context - **Data**: User-entered goal objectives and rubric criteria, full text loaded through `/rubric file`, and agent-written completion or blocker notes. These are persisted in checkpoint state and embedded in a synthetic `HumanMessage` whenever the current notice must be restored or re-pinned. - **Validation**: Direct, file-loaded, generated, and tool-authored goal-state paths enforce raw-character limits: 8,000 for an objective, 12,000 for a rubric, 12,000 for an accepted objective and criteria combined, 4,000 for a status note or prior blocker, and 16,000 across a notice. `goal_state_notice._embedded_text` then escapes `<`, `>`, and `&` to prevent boundary-tag forgery; lifecycle projection suppresses a paused or complete goal's objective. The scoped code has no content-safety, secret-detection, byte, token, or post-escape rendered-size limit. @@ -400,7 +428,7 @@ | T3 | DF7 | — | Unicode-homoglyph URL in LLM-generated tool args deceives user during approval | TB2 | Low | Disproven | `unicode_security.check_url_safety`, `agent._format_fetch_url_description` | | T4 | DF5, DF9 | — | Auto-approve mode bypasses all HITL gates; any LLM-initiated tool call executes | TB2 | Low | Verified | `agent.create_cli_agent` (`auto_approve` param), `agent._add_interrupt_on` | | T5 | DF13, DF14| DC2 | Local SQLite checkpoint file tampered with to inject adversarial content into future LLM context | None | Low | Unverified | `sessions.get_db_path` | -| T6 | DF3, DF4 | DC2 | Unauthenticated LangGraph dev server on localhost can be accessed by any local process | TB10 | Medium | Verified | `server._build_server_env`, `server._DEFAULT_HOST` | +| T6 | DF3, DF4, DF26 | DC2 | Unauthenticated LangGraph dev server on localhost can be accessed by any local process | TB10 | Medium | Verified | `server._build_server_env`, `server._DEFAULT_HOST`, `offload_api.app` | | T7 | DF19, DF20| DC3 | Makefile or project file content injected into system prompt via LocalContextMiddleware | TB9 | Low | Verified | `local_context._section_makefile`, `local_context.LocalContextMiddleware._get_modified_request` | | T8 | DF21 | DC3 | Custom subagent AGENTS.md body used verbatim as system_prompt without content validation | None | Low | Verified | `subagents._parse_subagent_file`, `agent.create_cli_agent` | | T9 | DF23 | — | `class_path` in config.toml triggers arbitrary Python code execution via `importlib.import_module()` | TB11 | Low | Verified | `config._create_model_from_class`, `model_config.ProviderConfig` | @@ -408,9 +436,9 @@ | T12 | DF10 | — | Project `.env` sets shell startup-hook variables (`BASH_ENV`, `ENV`) that run attacker-controlled scripts when `dcode` spawns Bash, before any HITL approval | TB11 | High | Verified | `config._load_dotenv`, `local_context.build_detect_script` | | T13 | DF7 | — | Configured shell allow-list checks only the first token, so an allow-listed interpreter/wrapper (`python3`, `bash`, `env`, `xargs`, …) runs arbitrary code via its arguments without approval in non-interactive mode | TB2 | Medium | Verified | `config.is_shell_command_allowed`, `config.contains_dangerous_patterns` | | T14 | DF7, DF9 | — | A weaker model configured for the Auto approval classifier reviews gated actions less reliably, including untrusted text carried in tool arguments and file content | TB2 | Low | Verified | `auto_mode.AutoModeHITLMiddleware._classifier_model`, `config.resolve_auto_classifier_model`, `config_manifest.resolve_auto_classifier_timeout` | -| T15 | DF26, DF27 | DC2 | Stored prompt injection through a goal, rubric, or status note influences later primary-model tool requests | TB12 | Medium | Likely | `goal_state_notice.build_goal_state_notice`, `goal_tools.GoalToolsMiddleware._request_with_goal_notice` | -| T16 | DF26, DF27 | DC2 | Sensitive local-file content, up to the 12,000-character rubric limit, is automatically persisted and transmitted to the configured model provider as rubric criteria | TB12 | Medium | Verified | `app.DeepAgentsApp._set_rubric_from_file`, `goal_state_notice.build_goal_state_notice` | -| T17 | DF26, DF27 | DC2 | Character-bounded goal/rubric/status-note text can still exceed provider context budgets after escaping or tokenization | TB12 | Medium | Verified | `goal_state_limits`, `goal_state_notice.build_goal_state_notice`, `goal_tools.GoalToolsMiddleware._request_with_goal_notice` | +| T15 | DF28, DF29 | DC2 | Stored prompt injection through a goal, rubric, or status note influences later primary-model tool requests | TB12 | Medium | Likely | `goal_state_notice.build_goal_state_notice`, `goal_tools.GoalToolsMiddleware._request_with_goal_notice` | +| T16 | DF28, DF29 | DC2 | Sensitive local-file content, up to the 12,000-character rubric limit, is automatically persisted and transmitted to the configured model provider as rubric criteria | TB12 | Medium | Verified | `app.DeepAgentsApp._set_rubric_from_file`, `goal_state_notice.build_goal_state_notice` | +| T17 | DF28, DF29 | DC2 | Character-bounded goal/rubric/status-note text can still exceed provider context budgets after escaping or tokenization | TB12 | Medium | Verified | `goal_state_limits`, `goal_state_notice.build_goal_state_notice`, `goal_tools.GoalToolsMiddleware._request_with_goal_notice` | ### Threat Details @@ -422,19 +450,19 @@ #### T15: Stored Prompt Injection Through Goal/Rubric State -- **Flow**: DF26/DF27 (user or local-file content → checkpointed notice → primary-model context) +- **Flow**: DF28/DF29 (user or local-file content → checkpointed notice → primary-model context) - **Description**: The goal-state notice embeds the full actionable objective, active criteria, and status note in a synthetic `HumanMessage`. A rubric loaded from an untrusted repository file, or a crafted status note, can therefore persist instructions that influence later model behavior. HTML escaping and boundary labels prevent literal tag forgery. They do not stop natural-language prompt injection. Interactive HITL still gates side-effecting tool calls. Auto and non-interactive configurations can reduce that protection. - **Preconditions**: (1) The user accepts a goal/rubric or loads a file containing attacker-controlled instructions; (2) the state is actionable or the rubric remains active; (3) the primary model follows the injected content; (4) for side effects, the resulting tool call is approved or an approval-bypassing mode is active. #### T16: Automatic Disclosure of File-Loaded Rubrics -- **Flow**: DF26/DF27 (`/rubric file` → checkpoint → primary-model request) +- **Flow**: DF28/DF29 (`/rubric file` → checkpoint → primary-model request) - **Description**: `/rubric file` reads the entire selected UTF-8 text file and persists its nonempty contents, up to the 12,000-character rubric limit (see TB12); a larger file is rejected outright rather than truncated. The notice then embeds the criteria into primary-model context. There is no warning or confirmation specific to provider transmission, so a user can inadvertently select a secret-bearing or proprietary file. This flow handles user content, not provider credentials. Credential storage and provider retention are outside the scoped implementation. - **Preconditions**: (1) A user selects a file with sensitive content; (2) it becomes an active rubric; (3) a model request is made while the rubric is active. #### T17: Provider Context Pressure Despite Character Limits -- **Flow**: DF26/DF27 (character-bounded text → escaped notice → model request) +- **Flow**: DF28/DF29 (character-bounded text → escaped notice → model request) - **Description**: Direct, file-loaded, generated, and tool-authored goal-state paths enforce raw-character limits before persistence or notice construction. HTML escaping happens afterward and can expand the rendered notice (for example, `&` becomes `&`), while provider tokenization and available context budgets vary. The middleware restores or re-pins the current notice after compaction. A valid near-limit notice therefore remains recurring model-request overhead. This increases spend. It can also contribute to a provider context-limit failure. - **Preconditions**: (1) A user, file, or model-supplied status note produces a valid near-limit notice; (2) its escaped or tokenized representation is large relative to the configured provider's available context; (3) the corresponding goal or rubric remains model-visible. @@ -464,8 +492,8 @@ #### T6: Unauthenticated LangGraph Dev Server on Localhost -- **Flow**: DF3/DF4 (CLI ↔ LangGraph dev server) -- **Description**: The CLI spawns a `langgraph dev` server subprocess with `LANGGRAPH_AUTH_TYPE=noop` (`client/launch/server.py:_build_server_env`). This disables all server-side authentication. The server binds to `127.0.0.1:{port}` (a free ephemeral port by default, so it no longer squats the well-known `langgraph dev` port 2024). Any local process that discovers the port can: send arbitrary inputs to the running agent thread, read the agent's conversation state (including tool results that may contain file contents or secrets), inject messages into the conversation history, or trigger state updates. The server is ephemeral — it lives only for the duration of the CLI session — but this is the entire attack window. Port discovery is feasible via localhost port scanning or by reading `/proc/{pid}/cmdline` which contains the `--port` argument. +- **Flow**: DF3/DF4/DF26 (CLI ↔ LangGraph dev server) +- **Description**: The CLI spawns a `langgraph dev` server subprocess with `LANGGRAPH_AUTH_TYPE=noop` (`client/launch/server.py:_build_server_env`). This disables all server-side authentication. The server binds to `127.0.0.1:{port}` (a free ephemeral port by default, so it no longer squats the well-known `langgraph dev` port 2024). Any local process that discovers the port can send inputs, read conversation state (including tool results that may contain file contents or secrets), inject messages, trigger state updates, or request server-owned offload for a known thread. The offload route does not accept conversation state and cannot write `messages`, so its direct impact is additional model/archive work plus a state-only summarization update. The server is ephemeral — it lives only for the duration of the CLI session — but this is the entire attack window. Port discovery is feasible via localhost port scanning or by reading `/proc/{pid}/cmdline` which contains the `--port` argument. - **Preconditions**: (1) Attacker has a local process running as the same user (or as root); (2) Attacker discovers the server port (port scan on localhost, or reads process arguments). #### T7: LocalContextMiddleware Injects Host File Contents into System Prompt @@ -537,12 +565,13 @@ | Configuration | DF10, DF12, DF15, DF23| T9, T10, T12, T14 | Dotenv shell-env precedence; TOML schema; MCP schema + allow/deny lists; JSON structure check; `class_path` format check; dotenv denylist for execution-hook env keys and project-`.env` trust vars | User | Dotenv denylist is best-effort — execution-hook env keys consumed by tools not yet enumerated still reach subprocesses; `class_path` executes module code before type check; MCP env dict unfiltered; Auto classifier strength is a user choice with no floor enforced (T14) | | Session restore | DF14 | T5 | OS file permissions; SQLite | Project | Unencrypted at rest | | Server IPC (env vars) | DF18 | T6 | `ServerConfig` serialization; parent env passed to child | Project | Provider API keys flow to server subprocess; system prompt in env | +| Offload operation | DF26, DF27 | T6 | Per-field request schema on consumed context keys; endpoint/transport keys stripped from client `model_params` (`offload_api._strip_transport_model_params`); idle/pending/checkpoint checks; state-only update typed to permitted channels; messages writes rejected | Project | `context.model`/`profile_overrides` are type-checked but their values flow to `config.create_model` (see C17/TB11 for `class_path`); unknown context keys pass through by design; local built-in route relies on loopback and `noop` auth; custom deployments own route auth and thread authorization | | Host environment | DF19, DF20 | T7 | Static script; exit code check; 30s timeout | Shared | Makefile content injected into system prompt without sanitization | | Custom subagents (FS) | DF21 | T8 | `yaml.safe_load`; HITL on `task` tool | User | Subagent body text not content-filtered | | Async subagent config | DF22 | None direct | TOML parse; type validation in `load_async_subagents` | User | URL and headers for remote subagents are user-controlled; no URL validation | | MCP subprocess env | DF24 | T10 | Dict type check only (`_validate_server_config`) | User | No key/value filtering; arbitrary env vars forwarded to subprocess | -| Offloaded history | DF25 | None direct | Thread ID is UUID7 (no path injection); backend handles storage | Shared | Raw conversation content written to sandbox filesystem | -| Goal/rubric state | DF26, DF27 | T15, T16, T17 | Lifecycle projection; notice fingerprinting; raw-character validation (8,000 objective; 12,000 rubric and objective-plus-criteria; 4,000 note/blocker; 16,000 notice); boundary-tag escaping | Shared | Untrusted instructions remain model-readable; file contents are automatically transmitted; no post-escape, byte, or token budget or provider-transmission warning | +| Offloaded history | DF25 | None direct | Filename is a framework-minted session id (no path injection); backend handles storage | Shared | Raw conversation content written to sandbox filesystem | +| Goal/rubric state | DF28, DF29 | T15, T16, T17 | Lifecycle projection; notice fingerprinting; raw-character validation (8,000 objective; 12,000 rubric and objective-plus-criteria; 4,000 note/blocker; 16,000 notice); boundary-tag escaping | Shared | Untrusted instructions remain model-readable; file contents are automatically transmitted; no post-escape, byte, or token budget or provider-transmission warning | --- @@ -590,7 +619,7 @@ Threats that appear valid in isolation but fall outside project responsibility b | D1 | Unsafe msgpack deserialization in langgraph checkpoint loading | Verified fix status — confirmed fixed and closed upstream. | Users on current `langgraph` versions are not exposed. | Upstream langgraph has patched the unsafe msgpack deserialization. No longer an active risk. | | D2 | Unicode URL homoglyph as project vulnerability | Traced `check_url_safety` + `strip_dangerous_unicode` + `format_warning_detail` → approval dialog display | `unicode_security.check_url_safety`, `agent._format_fetch_url_description` | Warning system is the intended control — the project correctly surfaces the risk to the user in the approval dialog. Not a project vulnerability; classified as mitigated by design (UI warning). | | D3 | SSRF via `http_request` / `fetch_url` to internal services | Traced `tools.http_request` and `tools.fetch_url` — no URL scheme or host blocklist. However, both tools require HITL approval in interactive mode. In non-interactive mode, only shell commands are auto-approved via the allow-list; HTTP tools still go through the HITL interrupt gate. | `tools.http_request`, `tools.fetch_url`, `agent._add_interrupt_on` | Not a project vulnerability in isolation — the HITL gate is the intended control for all HTTP tool calls. The user sees the full URL before approving. SSRF is only reachable if the user approves the request (interactive) or enables auto-approve (explicit opt-in). Classified as out-of-scope for the same reason as prompt injection in interactive mode. | -| D4 | Offload path injection via thread_id | Checked `sessions.generate_thread_id` — returns UUID7 string (alphanumeric + hyphens only, no path separators). | `sessions.generate_thread_id`, `offload.offload_messages_to_backend` | Thread IDs are UUID7 strings generated by the framework. No user-controlled path components reach the file path. Not exploitable. | +| D4 | Offload path injection via archive filename | Checked `SummarizationMiddleware._get_session_id`/`_get_history_path` — the leaf is `session_` plus a `uuid4().hex`. | `SummarizationMiddleware._get_session_id`, `SummarizationMiddleware._get_history_path` | The archive filename is minted by the framework from a UUID4 hex, so no user-controlled path component reaches the file path. Not exploitable. | --- @@ -610,9 +639,12 @@ Threats that appear valid in isolation but fall outside project responsibility b | 2026-07-28 | manual update | Extended the out-of-scope hooks row for plugin-contributed hooks: enabled plugins may supply `hooks/hooks.json`, gated by install plus enablement rather than workspace trust, with each handler's environment overlaid only by its own plugin path variables | | 2026-08-03 | manual update | Added T14 (a weaker Auto classifier model weakens action review) under TB2, covering the selectable classifier (`--auto-classifier-model`, `DEEPAGENTS_CODE_AUTO_CLASSIFIER_MODEL`, `[models].auto_classifier`, `/auto model`), its restriction to trusted config surfaces via `config._PROJECT_DOTENV_DENIED_ENV_KEYS`, and its fail-closed construction behavior (deny, then latch to human approval; never fall back to the main model). Extended the "LLM output" and "Configuration" input-coverage rows with T14 | | 2026-08-04 | manual update | Extended T14 for the configurable Auto classifier review deadline (`DEEPAGENTS_CODE_AUTO_CLASSIFIER_TIMEOUT`, `[models].auto_classifier_timeout`): bounded by `config_manifest.resolve_auto_classifier_timeout` between a floor and ceiling so the deadline cannot be removed, denied from a project `.env` via `config._PROJECT_DOTENV_DENIED_ENV_KEYS`, and fail-closed on expiry | -| 2026-08-11 | langster-threat-model (automated) | Added C18, TB12, and DF26/DF27 for persisted goal/rubric state injected into primary-model context after removal of the goal/rubric read tools. Updated DC2 and input-source coverage; added T15 (stored prompt injection), T16 (automatic disclosure of file-loaded criteria), and T17 (context-budget pressure), with provider-handling and trust-contract gaps recorded as Open Questions. | -| 2026-08-11 | langster-threat-model (automated) | Corrected TB12, DF26/DF27, and T17 to document the enforced raw-character limits and the narrower residual risk from post-escape expansion and provider-specific byte/token budgets. | +| 2026-08-11 | langster-threat-model (automated) | Added C19, TB12, and DF28/DF29 for persisted goal/rubric state injected into primary-model context after removal of the goal/rubric read tools. Updated DC2 and input-source coverage; added T15 (stored prompt injection), T16 (automatic disclosure of file-loaded criteria), and T17 (context-budget pressure), with provider-handling and trust-contract gaps recorded as Open Questions. | +| 2026-08-11 | langster-threat-model (automated) | Corrected TB12, DF28/DF29, and T17 to document the enforced raw-character limits and the narrower residual risk from post-escape expansion and provider-specific byte/token budgets. | +| 2026-08-17 | langster-threat-model (diff) | Added C18, a server-owned custom HTTP boundary, and updated the architecture, DC5, TB2, TB10, DF25-DF27, T6, and input coverage. The route owns checkpoint hydration, shared compaction/hooks/backend selection, state-only persistence typed to permitted channels, and cost rollback; the client sends no graph or checkpoint state. Recorded that `PreToolUse` `ask` is fail-closed (converted to a deny) on this path because the operation transport carries no HITL channel; no new threat was identified. | | 2026-08-17 | manual update | Noted that the goal-state character budget is re-validated when a one-shot `/rubric next` is consumed, not only when it is set: the goal state it is measured against is mutable between those points (`/goal amend`, an `update_goal` blocker note), so a set-time-only check could still degrade the notice and silently disable the promised grade. The degraded notice now also reports `Goal status: unavailable` rather than leaving a live status beside `Goal actionable: no`. | | 2026-08-17 | manual update | Extended T14 for `ask_user` question text as a classifier injection source: the receipt attests display and answer, not content, so a question claiming prior or blanket authorization is untrusted content; recorded the `_CLASSIFIER_POLICY` clauses that keep a paired question to an action/target description matched against canonical arguments. Narrowed the T14 "deterministic allow/deny" guard wording, which overstated the deny side: deterministic denies do not cover the Deny categories, and an affirmative classifier allow is not re-checked downstream | | 2026-08-19 | manual update | Recorded that a superseded goal-state notice is replaced in place rather than removed from a model request, because the summarizer derives its next cutoff from that list and persists it against the unfiltered checkpoint; and that a combined objective-plus-criteria overflow now ends the criteria turn with its character limit instead of retrying blind to the recursion limit. Noted that an unrecognized persisted goal status degrades to `paused` in the notice, so a corrupt or forward-version checkpoint cannot present itself to the model as a goal to work toward. Dropped the stale generated-commit pin and bounded T16's disclosure to the enforced rubric limit | | 2026-08-21 | manual update | Clarified that the dotenv execution-hook denylist (`config._DOTENV_DENIED_ENV_KEYS`) is a best-effort enumeration of known code-execution consumers, not a closed set: TB11's inside-detail and T12's description now state this explicitly, and the Configuration input-coverage gap was reworded from the stale "project `.env` can set shell startup-hook vars" (denylisted since #4288) to the actual residual — execution-hook keys consumed by not-yet-enumerated tools still reach subprocesses | +| 2026-08-24 | langster-threat-model (diff) | Removed the client-seeded `/offload` fallback. `/offload` is now available only through C18 on built-in servers; local in-process and ACP agents do not support it, and custom or older servers without the route fail at the HTTP boundary. Updated DC5, TB2, TB10, and T6 to remove the client self-approval and synthetic-message attack surface. The server route, hook behavior, archive guard, and state-only persistence controls are unchanged; no new threat was identified. | +| 2026-08-24 | manual update | The C18 boundary now strips endpoint/proxy/transport keys (`base_url`, `openai_proxy`, `http_client`, and similar) from client-supplied `model_params` before they reach `config.create_model` (`offload_api._strip_transport_model_params`), closing the credential-redirection consequence of T6 for this route. Client-supplied `model` and behavioral params still flow through; in-process `CLIContextSchema` model params remain trusted and unfiltered | diff --git a/libs/code/deepagents_code/_cli_context.py b/libs/code/deepagents_code/_cli_context.py index 422c43774d..d9583f478f 100644 --- a/libs/code/deepagents_code/_cli_context.py +++ b/libs/code/deepagents_code/_cli_context.py @@ -70,8 +70,6 @@ class CLIContextSchema: turn_id: str | None = None - offload_tool_call_id: str | None = None - hooks_snapshot_id: str | None = None hooks_server_events: list[str] = field(default_factory=list) @@ -139,13 +137,6 @@ class CLIContext(TypedDict, total=False): turn_id: str | None """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. - - This is set by the client, not graph state, so model-generated calls cannot - grant themselves permission to execute during the hidden compaction turn. - """ - hooks_snapshot_id: str | None """Canonical Hooks v2 configuration hash for this session. diff --git a/libs/code/deepagents_code/_testing_models.py b/libs/code/deepagents_code/_testing_models.py index d75b38eeb8..58e275c9e9 100644 --- a/libs/code/deepagents_code/_testing_models.py +++ b/libs/code/deepagents_code/_testing_models.py @@ -20,6 +20,17 @@ from langchain_core.callbacks import CallbackManagerForLLMRun +DCA_TEST_OFFLOAD_GATE_ENV = "DCA_TEST_OFFLOAD_GATE_DIR" +"""Env var pointing at a directory used to gate summary generation. + +When set, a summary request writes `/entered` and then polls for +`/release` before replying. File-based so the test process can hold the +server's compaction model call open without sharing Python state across the +server subprocess boundary. Only summary prompts are gated; ordinary turns pass +through, which is what lets a test launch a concurrent run *while* `/offload` +is blocked here. +""" + # Prompt markers that drive `ToolCallingIntegrationChatModel`. Each marker is the # full token (including the trailing `=`); the file path follows on the same line, # e.g. `DCA_TEST_WRITE_FILE=/tmp/out.txt`. These are the single source of truth @@ -94,6 +105,7 @@ def _generate( if (text := self._stringify_message(message)).strip() ) if self._looks_like_summary_request(prompt): + self._wait_at_summary_gate() content = "integration summary" else: excerpt = " ".join(prompt.split()[-18:]) @@ -103,7 +115,18 @@ def _generate( content = "integration reply" return ChatResult( - generations=[ChatGeneration(message=AIMessage(content=content))] + generations=[ + ChatGeneration( + message=AIMessage( + content=content, + usage_metadata={ + "input_tokens": 100, + "output_tokens": 20, + "total_tokens": 120, + }, + ) + ) + ] ) @property @@ -111,6 +134,39 @@ def _llm_type(self) -> str: """LangChain model type identifier.""" return "deterministic-integration" + @staticmethod + def _wait_at_summary_gate() -> None: + """Hold the summary call open until the test releases it. + + No-op unless `DCA_TEST_OFFLOAD_GATE_DIR` names a directory. When set, + write `/entered` (the test's signal that the offload operation is + mid-summary) and then poll for `/release`. Every summary request + rewrites the marker, so the test reads it as "a summary is in flight" + rather than "the first summary started". Bounded so a crashed test + cannot wedge the server subprocess indefinitely. + + Raises: + TimeoutError: If the gate is not released within 120 seconds. + """ + import os + import time + from pathlib import Path + + gate_dir = os.environ.get(DCA_TEST_OFFLOAD_GATE_ENV) + if not gate_dir: + return + gate = Path(gate_dir) + (gate / "entered").write_text("1") + deadline = time.monotonic() + 120 + while not (gate / "release").exists(): + if time.monotonic() > deadline: + msg = ( + "Offload test gate was never released; refusing to block " + "the server summary call forever." + ) + raise TimeoutError(msg) + time.sleep(0.05) + @staticmethod def _stringify_message(message: BaseMessage) -> str: """Flatten message content into plain text for deterministic responses. diff --git a/libs/code/deepagents_code/agent.py b/libs/code/deepagents_code/agent.py index 848faed6f4..6c17bb7b59 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 ( + OffloadOperation, + _create_cli_compaction_middleware, + attach_offload_operation, +) 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 @@ -116,6 +120,7 @@ logger = logging.getLogger(__name__) + _MEMORY_READONLY_SYSTEM_PROMPT = ( "\n" "{agent_memory}\n\n" @@ -2943,7 +2948,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 @@ -2973,7 +2978,16 @@ 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) + + # Publish the server operation on the backend shared with `server_graph`. + # The custom HTTP route owns checkpoint access and persistence, while this + # object retains the exact compaction and hook instances used by the agent. + attach_offload_operation( + composite_backend, + OffloadOperation(compaction_middleware, 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 3169a31c77..936d014e43 100644 --- a/libs/code/deepagents_code/app.py +++ b/libs/code/deepagents_code/app.py @@ -76,14 +76,11 @@ from deepagents_code._session_stats import ( USAGE_KIND_LABELS, USAGE_KIND_ORDER, - RecordedRequest, SessionStats, SpinnerStatus, - finalize_recorded_requests, format_cost, format_cost_estimate, format_token_count, - record_message_usage, ) # All config imports — settings, create_model, detect_provider, is_ascii_mode, @@ -124,7 +121,6 @@ latest_goal_state_message_index, latest_goal_state_notice, log_malformed_summarization_event as _log_malformed_summarization_event, - summarization_cutoff as _summarization_cutoff, validated_summarization_cutoff as _validated_summarization_cutoff, ) from deepagents_code.iterm_cursor_guide import restore_iterm_cursor_guide @@ -583,19 +579,6 @@ def _warn_discarded_goal_channels(state_values: dict[str, Any]) -> list[str]: return discarded -_OFFLOAD_WEDGE_WARNING = ( - "Offload failed and the conversation may be left in an inconsistent state " - "(a compaction request could not be cleaned up). If your next message " - "errors, start a new thread." -) -"""Shown when a failed `/offload` could not remove its unanswered seed. - -A dangling `compact_conversation` tool call the model API later rejects would -otherwise wedge the thread with only a log warning; surfacing this tells the -user why an unrelated next turn might fail and how to recover. -""" - - def _effective_conversation(messages: list[Any], event: Any) -> list[Any]: # noqa: ANN401 """Reconstruct the effective conversation the model would see. @@ -613,11 +596,10 @@ def _effective_conversation(messages: list[Any], event: Any) -> list[Any]: # no The bounds case diverges deliberately from the SDK. The SDK reads a cutoff past the end as "everything was summarized" and returns `[summary]` alone. - Here, a list shorter than the cutoff means history was removed after the - summary was written; the `/offload` failure paths issue `RemoveMessage`. The - survivors are therefore recent messages, and `[summary]` would hide live - turns. Callers size context and detect dangling tool calls, where - over-reporting the window is the safe direction. + Here, a list shorter than the cutoff means history may have been removed + after the summary was written. The survivors are therefore treated as recent + messages, and `[summary]` would hide live turns. Callers size context and + detect dangling tool calls, where over-reporting the window is safer. `validated_summarization_cutoff` holds the matching decision for the notice predicate. @@ -644,37 +626,6 @@ def _effective_conversation(messages: list[Any], event: Any) -> list[Any]: # no return [summary, *messages[cutoff:]] -def _message_text(msg: Any) -> str: # noqa: ANN401 - """Extract the text content of a message object or serialized dict. - - Handles the shapes `/offload` sees across the LangGraph server boundary: - a message object with `.content`, or a serialized dict with `"content"`. - A string content is returned as-is; a list of content blocks has its text - parts concatenated (so a `ToolMessage` whose content is a block list is not - stringified to `"[{...}]"`, which would defeat prefix matching). - - Args: - msg: A message object or serialized message dict. - - Returns: - The concatenated text content, or an empty string when there is none. - """ - content = ( - msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "") - ) - if isinstance(content, str): - return content - if isinstance(content, list): - parts: list[str] = [] - for block in content: - if isinstance(block, str): - parts.append(block) - elif isinstance(block, dict) and isinstance(block.get("text"), str): - parts.append(block["text"]) - return "".join(parts) - return "" if content is None else str(content) - - def _is_tool_message(msg: Any) -> bool: # noqa: ANN401 """Return whether `msg` is a tool message in object or serialized form.""" if isinstance(msg, dict): @@ -685,44 +636,6 @@ def _is_tool_message(msg: Any) -> bool: # noqa: ANN401 return isinstance(msg, ToolMessage) -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". - - The caller passes only the messages produced by the *current* `/offload` - attempt (the tail after the pre-seed prefix). This matters because the - failure prefix is shared with the SDK's own compaction-failure wording, so - an unbounded scan could match a stale failure from an unrelated prior turn; - slicing to the current attempt keeps detection specific to this run. - - Args: - messages: The messages produced by this `/offload` attempt (objects or - serialized dicts), i.e. committed state beyond the pre-seed prefix. - - Returns: - The failure message text, or `None` if no failure marker is found. - """ - from deepagents_code.offload_middleware import COMPACTION_FAILURE_PREFIX - - for msg in reversed(messages): - if not _is_tool_message(msg): - continue - text = _message_text(msg) - if text.startswith(COMPACTION_FAILURE_PREFIX): - return text - return None - - -def _message_id(msg: Any) -> str | None: # noqa: ANN401 - """Return a message's id from object or serialized-dict form.""" - return msg.get("id") if isinstance(msg, dict) else getattr(msg, "id", None) - - def _message_tool_call_id(msg: Any) -> str | None: # noqa: ANN401 """Return the `tool_call_id` a tool message answers, if any.""" return ( @@ -12111,9 +12024,10 @@ async def _flush_pending_shell_messages(self) -> None: await remote.aensure_thread(remote_config) await self._agent.aupdate_state(config, {"messages": messages}) except Exception: # best-effort; UI already showed the output - # Parity with the offload path's `aupdate_state` failure handling: - # log the traceback and surface a non-blocking toast, since the - # model silently lacking output the user expects is confusing. + # Log the traceback and surface a non-blocking toast: the shell + # output is already on screen, so the only loss is the model not + # seeing it, and silently lacking output the user expects is + # confusing. logger.exception("Failed to flush shell command into model context") with suppress(Exception): self.notify( @@ -16566,218 +16480,96 @@ async def _get_conversation_token_count(self) -> int | None: _, conversation = await self._get_context_usage_counts() return conversation - async def _run_offload_task(self) -> None: - """Run a synchronously reserved `/offload` outside the App message pump.""" - self._offload_task_started = True - try: - await self._handle_offload(reserved=True) - finally: - self._offload_task_started = False - self._offload_worker = None - if not self._startup_sequence_running: - await self._process_next_from_queue() - - async def _handle_offload(self, *, reserved: bool = False) -> None: - """Run `/offload`, always releasing its busy-state reservation. - - Args: - reserved: Whether the caller already reserved the turn synchronously. - """ - if not reserved: - self._set_agent_running(True) - try: - await self._offload_impl() - finally: - self._set_agent_running(False) - try: - await self._set_spinner(None) - except Exception: # best-effort spinner cleanup - logger.exception("Failed to dismiss spinner after offload") - - async def _offload_impl(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`. - - Raises: - CancelledError: If the offload worker is interrupted. - """ - from langchain_core.messages.utils import count_tokens_approximately - + async def _handle_server_offload(self, config: RunnableConfig) -> None: + """Request and render the built-in server-owned offload operation.""" from deepagents_code.hooks.client_lifecycle import ClientHookStopError - if not self._agent or not self._lc_thread_id: + remote = self._remote_agent() + if remote is None: await self._mount_message( - AppMessage("Nothing to offload \u2014 start a conversation first"), + ErrorMessage("Offload failed: no dcode server is connected.") ) return - - config: RunnableConfig = {"configurable": {"thread_id": self._lc_thread_id}} - - try: - state_values = await self._get_thread_state_values(self._lc_thread_id) - except Exception as exc: # noqa: BLE001 - await self._mount_message(ErrorMessage(f"Failed to read state: {exc}")) - return - - if not state_values: - await self._mount_message( - AppMessage("Nothing to offload \u2014 start a conversation first"), - ) - return - + # Set once the server reports a committed compaction. Everything after + # that point is local reporting, and a failure there must not tell the + # user the offload failed -- they would run it again on an already + # compacted conversation. + committed = False try: await self._set_spinner("Offloading") - - prior_event = state_values.get("_summarization_event") - before_messages = state_values.get("messages", []) - # Bounds-checked against the list it indexes, matching - # `_effective_conversation` below. Without `message_count`, an - # out-of-bounds cutoff is trusted here while that call rejects it, and - # the two disagree: `messages_offloaded` inflates, `messages_kept` - # collapses to 0, and the report claims a large offload beside roughly - # zero token savings. - prior_cutoff = _summarization_cutoff( - prior_event, - message_count=len(before_messages), - ) - conversation_tokens_before = count_tokens_approximately( - _effective_conversation(before_messages, prior_event) + from deepagents_code._cli_context import CLIContext + from deepagents_code.config import settings + + 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, + # The operation runs the agent's `PreCompact` and `PreToolUse` + # hooks, and the server defaults a missing mode to `manual`. + # Without these a configured hook would see Manual during + # `/offload` even in Auto-Accept or YOLO, so a hook that keys + # its decision on the mode behaves differently here than on + # every interactive turn. + approval_mode=self._approval_mode.value, + auto_approve=self._auto_approve, + ) + self._hooks.apply_graph_context(context) + result = await remote.aoffload( + config=config, + context=context, + fulfill_hook=self._hooks.fulfill_interrupt, ) - reported_tokens_before = _persisted_context_tokens(state_values) + await self._sync_session_cost_from_checkpoint() - # 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()) - - try: - tool_error = await self._drive_server_side_compaction( - config, seed_tool_call_id + status = result.get("status") + if status == "empty": + await self._mount_message( + AppMessage("Nothing to offload — start a conversation first") ) - except ClientHookStopError: return - except (asyncio.CancelledError, Exception) as stream_error: - # A server graph can checkpoint the tool-node update before a - # later stream transport failure reaches this client. Reconcile - # the durable event before reporting the operation as failed. - logger.warning( - "Offload stream failed; checking for committed compaction state", - exc_info=True, - ) - try: - new_state = await self._get_thread_state_values(self._lc_thread_id) - except Exception as state_error: - logger.warning( - "Failed to reconcile state after offload stream error", - exc_info=True, - ) - if 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") - reconciled_cutoff = _summarization_cutoff( - reconciled_event, - message_count=len(new_state.get("messages", [])), - ) - if reconciled_cutoff <= 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 - # `tool_use` that the model API rejects on the next turn. - if not await self._remove_unanswered_offload_seed( - config, seed_tool_call_id - ): - await self._mount_message(ErrorMessage(_OFFLOAD_WEDGE_WARNING)) - raise - else: - if tool_error is not None: - # Tool failure can follow completed summary/model requests. - # Settle the display from the graph before returning the - # failure; local detailed stats never write this total. - await self._sync_session_cost_from_checkpoint() - await self._mount_message(ErrorMessage(tool_error)) - return - - # Read the persisted result back so the UI reflects server state - # (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) - # 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) - self._sync_cache_state_from_state(new_state) - new_event = new_state.get("_summarization_event") - new_cutoff = _summarization_cutoff( - new_event, - message_count=len(new_state.get("messages", [])), - ) - - 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. - 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 - ) - # `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. + if status == "noop": await self._mount_message( AppMessage( - "Nothing to offload \u2014 the conversation is already " - "compact.", - ), + "Nothing to offload — the conversation is already compact." + ) ) return - - archive_path = ( - new_event.get("file_path") - if isinstance(new_event, dict) - else getattr(new_event, "file_path", None) - ) - # Recompute the post-offload conversation from the original pre-seed - # messages plus the new event. This excludes the compact tool's own - # machinery while preserving the provider-reported system/tool overhead - # from the last ordinary turn when that total is available. - conversation_tokens_after = count_tokens_approximately( - _effective_conversation(before_messages, new_event) + if status in {"denied", "failed"}: + error = result.get("error") or "The server rejected the operation." + await self._mount_message(ErrorMessage(f"Offload failed: {error}")) + return + if status != "compacted": + await self._mount_message( + ErrorMessage( + "Offload failed: the server returned an invalid result." + ) + ) + return + committed = True + + # The server reports `count_tokens_approximately` over the effective + # conversation, so these are conversation-scale estimates. + # `usage_label` tells the user which metric they are reading because + # "Conversation" excludes the system/tool overhead that "Context" + # includes and the two percentages are not comparable across + # offloads. + conversation_tokens_before = result["tokens_before"] + conversation_tokens_after = result["tokens_after"] + # A cached count from a real model turn is the provider's own total; + # an approximate one is our own estimate, and rescaling an estimate + # by an estimate would compound the error. Only the former promotes + # the report to context scale. + reported_tokens_before = ( + 0 if self._tokens_approximate else self._context_tokens ) if reported_tokens_before: # Subtract the *delta* from the provider total rather than - # rebuilding the total as `overhead + conversation_after`. The two - # are algebraically equal, but only this form keeps both figures on - # the provider's scale: `count_tokens_approximately` need only - # overshoot the provider count by a token for an - # `overhead = max(0, reported - conversation_before)` clamp to - # collapse the overhead to zero, which would silently report the - # whole system prompt and tool schema as freed context. - # - # The estimator's error appears with opposite signs in the two - # conversation counts and largely cancels in the difference, which - # is why `before` prints exact and only `after` carries a `~`. + # rebuilding it as `overhead + conversation_after`; this keeps + # both figures on the provider's scale. The estimator's error + # appears with opposite signs in the two conversation counts and + # largely cancels in the difference, so only `after` carries a `~`. tokens_before = reported_tokens_before tokens_after = max( 0, @@ -16788,515 +16580,161 @@ async def _offload_impl(self) -> None: before = format_token_count(tokens_before) after = f"~{format_token_count(tokens_after)}" else: - # No usable provider total (never set, or a checkpoint value - # `_persisted_context_tokens` rejected), so fall back to a - # conversation-only estimate. `usage_label` is what tells the user - # which metric they are reading: "Conversation" excludes the - # system/tool overhead that "Context" includes, so the two - # percentages are not comparable across offloads. tokens_before = conversation_tokens_before tokens_after = conversation_tokens_after usage_label = "Conversation" before = f"~{format_token_count(tokens_before)}" after = f"~{format_token_count(tokens_after)}" - - # Message and turn counts are derived purely from the absolute cutoffs - # into the ORIGINAL pre-seed `before_messages`, never from the post-run - # `new_state["messages"]`. The compact tool's own machinery (the seeded - # tool call, its result, and the trailing model turn) lands at/after - # `new_cutoff` in the post-run list, so slicing that list instead would - # count those artifacts as kept conversation. - messages_offloaded = max(0, new_cutoff - prior_cutoff) - messages_kept = max(0, len(before_messages) - new_cutoff) - turns_offloaded = sum( - is_human_message(message) and not is_internal_message(message) - for message in before_messages[prior_cutoff:new_cutoff] - ) - turns_kept = sum( - is_human_message(message) and not is_internal_message(message) - for message in before_messages[new_cutoff:] - ) - offloaded_message_label = ( - "message" if messages_offloaded == 1 else "messages" - ) - kept_message_label = "message" if messages_kept == 1 else "messages" - offloaded_turn_label = "turn" if turns_offloaded == 1 else "turns" - kept_turn_label = "turn" if turns_kept == 1 else "turns" - # Floored at zero: a summary can come out larger than the messages it - # replaced, and in the reported branch `tokens_after` mixes an exact - # provider total with an estimated delta. Neither should ever render as - # a negative "decrease". + # Floored at zero: a summary can come out larger than the messages + # it replaced, and in the reported branch `tokens_after` mixes an + # exact provider total with an estimated delta. Neither should ever + # render as a negative "decrease". pct = ( max(0, round((tokens_before - tokens_after) / tokens_before * 100)) if tokens_before > 0 else 0 ) - - offloaded_counts = ( - f"{messages_offloaded} older {offloaded_message_label} " - f"({turns_offloaded} conversation {offloaded_turn_label})" + kept_message_label = ( + "message" if result["messages_kept"] == 1 else "messages" ) if tokens_after <= tokens_before: stats_line = ( f"{usage_label}: {before} → {after} tokens ({pct}% decrease), " - f"{messages_kept} {kept_message_label} " - f"({turns_kept} conversation {kept_turn_label}) kept." - ) - outcome = ( - f"Offloaded {offloaded_counts}, freeing up context window space." + f"{result['messages_kept']} {kept_message_label} kept." ) else: stats_line = ( f"{usage_label}: {before} → {after} tokens (increase), " - f"{messages_kept} {kept_message_label} " - f"({turns_kept} conversation {kept_turn_label}) kept." + f"{result['messages_kept']} {kept_message_label} kept." ) - outcome = ( + offloaded_message_label = ( + "message" if result["messages_offloaded"] == 1 else "messages" + ) + offloaded_counts = ( + f"{result['messages_offloaded']} older {offloaded_message_label}" + ) + outcome = ( + f"Offloaded {offloaded_counts}, freeing up context window space." + if tokens_after <= tokens_before + else ( f"Offloaded {offloaded_counts}, but the summary was larger " "than the messages it replaced, so context increased." ) - if archive_path: - from deepagents_code.offload import offload_storage_is_ephemeral - - # In local mode the archive may have landed in a temp fallback - # directory (persistent `~/.deepagents` was unwritable). The - # write succeeded, so context was freed and history is readable - # now, but it may not survive a restart -- say so rather than - # imply durable storage. + ) + if result.get("archive_path"): caveat = ( "\nNote: history was saved to temporary storage and may not " "survive a restart." - if offload_storage_is_ephemeral() + if result.get("archive_ephemeral") else "" ) await self._mount_message( - AppMessage( - f"{outcome}\n{stats_line}{caveat}", - ), + AppMessage(f"{outcome}\n{stats_line}{caveat}") ) else: - # Context was still freed (the summary is in-context), but the - # archive write failed, so the offloaded messages are not - # recoverable. Surface both facts in one message rather than a - # separate warning immediately followed by a success line. await self._mount_message( ErrorMessage( - f"{outcome} The " - "conversation history could not " - "be saved to storage, so those messages are not " - f"recoverable. Check logs for details.\n{stats_line}", + f"Offloaded {offloaded_counts} " + "and freed context, but the conversation history could not " + "be saved to storage, so those messages are not recoverable. " + f"Check logs for details.\n{stats_line}" ) ) - + # Flagged approximate in both branches: the conversation-scale + # figure is an estimate outright, and the provider-scale one mixes an + # exact total with an estimated delta. self._on_tokens_update(tokens_after, approximate=True) - except Exception as exc: # surface offload errors to user - logger.exception("Offload failed") - await self._mount_message(ErrorMessage(f"Offload failed: {exc}")) - - async def _drive_server_side_compaction( - self, config: RunnableConfig, seed_tool_call_id: str | None = None - ) -> str | None: - """Trigger the server-side `compact_conversation` tool with `force=True`. - - 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. - - A first-turn `Command(update=..., goto=...)` is intentionally avoided: - the LangGraph API server rebuilds it with `goto=None` and crashes - `_control_branch`. The `aupdate_state(as_node="model")` + `astream` - continuation is the stable path. - - Args: - config: Config with `configurable.thread_id`. - seed_tool_call_id: Id for the seeded tool call. Supplied by - `_handle_offload` so it can remove the seed if the run fails; - a fresh id is generated when omitted (e.g. direct callers). - - Returns: - An error string when the tool reported a compaction failure, or - `None` when the run completed (whether it compacted or was a - no-op — the caller distinguishes those from persisted state). - Note the `None` return also covers the bounded-drain-exceeded - path, which has already mounted its own user-facing message - before returning. - """ - from langchain.agents.middleware.human_in_the_loop import ( - ApproveDecision, - RejectDecision, - ) - from langchain_core.messages import AIMessage - from langgraph.types import Command - - from deepagents_code._tracing import stream_trace_config - from deepagents_code.config import settings - from deepagents_code.hooks.client_lifecycle import ClientHookStopError - from deepagents_code.hooks.interrupt import is_hook_interrupt_payload - from deepagents_code.hooks.models.domain import SessionStartCause - from deepagents_code.offload_middleware import ( - COMPACTION_FAILURE_PREFIX, - _offload_seed_message_id, - ) - - agent = self._agent - if agent is None: - return None - - offload_stats = SessionStats() - offload_thread_id = self._lc_thread_id - recorded_usage_requests: dict[str, RecordedRequest] = {} - model_spec = self._effective_model_spec() or "" - fallback_provider, separator, fallback_model = model_spec.partition(":") - if not separator: - fallback_model = fallback_provider - fallback_provider = "" - - tool_call_id = seed_tool_call_id or str(uuid.uuid4()) - # Stable message id so a failed run can address the seed for removal. - seed = AIMessage( - content="", - id=_offload_seed_message_id(tool_call_id), - tool_calls=[ - { - "name": "compact_conversation", - "args": {"force": True}, - "id": tool_call_id, - } - ], - ) - - # Remote dev servers separate checkpoint persistence from HTTP thread - # registration; register before mutating state so the write lands. - if remote := self._remote_agent(): - await remote.aensure_thread( - {"configurable": {"thread_id": self._lc_thread_id}} - ) - await agent.aupdate_state(config, {"messages": [seed]}, as_node="model") - - tool_error: str | None = None - # `self._agent` includes local graphs whose generic context defaults to - # `None`, but the graph is built with `CLIContextSchema` at runtime. - streaming_agent = cast("Any", agent) - - seeded_compaction_approved = False - compact_boundary_fired = False - 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, - classifier_model=self._auto_classifier_context_value(), - thread_id=self._lc_thread_id, - offload_tool_call_id=tool_call_id, - ) - self._hooks.apply_graph_context(stream_context) - - def _decisions_for_interrupt(interrupt_obj: Any) -> list[Any]: # noqa: ANN401 - """Approve the forced compaction; reject any other gated tool call. - - HITL action requests do not expose tool-call IDs, so the seeded - request is identified by its exact forced arguments and approved - at most once. Any repeated compaction request fails closed. - - Args: - interrupt_obj: The interrupt surfaced by the HITL middleware. + # Fired only after the outcome is on screen. Compaction has already + # committed server-side, so a hook that raises must not be allowed + # to replace the user's result with "Offload failed" and leave the + # status bar on pre-offload counts. `_run_session_start_hook` mounts + # a stop reason itself, so the user sees both facts in order. + from deepagents_code.hooks.models.domain import SessionStartCause - Returns: - One decision per `action_request`, in order, as the HITL - middleware requires. - """ - nonlocal seeded_compaction_approved - value = getattr(interrupt_obj, "value", None) - action_requests = ( - value.get("action_requests") if isinstance(value, dict) else None - ) - if not action_requests: - # Without an identifiable action, approving could execute a - # different gated tool. A singleton rejection safely answers - # the surfaced interrupt. - return [ - RejectDecision( - type="reject", - message=( - "Not executed: /offload could not identify the " - "requested action." - ), - ) - ] - decisions: list[Any] = [] - for req in action_requests: - name = req.get("name") if isinstance(req, dict) else None - args = req.get("args") if isinstance(req, dict) else None - is_seeded_request = ( - not seeded_compaction_approved - and name == "compact_conversation" - and isinstance(args, dict) - and args.get("force") is True - ) - if is_seeded_request: - decisions.append(ApproveDecision(type="approve")) - seeded_compaction_approved = True - else: - decisions.append( - RejectDecision( - type="reject", - message=( - "Not executed: /offload only performs " - "conversation compaction." - ), - ) + try: + await self._run_session_start_hook(SessionStartCause.COMPACT) + except ClientHookStopError: + # The stop reason is already mounted; the offload itself stands. + pass + except Exception: + logger.exception("SessionStart hook failed after server offload") + await self._mount_message( + ErrorMessage( + "The conversation was offloaded, but a configured " + "SessionStart hook failed. Check logs for details." ) - return decisions - - async def _drain(stream_input: Any) -> list[tuple[str, dict[str, Any]]]: # noqa: ANN401 - """Advance the graph, collecting interrupts that need a resume. - - Sets `tool_error` if the compaction tool reported a failure. - - Args: - stream_input: `None` to advance, or a `Command(resume=...)` to - answer pending interrupts. - - Returns: - `(interrupt_id, resume_value)` pairs for every interrupt - surfaced during this stream. + ) + except ClientHookStopError: + # Raised before the compaction committed (e.g. from a hook the + # operation routed back to this client); the reason is mounted by + # the hook executor, so adding "Offload failed" would double-report. + return + except Exception as exc: + from deepagents_code.client.remote_client import format_agent_exception - Raises: - ClientHookStopError: If compact-session startup is blocked. - """ - nonlocal compact_boundary_fired, tool_error - pending: list[tuple[str, dict[str, Any]]] = [] - async for chunk in streaming_agent.astream( - stream_input, - stream_mode=["messages", "updates"], - subgraphs=True, - config=stream_trace_config(config, stream_input), - context=stream_context, - durability="exit", - ): - if not isinstance(chunk, tuple) or len(chunk) != 3: # noqa: PLR2004 # (namespace, mode, data) - continue - _namespace, mode, data = chunk - if mode == "updates" and isinstance(data, dict): - for interrupt_obj in data.get("__interrupt__") or []: - iid = getattr(interrupt_obj, "id", None) - if iid: - value = getattr(interrupt_obj, "value", None) - if is_hook_interrupt_payload(value): - resume = await self._hooks.fulfill_interrupt(value) - pending.append((iid, resume)) - continue - decisions = _decisions_for_interrupt(interrupt_obj) - pending.append((iid, {"decisions": decisions})) - elif mode == "messages" and isinstance(data, tuple): - msg = data[0] - recorded_usage = record_message_usage( - offload_stats, - msg, - fallback_model=fallback_model, - fallback_provider=fallback_provider, - request_metadata=( - data[1] - if len(data) > 1 and isinstance(data[1], dict) - else None - ), - kind="offload", - recorded_requests=recorded_usage_requests, - ) - if ( - recorded_usage is not None - and recorded_usage.cost_usd is not None - ): - # Display-only until the graph's next absolute total - # arrives. The ledger is closed at each round boundary, - # so a replayed resume round cannot add this twice. - self._add_provisional_cost(recorded_usage.cost_usd) - if _is_tool_message(msg): - text = _message_text(msg) - if text.startswith(COMPACTION_FAILURE_PREFIX) or ( - getattr(msg, "name", None) == "compact_conversation" - and getattr(msg, "status", None) == "error" - ): - tool_error = text - elif ( - getattr(msg, "name", None) == "compact_conversation" - and text.startswith("Conversation compacted.") - and not compact_boundary_fired - ): - compact_boundary_fired = True - if not await self._run_session_start_hook( - SessionStartCause.COMPACT - ): - msg = "Compact continuation stopped by hook" - raise ClientHookStopError(msg) - # Close the ledger at the round boundary: the next resume round - # replays these chunks, which would otherwise revise their requests - # a second time and double the offload's tokens and cost. - finalize_recorded_requests(recorded_usage_requests) - return pending - - # Bound the resume loop: after compaction the model runs again, and a - # 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: - logger.warning( - "Offload exceeded %d resume rounds; leaving %d interrupt(s) " - "unresolved", - max_resume_rounds, - len(pending), + if committed: + # The server already compacted and persisted. Reporting the + # rendering failure as "Offload failed" would send the user to + # offload a second time, so name what actually broke. + logger.exception("Server offload committed but reporting failed") + await self._mount_message( + ErrorMessage( + "The conversation was offloaded, but the result could " + "not be displayed. Check logs for details." ) - # Compaction itself already committed in round 1, so the caller - # still reports the offload. Surface the abandoned drain so the - # user knows the thread was left paused mid-run and may need a - # fresh message to reset. Skip this when a tool failure is - # already pending, so the caller shows that error instead of - # the user seeing two conflicting messages. - if tool_error is None: - await self._mount_message( - ErrorMessage( - "Offload completed, but the agent kept requesting " - "tools afterward and the run could not be fully " - "drained. Send a new message to continue; the " - "thread may need to reset." - ) - ) - break - resume_payload = dict(pending) - pending = await _drain(Command(resume=resume_payload)) + ) + return + logger.exception("Server offload failed") + await self._mount_message( + ErrorMessage(f"Offload failed: {format_agent_exception(exc)}") + ) - return tool_error + async def _run_offload_task(self) -> None: + """Run a synchronously reserved `/offload` outside the App message pump.""" + self._offload_task_started = True + try: + await self._handle_offload(reserved=True) finally: - # Usage can be incurred even when compaction fails, does nothing, or - # raises after a completed model request. Keep the graph-owned cost - # total untouched; these merges supply only the local breakdown. - self._session_stats.merge(offload_stats) - if offload_thread_id == self._lc_thread_id: - self._thread_stats.merge(offload_stats) - self._refresh_cache_display() - - async def _remove_offload_artifacts( - self, - config: RunnableConfig, - messages: list[Any], - prior_event: object, - ) -> None: - """Restore state changed by a no-op `/offload` graph run. + self._offload_task_started = False + self._offload_worker = None + if not self._startup_sequence_running: + await self._process_next_from_queue() - 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 - committed seed/result pair left behind is harmless (unlike an unanswered - seed); letting the write raise here would misreport a working offload as - "Offload failed" via the caller's outer handler. + async def _handle_offload(self, *, reserved: bool = False) -> None: + """Run `/offload`, always releasing its busy-state reservation. Args: - config: Config with `configurable.thread_id`. - messages: Messages appended after the pre-run state snapshot. - prior_event: Summarization event from the pre-run state snapshot. + reserved: Whether the caller already reserved the turn synchronously. """ - from langchain_core.messages import RemoveMessage - - agent = self._agent - if agent is None: - return - removals = [ - RemoveMessage(id=message_id) - for message in messages - if (message_id := _message_id(message)) is not None - ] + if not reserved: + self._set_agent_running(True) try: - await agent.aupdate_state( - config, - { - "messages": removals, - "_summarization_event": prior_event, - }, - as_node="model", - ) - except Exception: # best-effort restoration; keep the no-op report - logger.warning( - "Failed to restore state after a no-op offload run", exc_info=True - ) - - async def _remove_unanswered_offload_seed( - self, config: RunnableConfig, seed_tool_call_id: str - ) -> bool: - """Remove a committed `/offload` seed whose tool call was never answered. - - The seed `AIMessage` carrying the forced `compact_conversation` call is - committed via `aupdate_state` before the run advances — independently of - the stream's durability. If the run then fails before the tool produces - a `ToolMessage`, the seed is left as an unanswered `tool_use` in - committed state, which the model API rejects on the next turn - ("tool_use ids ... without tool_result"), potentially wedging the - thread. This best-effort removes that seed so a failed `/offload` leaves - a valid conversation. - - If the tool *did* run (a `ToolMessage` answers the call), the seed and - its result form a valid pair and are left untouched — removing the seed - alone would orphan the `ToolMessage`. - - Args: - config: Config with `configurable.thread_id`. - seed_tool_call_id: The id of the seeded `compact_conversation` call. - - Returns: - True if the thread is known to be free of a dangling seed (removed, - validly answered, or absent). False if a dangling seed may - remain because the state read or the removal write failed — the - caller should warn the user the thread may be inconsistent. - """ - from langchain_core.messages import RemoveMessage + await self._offload_impl() + finally: + self._set_agent_running(False) + try: + await self._set_spinner(None) + except Exception: # best-effort spinner cleanup + logger.exception("Failed to dismiss spinner after offload") - agent = self._agent - if agent is None or not self._lc_thread_id: - return True - try: - state = await self._get_thread_state_values(self._lc_thread_id) - except Exception: # best-effort cleanup; keep the original error - logger.warning( - "Could not read state to clean up offload seed", exc_info=True + async def _offload_impl(self) -> None: + """Offload older messages through the dcode server operation.""" + if not self._agent or not self._lc_thread_id: + await self._mount_message( + AppMessage("Nothing to offload — start a conversation first"), ) - return False - - messages = state.get("messages", []) - # An answering ToolMessage means the tool ran; the pair is valid. - if any( - _is_tool_message(msg) and _message_tool_call_id(msg) == seed_tool_call_id - for msg in messages - ): - return True + return - seed_id = next( - ( - _message_id(msg) - for msg in messages - if seed_tool_call_id in _message_tool_call_ids(msg) - ), - None, - ) - if not seed_id: - return True - try: - await agent.aupdate_state( - config, {"messages": [RemoveMessage(id=seed_id)]}, as_node="model" + remote = self._remote_agent() + if remote is None: + await self._mount_message( + AppMessage("Offload is not supported for local agents."), ) - except Exception: # best-effort cleanup; keep the original error - logger.warning("Failed to remove dangling offload seed", exc_info=True) - return False - return True + return + + config: RunnableConfig = {"configurable": {"thread_id": self._lc_thread_id}} + await self._handle_server_offload(config) def _start_plugin_auto_update(self) -> None: """Start the plugin auto-update worker.""" diff --git a/libs/code/deepagents_code/client/launch/server.py b/libs/code/deepagents_code/client/launch/server.py index 05f6195979..1be641391b 100644 --- a/libs/code/deepagents_code/client/launch/server.py +++ b/libs/code/deepagents_code/client/launch/server.py @@ -40,6 +40,10 @@ `langgraph dev` projects alongside `deepagents-code` without a port collision. """ +_DCODE_GRAPH_REF = "deepagents_code.server_graph:make_graph" +"""Built-in graph reference. Also gates registration of the offload HTTP app: +a custom `graph_ref` gets no `http` block and does not support `/offload`.""" + _HEALTH_POLL_INTERVAL_LOCAL = 0.1 _HEALTH_POLL_INTERVAL_REMOTE = 0.3 @@ -161,20 +165,35 @@ def _extract_startup_error_marker(output: str) -> str | None: def generate_langgraph_json( output_dir: str | Path, *, - graph_ref: str = "deepagents_code.server_graph:make_graph", + graph_ref: str = _DCODE_GRAPH_REF, env_file: str | None = None, checkpointer_path: str | None = None, + auth_path: str | None = None, ) -> Path: """Generate a `langgraph.json` config file for `langgraph dev`. + Registers the interactive `agent` graph and dcode's custom HTTP operations, + which opt into LangGraph's route-auth layer (`enable_custom_route_auth`) so a + deployment that configures auth gates them. Production runs `noop` auth and + relies on the loopback bind, so that gate is inert there -- see + `auth_path` below and THREAT_MODEL.md TB10. `/offload` is served by that + backend boundary rather than exposed as another client-addressable graph. + 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. + Custom graphs omit the built-in offload service, so `/offload` is + unsupported when one is supplied. 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 to disk instead of in-memory. + auth_path: Import path to a LangGraph `Auth` instance, emitted as the + config's `auth.path`. Production servers run with + `LANGGRAPH_AUTH_TYPE=noop` and no auth backend; this exists so + tests can prove `enable_custom_route_auth` gates the operation + routes when a deployment *does* configure one. Returns: Path to the generated config file. @@ -183,6 +202,13 @@ def generate_langgraph_json( "dependencies": ["."], "graphs": {"agent": graph_ref}, } + if graph_ref == _DCODE_GRAPH_REF: + config["http"] = { + "app": "deepagents_code.offload_api:app", + "enable_custom_route_auth": True, + } + if auth_path: + config["auth"] = {"path": auth_path} 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..a50b7d3556 100644 --- a/libs/code/deepagents_code/client/remote_client.py +++ b/libs/code/deepagents_code/client/remote_client.py @@ -10,10 +10,12 @@ import asyncio import logging -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable, Mapping + from collections.abc import AsyncIterator, Awaitable, Callable, Mapping + + from deepagents_code.offload_middleware import OffloadResult logger = logging.getLogger(__name__) @@ -26,6 +28,124 @@ how many runs are active. """ +_OFFLOAD_CANCEL_WAIT_SECONDS = 10.0 +"""Bound the Esc path while the server confirms offload termination.""" + +_OFFLOAD_MAX_RESUME_ROUNDS = 32 +"""Bound hook transport rounds for a single server operation. + +Counts *hook fulfillments*, not POSTs: the loop runs one extra iteration so the +round that follows the last fulfillment can observe completion. Exceeding this +many distinct invocations means either genuinely many hooks or an unstable +invocation-id derivation server-side. +""" + +_OFFLOAD_RESULT_INT_FIELDS = ( + "messages_offloaded", + "messages_kept", + "tokens_before", + "tokens_after", +) + + +async def _join_task_deferring_cancellation[T](task: asyncio.Task[T]) -> None: + """Join a task despite repeated cancellation of the waiting caller.""" + while not task.done(): + try: + await asyncio.wait((task,)) + except asyncio.CancelledError: + continue + + +async def _cancel_server_offload( + graph: Any, # noqa: ANN401 # untyped RemoteGraph client + thread_id: str, + operation_id: str, +) -> str: + """Request cancellation and wait for the server's terminal acknowledgement. + + Returns: + Server terminal status (`cancelled` or `finished`). + + Raises: + RuntimeError: If the server returns an invalid acknowledgement. + """ + response = await asyncio.wait_for( + graph.client.http.post( + f"/dcode/threads/{thread_id}/offload/{operation_id}/cancel", + json={}, + ), + timeout=_OFFLOAD_CANCEL_WAIT_SECONDS, + ) + status = response.get("status") if isinstance(response, dict) else None + if status not in {"cancelled", "finished"}: + msg = "Offload server returned an invalid cancellation acknowledgement." + raise RuntimeError(msg) + return status + + +async def _await_offload_step[T]( + awaitable: Awaitable[T], + *, + graph: Any, # noqa: ANN401 # untyped RemoteGraph client + thread_id: str, + operation_id: str, +) -> T: + """Await one offload step and confirm server termination if cancelled. + + Returns: + The awaited step's result. + + Raises: + asyncio.CancelledError: After the server confirms the operation ended. + """ + try: + return await awaitable + except asyncio.CancelledError: + cancellation = asyncio.create_task( + _cancel_server_offload(graph, thread_id, operation_id) + ) + await _join_task_deferring_cancellation(cancellation) + cancellation.result() + raise + + +def _validated_offload_result(result: object) -> OffloadResult: + """Check a server offload result before any caller indexes it. + + The renderer subscripts these fields by key and unguarded. Validating here + means a protocol skew fails with a message naming the problem, instead of a + `KeyError` reported as a generic reporting failure for an offload the server + already committed. + + Args: + result: The `result` object from a `complete` operation response. + + Returns: + The same mapping, once its required fields are known to be present. + + Raises: + RuntimeError: If a required field is missing or has the wrong type. + """ + if not isinstance(result, dict): + msg = "Offload server completed without a typed result." + raise RuntimeError(msg) # noqa: TRY004 # protocol fault, not a type misuse + status = result.get("status") + if not isinstance(status, str) or not status: + msg = "Offload result has no status." + raise RuntimeError(msg) + # Statistics are only meaningful (and only read) for a committed compaction. + if status == "compacted": + for field in _OFFLOAD_RESULT_INT_FIELDS: + value = result.get(field) + if not isinstance(value, int) or isinstance(value, bool): + msg = ( + f"Offload result field {field!r} must be an integer, got " + f"{type(value).__name__}." + ) + raise RuntimeError(msg) # noqa: TRY004 # protocol fault + return cast("OffloadResult", result) + def _require_thread_id(config: Mapping[str, Any] | None) -> str: """Extract and validate that `thread_id` is present in config. @@ -73,7 +193,7 @@ def agent_error_type(exc: BaseException) -> str: def format_agent_exception(exc: BaseException) -> str: - """Render an exception from `RemoteAgent.astream` for the UI. + """Render an exception from any `RemoteAgent` call for the UI. The LangGraph server serializes non-allowlisted exceptions as `{"error": , "message": }` @@ -82,7 +202,8 @@ def format_agent_exception(exc: BaseException) -> str: Python dict repr in the UI. Args: - exc: The exception caught from the agent stream. + exc: The exception caught from an agent call -- the SSE stream, or an + HTTP operation such as the offload route. Returns: `": "` for `RemoteException` dict payloads, @@ -154,6 +275,129 @@ def _get_graph(self) -> Any: # noqa: ANN401 ) return self._graph + async def aoffload( + self, + *, + config: Mapping[str, Any], + context: Mapping[str, Any], + fulfill_hook: Callable[[object], Awaitable[dict[str, object]]], + ) -> OffloadResult: + """Request server-owned offload and fulfill its hook callbacks. + + Args: + config: Runnable config identifying the thread. + context: Runtime model and Hooks v2 context. + fulfill_hook: Client hook executor for server requests. + + Returns: + Typed offload result from the server operation. + + Raises: + TypeError: If the server response is not a JSON object. + RuntimeError: If the server does not provide the offload route, + returns an invalid protocol response, or exceeds the hook round + limit. + APIStatusError: From the server operation. A 409 (conflict) or 422 + (malformed request) means no state was committed; a 500 raised + as indeterminate means a commit may have landed and carries + user-actionable text the caller must surface rather than + swallow. + """ # noqa: DOC502 -- APIStatusError is raised by the SDK transport + from uuid import uuid4 + + from langgraph_sdk.errors import NotFoundError + + from deepagents_code.hooks.interrupt import is_hook_interrupt_payload + + thread_id = _require_thread_id(config) + # The operation reads and writes thread state over HTTP, so the thread's + # live row must exist first. Checkpoint persistence and registration are + # separate on the dev server (see `aensure_thread`), so a resumed thread + # -- or one whose server restarted mid-session -- has state on disk and + # no row, and every request below would 404. + await self.aensure_thread({"configurable": {"thread_id": thread_id}}) + operation_id = str(uuid4()) + hook_responses: dict[str, object] = {} + graph = self._get_graph() + for round_index in range(_OFFLOAD_MAX_RESUME_ROUNDS + 1): + try: + response = await _await_offload_step( + graph.client.http.post( + f"/dcode/threads/{thread_id}/offload", + json={ + "operation_id": operation_id, + "context": dict(context), + "hook_responses": hook_responses, + }, + ), + graph=graph, + thread_id=thread_id, + operation_id=operation_id, + ) + except NotFoundError as exc: + # The route itself is missing. An unregistered thread cannot + # reach here as a 404: the server catches that and answers 409 + # with its own message. A custom `graph_ref` server, or one + # older than this operation, never registers the dcode HTTP app, + # and the SDK's bare "404 Not Found" names neither the cause nor + # a fix. + msg = ( + "This server does not provide dcode's /offload operation. " + "Use the built-in dcode server, or upgrade the server to a " + "version that registers it." + ) + raise RuntimeError(msg) from exc + if not isinstance(response, dict): + msg = "Offload server returned a non-object response." + raise TypeError(msg) + status = response.get("status") + if status == "complete": + return _validated_offload_result(response.get("result")) + request = response.get("request") + if status != "interrupt" or not is_hook_interrupt_payload(request): + msg = "Offload server returned an invalid operation response." + raise RuntimeError(msg) + invocation = request.get("request") + invocation_id = ( + invocation.get("invocation_id") + if isinstance(invocation, dict) + else None + ) + if not isinstance(invocation_id, str) or not invocation_id: + msg = "Offload hook request has no invocation id." + raise RuntimeError(msg) + logger.debug( + "Offload round %d fulfilling hook invocation %s", + round_index, + invocation_id, + ) + if round_index == _OFFLOAD_MAX_RESUME_ROUNDS: + # The extra iteration exists to POST the last fulfillment and + # read the result, not to answer one more hook. Without this the + # loop fulfills 33 hooks and then reports "after 32 rounds". + break + hook_responses[invocation_id] = await _await_offload_step( + fulfill_hook(request), + graph=graph, + thread_id=thread_id, + operation_id=operation_id, + ) + # The server only re-requests an invocation id it has not been given, so + # exhaustion means this many *distinct* ids. That is either genuinely + # many hooks or an unstable invocation-id derivation server-side; log the + # ids so the two are distinguishable, and do not assert a cause in the + # user-facing message. + logger.warning( + "Offload exceeded %d hook rounds; fulfilled invocation ids: %s", + _OFFLOAD_MAX_RESUME_ROUNDS, + sorted(hook_responses), + ) + msg = ( + f"Offload did not complete after {_OFFLOAD_MAX_RESUME_ROUNDS} hook " + "rounds. Check the server log for the hook invocations it requested." + ) + raise RuntimeError(msg) + async def astream( self, input: dict | Any, # noqa: A002, ANN401 diff --git a/libs/code/deepagents_code/configurable_model.py b/libs/code/deepagents_code/configurable_model.py index d618d082e0..c88150b9d6 100644 --- a/libs/code/deepagents_code/configurable_model.py +++ b/libs/code/deepagents_code/configurable_model.py @@ -293,7 +293,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 @@ -366,15 +366,16 @@ def _get_context(request: ModelRequest) -> CLIContextSchema | None: def _model_spec_from_model(model: BaseChatModel) -> str | None: """Return a resumable `provider:model` spec for a model object.""" - provider = _get_ls_provider(model) model_name = get_model_identifier(model) - if provider and model_name: - return f"{provider}:{model_name}" - from deepagents_code.config import settings settings_provider = settings.model_provider or "" settings_model = settings.model_name or "" + if settings_provider and settings_model and model_name == settings_model: + return f"{settings_provider}:{settings_model}" + provider = _get_ls_provider(model) + if provider and model_name: + return f"{provider}:{model_name}" if settings_provider and settings_model: return f"{settings_provider}:{settings_model}" return None diff --git a/libs/code/deepagents_code/cost_tracking.py b/libs/code/deepagents_code/cost_tracking.py index 67ef435fc7..2ff8bf12aa 100644 --- a/libs/code/deepagents_code/cost_tracking.py +++ b/libs/code/deepagents_code/cost_tracking.py @@ -1,8 +1,12 @@ """Estimate and persist cumulative model cost for each thread. -The graph owns the durable total. `CostTrackingMiddleware` is the only writer of -`_session_cost_usd`, so each cost update rides the model checkpoint and works for -local, headless, and remote graph execution without a client-side state update. +The graph owns the durable total. `CostTrackingMiddleware` writes ordinary graph +deltas, while `prepare_operation_cost` gives server-owned operations a +rollback-safe delta to commit with their state update. Each cost update therefore +rides a graph checkpoint and works for local, headless, and remote execution +without a client-side state update -- the middleware also runs in local +in-process agents, where there is no server and the delta rides the local +checkpoint. The client is a reader: it renders the streamed total and never maintains its own lifetime figure. @@ -55,7 +59,7 @@ from collections import OrderedDict from collections.abc import Mapping, Sequence from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Annotated, Any, NotRequired, TypedDict from langchain.agents.middleware.types import ( @@ -1966,6 +1970,142 @@ def _restore_recorded_costs( return True +@dataclass(slots=True) +class PreparedOperationCost: + """A priced side-operation delta pending checkpoint persistence. + + Use exactly once: every prepare must be either committed (its `update` + persisted) or rolled back. `_drain_recorded_costs` is destructive, so a + prepare that is neither deletes that spend from the thread's lifetime total + permanently, and nothing can detect the loss afterwards. + + `_settled` is `init=False` so a caller cannot construct a pre-neutralized + instance whose `rollback` is already a no-op. The class stays unfrozen only + because marking settlement on a frozen dataclass needs + `object.__setattr__`, which this project's lint rules reject. + """ + + thread_id: str + records: list[_ModelCallRecord] + delta_usd: float + _settled: bool = field(default=False, init=False) + + @property + def update(self) -> dict[str, float]: + """Additive checkpoint update for this prepared charge.""" + return {"_session_cost_usd": self.delta_usd} if self.delta_usd > 0 else {} + + def rollback(self) -> None: + """Return claimed records to the recorder when no charge was committed. + + `_drain_recorded_costs` **removes** the entries from the process-wide + recorder, so a prepare that is neither committed nor rolled back deletes + that spend from the thread's lifetime total permanently. Restoring lets + the next drain price it again. + + Only call this when the checkpoint write did not land: restoring records + for a write that *did* commit double-charges the thread on the next + drain. Repeat calls are ignored for the same reason. + """ + if self._settled: + return + self._settled = True + if not _restore_recorded_costs(self.thread_id, self.records): + logger.warning( + "Could not restore %d operation cost record(s) after a failed " + "checkpoint update; $%.6f is dropped from the thread total", + len(self.records), + self.delta_usd, + ) + + def commit(self) -> None: + """Mark the prepared charge as persisted. + + Records nothing: the delta reaches the thread through `update`, which + the caller writes to the checkpoint. This only settles the instance so + an abandoned prepare can be told apart from a completed one. + """ + self._settled = True + + def __del__(self) -> None: + """Warn when a prepare was abandoned without settling. + + The drain is destructive, so an instance collected while unsettled has + silently deleted its spend from the thread's lifetime total. Committing + and rolling back are both observable; only the leak was not, which is + the one case the class docstring calls unrecoverable. + """ + if not self._settled: + logger.warning( + "Operation cost prepare for thread %s was abandoned without " + "commit or rollback; $%.6f across %d record(s) is lost from the " + "thread total", + self.thread_id, + self.delta_usd, + len(self.records), + ) + + +def prepare_operation_cost( + state: CostState, + thread_id: str, +) -> PreparedOperationCost: + """Price model calls made by a server operation without committing them. + + The caller must persist `PreparedOperationCost.update` atomically with the + operation state, or call `rollback()` if that write fails or is abandoned. + A prepare with a zero delta still consumes its records, so an abandoned + prepare must roll back even when it has nothing to write. + + Args: + state: Current thread state used for model/provider fallback metadata. + thread_id: Thread that owns the side-operation model calls. + + Returns: + Prepared additive cost delta and the claimed recorder entries. + + """ + records = _drain_recorded_costs(thread_id) + fallback = _checkpointed_model_spec(state) + delta_usd = 0.0 + try: + for record in records: + cost_usd = estimate_cost( + record.usage_metadata, + *_pricing_target(record.model_name, record.provider, fallback), + ) + if cost_usd is None: + # Matches `CostTrackingMiddleware`: silently omitting an + # unpriceable call leaves the total quietly short, so name what + # could not be priced. + logger.warning( + "No pricing for operation model call %r (provider %r); " + "its cost is omitted from the thread total", + record.model_name, + record.provider, + ) + continue + delta_usd += cost_usd + except BaseException: + if not _restore_recorded_costs(thread_id, records): + # `_restore_recorded_costs` returns `bool` so callers can report a + # failed restore; the other two call sites both log. Without this a + # pricing crash could drop the drained spend with no record. + logger.warning( + "Could not restore %d drained cost record(s) for thread %s " + "after a pricing failure; that spend is lost from the " + "lifetime total", + len(records), + thread_id, + ) + raise + return PreparedOperationCost( + thread_id=thread_id, + records=records, + delta_usd=delta_usd, + ) + + class _CostTransfer(TypedDict): """One completed nested total addressed to its owning parent graph.""" diff --git a/libs/code/deepagents_code/hooks/server_middleware.py b/libs/code/deepagents_code/hooks/server_middleware.py index a3961c566f..949da867ed 100644 --- a/libs/code/deepagents_code/hooks/server_middleware.py +++ b/libs/code/deepagents_code/hooks/server_middleware.py @@ -9,9 +9,11 @@ import hashlib import json +import logging import time from collections.abc import Mapping, Sequence from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass, field, replace from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Annotated, Any, Literal, NotRequired, TypeGuard, cast @@ -90,6 +92,63 @@ _COMPACT_TOOL_NAME = "compact_conversation" _INVOCATION_NAMESPACE = UUID("f2896d18-cf2a-4e7d-b11a-d5b10fc0e335") + +class HookTransportInterruptError(BaseException): + """Carry a hook request across a non-graph server operation boundary. + + Derives from `BaseException`, not `Exception`, for the same reason + `asyncio.CancelledError` does: it is a control signal that must reach the + HTTP boundary intact. The compaction chain it crosses is lined with broad + `except Exception` handlers, any of which would otherwise turn a resumable + hook request into a permanent `"failed"` result. + """ + + def __init__(self, request: HookInvocationRequest) -> None: + """Initialize the transport interrupt. + + Args: + request: Hook invocation the client must fulfill. + """ + super().__init__(str(request.invocation_id)) + self.request = request + + +logger = logging.getLogger(__name__) + +_HOOK_RESPONSES: ContextVar[Mapping[str, object] | None] = ContextVar( + "deepagents_code_hook_responses", + default=None, +) + + +@contextmanager +def operation_hook_responses( + responses: Mapping[str, object], +) -> Iterator[None]: + """Serve hook responses while a server operation replays from the top. + + Args: + responses: Resume payloads keyed by deterministic hook invocation ID. + """ + token = _HOOK_RESPONSES.set(responses) + try: + yield + finally: + _HOOK_RESPONSES.reset(token) + + +def _in_server_operation() -> bool: + """Report whether hooks are running under a non-graph server operation. + + `None` means graph mode; an empty mapping means operation mode with no + answers accumulated yet, which is why this cannot be a truthiness check. + + Returns: + `True` when the caller is inside `operation_hook_responses`. + """ + return _HOOK_RESPONSES.get() is not None + + type PreToolBehavior = Literal["allow", "deny", "none"] _DEFAULT_DENY_REASON = "Blocked by PreToolUse hook" @@ -861,7 +920,20 @@ def _invoke_hook( invocation=HookInvocation(context=context, event=event), deadline=datetime.now(UTC) + deadline, ) - raw = interrupt(build_hook_interrupt_payload(request)) + operation_responses = _HOOK_RESPONSES.get() + if operation_responses is None: + raw = interrupt(build_hook_interrupt_payload(request)) + else: + # Operation mode: `interrupt()` is unusable outside a Pregel task, so a + # request the client has not answered yet is raised out to the HTTP + # boundary instead. Because the operation re-executes from the top on + # every resume round, an already-answered invocation is replayed from + # this mapping rather than re-invoked -- that is what makes an operation + # with several hooks terminate instead of looping forever. + key = str(request.invocation_id) + if key not in operation_responses: + raise HookTransportInterruptError(request) + raw = operation_responses[key] try: response = parse_hook_resume_value( raw, @@ -871,6 +943,16 @@ def _invoke_hook( except ValidationError: # Only shape errors degrade to a neutral decision. A plain `ValueError` # means the client answered a different request, so it stays fatal. + # + # Log it too: the diagnostic is only rendered by the client-side hook + # presenter, and the offload operation reads just the pre-tool channel + # from this update, so on that path the diagnostic is dropped and the + # hook is silently ignored. + logger.warning( + "Malformed hook resume value for invocation %s; treating it as no decision", + request.invocation_id, + exc_info=True, + ) diagnostic = HookDiagnostic( code="invalid_resume", severity="warning", @@ -1091,6 +1173,24 @@ def _ask_permission_via_hitl( Returns: A deny ToolMessage when the user rejects, otherwise `None` to proceed. """ + if _in_server_operation(): + # `interrupt()` is only usable inside a Pregel task: it reaches into the + # run's scratchpad, which a server operation's fabricated config has no + # equivalent of. Deny with an actionable reason instead of raising a + # `KeyError` on an internal LangGraph config key. The operation + # transport carries hook *invocations*, not HITL review requests, so + # there is no channel to prompt the user on here. + return _denied_tool_message( + call, + PermissionEffect( + behavior="deny", + reason=( + f"PreToolUse returned `ask` for {call.name}, which cannot " + "prompt for approval during a server-side operation such as " + "/offload. Return `allow` or `deny` for this tool instead." + ), + ), + ) description = permission.reason or "PreToolUse hook requested approval" response = interrupt( HITLRequest( diff --git a/libs/code/deepagents_code/offload_api.py b/libs/code/deepagents_code/offload_api.py new file mode 100644 index 0000000000..52df15aa70 --- /dev/null +++ b/libs/code/deepagents_code/offload_api.py @@ -0,0 +1,932 @@ +"""dcode-owned HTTP boundary for server-side thread offload.""" + +from __future__ import annotations + +import asyncio +import logging +from collections import OrderedDict +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Literal, cast +from weakref import WeakValueDictionary + +from langchain_core.messages import convert_to_messages +from langchain_core.runnables.config import var_child_runnable_config +from langgraph.runtime import ExecutionInfo, Runtime +from langgraph_sdk import get_client +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +from deepagents_code._cli_context import CLIContextSchema +from deepagents_code.cost_tracking import prepare_operation_cost +from deepagents_code.hooks.interrupt import build_hook_interrupt_payload +from deepagents_code.hooks.server_middleware import ( + HookTransportInterruptError, + operation_hook_responses, +) +from deepagents_code.offload_middleware import ( + OffloadStateUpdate, + _archive_lock, + unchanged_offload_result, +) +from deepagents_code.server_graph import get_server_runtime + +if TYPE_CHECKING: + from langchain_core.runnables import RunnableConfig + from starlette.requests import Request + + from deepagents_code.cost_tracking import PreparedOperationCost + from deepagents_code.offload_middleware import ( + OffloadExecution, + OffloadResponse, + _OffloadState, + ) + +logger = logging.getLogger(__name__) + +_WRITABLE_STATE_CHANNELS = frozenset(OffloadStateUpdate.__annotations__) +"""Checkpoint channels a server-owned offload may write. + +Derived from `OffloadStateUpdate` so the runtime guard and the type cannot +drift: adding a channel to the type is the only way to permit writing it. +""" +_OFFLOADABLE_THREAD_STATUSES = frozenset({"idle", "error"}) +"""Thread statuses that hold no in-flight work, so offload may proceed. + +`error` is included deliberately. A run that raises anything other than an +interrupt or rollback leaves the thread row on `error` until the *next* run +completes, and `RemoteAgent.aensure_thread` uses `if_exists="do_nothing"`, so +it does not clear it. Excluding `error` would refuse `/offload` for the whole +window after a failed turn -- exactly when a user reaches for it to recover +from a context overflow. Quiescence is checked separately against the +checkpoint's `next`/`tasks`/`interrupts`, which still catches an errored run +that left a pending node. +""" +_thread_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary() +type _OperationKey = tuple[str, str] +type _OperationOutcome = Literal["cancelled", "finished"] +_active_operations: dict[_OperationKey, asyncio.Task[object]] = {} +_operation_outcomes: OrderedDict[_OperationKey, _OperationOutcome] = OrderedDict() +_MAX_OPERATION_OUTCOMES = 1024 +"""Bound completed/cancelled ids retained to close request/cancel races.""" + +# One client for the process. `get_client` builds a fresh `httpx.AsyncClient` +# (with its own connection pool) per call and exposes no close hook we own, so +# calling it per request -- and this route runs once per hook resume round -- +# would leak a pool for the lifetime of the server. +_client: Any = None + + +def _thread_client() -> Any: # noqa: ANN401 # untyped LangGraph SDK client + """Return the process-wide in-process LangGraph SDK client.""" + global _client # noqa: PLW0603 # module-level singleton by design + if _client is None: + _client = get_client(url=None, api_key=None) + return _client + + +def _thread_lock(thread_id: str) -> asyncio.Lock: + """Return one live lock per thread without retaining inactive threads.""" + lock = _thread_locks.get(thread_id) + if lock is None: + lock = asyncio.Lock() + _thread_locks[thread_id] = lock + return lock + + +def _remember_operation(key: _OperationKey, outcome: _OperationOutcome) -> None: + """Retain a bounded terminal outcome for late or reordered cancellation.""" + _operation_outcomes[key] = outcome + _operation_outcomes.move_to_end(key) + while len(_operation_outcomes) > _MAX_OPERATION_OUTCOMES: + _operation_outcomes.popitem(last=False) + + +def _register_operation(key: _OperationKey) -> str | None: + """Register the current request task. + + Returns: + A refusal reason, or `None` when registration succeeds. + """ + outcome = _operation_outcomes.get(key) + if outcome == "cancelled": + return "The offload operation was cancelled." + if outcome == "finished": + return "The offload operation already finished." + if key in _active_operations: + return "This offload operation already has an active request." + task = asyncio.current_task() + if task is None: + return "The server could not register the offload operation." + _active_operations[key] = cast("asyncio.Task[object]", task) + return None + + +def _finish_operation(key: _OperationKey, outcome: _OperationOutcome | None) -> None: + """Release an active round and optionally retain its terminal outcome.""" + task = asyncio.current_task() + if _active_operations.get(key) is task: + _active_operations.pop(key, None) + if outcome is not None: + _remember_operation(key, outcome) + + +class _OffloadConflictError(RuntimeError): + """The thread changed or became active during an offload attempt.""" + + +class _OffloadUnavailableError(RuntimeError): + """The server runtime could not be built, so no operation can run. + + `get_server_runtime` converts a construction failure into a startup-error + marker and `sys.exit(1)`. That barrier was written for the `langgraph.json` + graph factory, where exiting is right; reached from a request handler it + would kill the server process mid-request, and `SystemExit` is a + `BaseException`, so the route's own handler could not turn it into a + response. Server-owned offload cannot run without that runtime, so report + the condition instead. + """ + + +class _OffloadIndeterminateError(RuntimeError): + """The state write may or may not have landed; the outcome is unknown. + + Raised only when the checkpoint write itself failed *and* a follow-up read + shows the thread advanced anyway, so the operation cannot honestly claim + either that it committed or that it did not. + """ + + +# Context fields validated at this boundary (everything else in +# `CLIContextSchema` — turn ids and the approval-mode key, for example — drives +# interactive-run machinery this operation never touches). Validated here so a +# malformed client request fails with a 422 naming the field instead of a 500 +# deep in model resolution or hook dispatch. Validated is not the same as read: +# `classifier_model` is checked for shape but feeds only auto mode, which this +# operation never enters. +_CONTEXT_STR_OR_NONE_FIELDS = ( + "model", + "classifier_model", + "approval_mode", + "thread_id", + "hooks_snapshot_id", + "prompt_id", +) +_CONTEXT_DICT_FIELDS = ("model_params", "profile_overrides") + +_TRANSPORT_MODEL_PARAM_KEYS = frozenset( + { + # Endpoint selection. + "base_url", + "api_base", + "openai_api_base", + "anthropic_api_url", + "azure_endpoint", + "azure_openai_api_base", + "api_endpoint", + # Proxy routing. + "openai_proxy", + "anthropic_proxy", + "proxy", + "proxies", + # Outbound transport injection: these keys hand whole HTTP clients, + # transports, or header maps to the model constructor. + "http_client", + "http_async_client", + "transport", + "default_headers", + "custom_headers", + } +) +"""`model_params` keys stripped from client-supplied offload context. + +`create_model` merges these params verbatim into the model constructor +(`offload_middleware._summarization_for_runtime`), and the summarizer's +outbound provider calls carry the server's credentials. A client that sets an +endpoint/proxy/transport key therefore chooses where those credentials are +sent. The in-process paths trust `model_params` (the user supplied them +through their own flags and config); the HTTP boundary does not -- the dev +server accepts connections from any local process, so the request's model +selection must not extend to its network plumbing. + +A backstop, not the primary control. `_checkpoint_model_context` discards the +request's `model` and `model_params` outright and substitutes the checkpointed +values, so a client-supplied endpoint cannot reach `create_model` even without +this filter. It is kept for the case that control cannot cover: a future path +that resolves a model before, or instead of, reading the checkpoint. Treat a +warning from here as a client sending params it should not, not as a breach. + +Denylist rather than allowlist: `create_model` serves arbitrary providers, so a +fixed allowlist would silently drop legitimate provider-specific params. +""" + + +def _strip_transport_model_params(context: dict[str, Any]) -> dict[str, Any]: + """Return a context copy with endpoint/transport model params removed. + + Args: + context: The request's already type-checked `context` object. + + Returns: + The same dict when `model_params` holds no stripped keys, otherwise a + shallow copy whose `model_params` omits them. + """ + params = context.get("model_params") + if not isinstance(params, dict): + return context + stripped = { + key: value + for key, value in params.items() + if key not in _TRANSPORT_MODEL_PARAM_KEYS + } + if len(stripped) == len(params): + return context + # Logged, not silent: dropping these changes where the summarizer's + # credentialed calls go, so a user whose gateway config is being ignored has + # something to find. Key names only -- the values are endpoints and headers. + logger.warning( + "Dropped transport key(s) %s from offload model_params; a server-owned " + "operation does not accept a client-chosen endpoint", + sorted(set(params) - set(stripped)), + ) + return {**context, "model_params": stripped} + + +def _validate_context(context: dict[str, Any]) -> None: + """Check the context fields the offload operation consumes. + + Only the listed keys are type-checked; unknown keys pass through so a + newer client can keep talking to this server version. + + Args: + context: The request's `context` object. + + Raises: + TypeError: If a consumed field has the wrong type, naming the field. + """ + for key in _CONTEXT_STR_OR_NONE_FIELDS: + value = context.get(key) + if value is not None and not isinstance(value, str): + msg = f"context.{key} must be a string or null, got {type(value).__name__}." + raise TypeError(msg) + for key in _CONTEXT_DICT_FIELDS: + value = context.get(key) + if value is not None and not isinstance(value, dict): + msg = f"context.{key} must be an object, got {type(value).__name__}." + raise TypeError(msg) + limit = context.get("model_context_limit") + # bool is an int subclass, so exclude it explicitly: JSON `true` is not a + # token limit. + if limit is not None and (isinstance(limit, bool) or not isinstance(limit, int)): + msg = ( + "context.model_context_limit must be an integer or null, " + f"got {type(limit).__name__}." + ) + raise TypeError(msg) + auto_approve = context.get("auto_approve") + if auto_approve is not None and not isinstance(auto_approve, bool): + msg = ( + f"context.auto_approve must be a boolean or null, " + f"got {type(auto_approve).__name__}." + ) + raise TypeError(msg) + events = context.get("hooks_server_events") + if events is not None and ( + not isinstance(events, list) + or any(not isinstance(event, str) for event in events) + ): + msg = "context.hooks_server_events must be a list of strings or null." + raise TypeError(msg) + + +def _checkpoint_id(state: Mapping[str, object]) -> str: + checkpoint = state.get("checkpoint") + value = checkpoint.get("checkpoint_id") if isinstance(checkpoint, Mapping) else None + if not isinstance(value, str) or not value: + msg = "The thread has no checkpoint to offload." + raise _OffloadConflictError(msg) + return value + + +def _operation_payload( + payload: object, +) -> tuple[str, dict[str, Any], dict[str, object]]: + """Validate the narrow client-to-operation request shape. + + Args: + payload: Decoded request JSON. + + Returns: + Operation id, runtime context, and accumulated hook responses. + + Raises: + TypeError: If the payload or a structured field has the wrong shape. + """ + if not isinstance(payload, dict): + msg = "Offload request must be a JSON object." + raise TypeError(msg) + operation_id = payload.get("operation_id") + context = payload.get("context") + responses = payload.get("hook_responses", {}) + if not isinstance(operation_id, str) or not operation_id: + msg = "operation_id must be a non-empty string." + raise TypeError(msg) + if not isinstance(context, dict): + msg = "context must be a JSON object." + raise TypeError(msg) + if not isinstance(responses, dict): + msg = "hook_responses must be a JSON object." + raise TypeError(msg) + validated_context = {str(key): value for key, value in context.items()} + _validate_context(validated_context) + return ( + operation_id, + _strip_transport_model_params(validated_context), + {str(key): value for key, value in responses.items()}, + ) + + +def _hydrate_state(values: object) -> _OffloadState: + """Hydrate serialized checkpoint messages for the compaction service. + + Args: + values: State values returned by LangGraph Server. + + Returns: + A shallow state copy containing LangChain message objects. + + Raises: + TypeError: If the server returns an unexpected state shape. + """ + if not isinstance(values, dict): + msg = "LangGraph returned non-object thread state." + raise TypeError(msg) + state = dict(values) + messages = state.get("messages", []) + if not isinstance(messages, list): + msg = "LangGraph returned a non-list messages channel." + raise TypeError(msg) + state["messages"] = convert_to_messages(messages) + + # LangGraph serializes the summary stored inside the private event channel + # independently of the top-level `messages` channel. The summarization SDK + # prepends it to the effective conversation, so it must be a message object + # too rather than the serialized dict returned by the thread API. + event = state.get("_summarization_event") + if isinstance(event, Mapping) and "summary_message" in event: + hydrated_event = dict(event) + summary_message = hydrated_event["summary_message"] + hydrated_event["summary_message"] = convert_to_messages([summary_message])[0] + state["_summarization_event"] = hydrated_event + return cast("_OffloadState", state) + + +def _checkpoint_model_context( + context: dict[str, Any], state: Mapping[str, object] +) -> dict[str, Any]: + """Replace request model selection with server-checkpointed values. + + The client still supplies hook and profile context, but it cannot choose + the model's outbound transport for this server-owned operation. Successful + agent turns checkpoint the resolved model spec and the runtime overrides + they actually used, so those values preserve trusted launch/model-switch + settings such as a private `base_url` without accepting an arbitrary + offload request's endpoint override. + + Args: + context: Validated request context. + state: Server-read checkpoint values for the target thread. + + Returns: + Context using checkpointed model settings, or no model override when + the thread predates model checkpointing so the startup summarizer + is reused. + """ + trusted = dict(context) + trusted.pop("model", None) + trusted.pop("model_params", None) + model = state.get("_model_spec") + params = state.get("_model_params") + if isinstance(model, str) and model: + trusted["model"] = model + if isinstance(params, dict): + trusted["model_params"] = dict(params) + return trusted + + +async def _require_idle_thread(client: Any, thread_id: str) -> None: # noqa: ANN401 + """Reject offload while LangGraph reports an active thread. + + Args: + client: In-process LangGraph SDK client. + thread_id: Thread being compacted. + + Raises: + _OffloadConflictError: If the thread has work in flight, or is not + registered on the server at all. + """ + from langgraph_sdk.errors import NotFoundError + + try: + thread = await client.threads.get(thread_id) + except NotFoundError as exc: + # Checkpoint persistence and HTTP thread registration are separate on + # the dev server, so a thread can hold on-disk state while its live row + # is absent (see `RemoteAgent.aensure_thread`). The client registers + # before requesting the operation; reaching here means it could not, so + # name the condition instead of letting a 404 become an opaque 500. + msg = ( + "This thread is not registered on the server; send a message " + "before offloading." + ) + raise _OffloadConflictError(msg) from exc + if thread.get("status") not in _OFFLOADABLE_THREAD_STATUSES: + msg = "Cannot offload while the thread has an active or interrupted run." + raise _OffloadConflictError(msg) + + +async def _write_landed( + client: Any, # noqa: ANN401 # untyped LangGraph SDK client + thread_id: str, + checkpoint_id: str, +) -> Literal["advanced", "unchanged", "unreadable"]: + """Classify a failed `update_state` against the checkpoint we read. + + A new checkpoint means the write most likely applied despite the error. A + concurrent run could also have advanced the thread, so this is a bias, not a + proof -- it biases toward keeping cost records claimed (understating spend at + worst) over restoring them (which would double-charge). + + `unreadable` is reported separately from `advanced` so the caller can say + which one happened. Both keep the records claimed, but only `advanced` has + evidence the write landed; conflating them would log a thread advance that + was never observed. + + Args: + client: In-process LangGraph SDK client. + thread_id: Thread that was being compacted. + checkpoint_id: Checkpoint the operation read and validated against. + + Returns: + `advanced` if the checkpoint changed, `unchanged` if it did not, or + `unreadable` if the thread could not be read back. + """ + try: + current = await client.threads.get_state(thread_id) + except BaseException: + # `BaseException`, not `Exception`: this runs inside the caller's + # settlement handler, so an escape here -- a `CancelledError` from a + # disconnect or a shutdown re-delivering cancellation while that handler + # unwinds -- would skip the rollback entirely and delete the drained + # cost records from the thread's lifetime total with no trace. + # + # An unreadable thread cannot rule the write out, so stay on the + # conservative side and treat the outcome as indeterminate. + logger.exception( + "Could not read thread %s back to classify a failed offload write", + thread_id, + ) + return "unreadable" + return "advanced" if _checkpoint_id(current) != checkpoint_id else "unchanged" + + +async def _commit_state_update( + client: Any, # noqa: ANN401 # untyped LangGraph SDK client + thread_id: str, + checkpoint_id: str, + update: dict[str, Any], + prepared: PreparedOperationCost, +) -> None: + """Persist the summary reservation and settle its claimed model cost. + + Raises: + _OffloadIndeterminateError: If the write failed after the thread + advanced and its outcome cannot be determined. + """ + try: + await client.threads.update_state(thread_id, update) + except BaseException as exc: + outcome = await _write_landed(client, thread_id, checkpoint_id) + if outcome != "unchanged": + if outcome == "advanced": + logger.exception( + "Offload state write for thread %s failed after the thread " + "advanced past checkpoint %s; keeping %d cost record(s) " + "claimed", + thread_id, + checkpoint_id, + len(prepared.records), + ) + else: + # Distinct from `advanced`: no thread advance was observed, so + # the write may never have landed. Naming the amount makes an + # otherwise undetectable loss auditable. + logger.exception( + "Offload state write for thread %s failed and the thread " + "could not be read back; keeping %d cost record(s) claimed " + "to avoid double-charging, so $%.6f may be lost from the " + "thread total", + thread_id, + len(prepared.records), + prepared.delta_usd, + ) + # Deliberately settled rather than rolled back: the delta is + # treated as persisted, so restoring the records would double-charge + # the next drain. + prepared.commit() + if isinstance(exc, asyncio.CancelledError): + raise + msg = ( + "Offload compacted the conversation but could not confirm " + "the state write. Run /context to check whether the " + "conversation was compacted before offloading again." + ) + raise _OffloadIndeterminateError(msg) from None + logger.warning( + "Offload state write for thread %s failed with no thread advance; " + "restoring %d cost record(s)", + thread_id, + len(prepared.records), + ) + prepared.rollback() + raise + prepared.commit() + + +async def _archive_path_landed( + client: Any, # noqa: ANN401 # untyped LangGraph SDK client + thread_id: str, + path: str, +) -> bool | None: + """Check whether the follow-up checkpoint links the completed archive. + + Returns: + `True` when linked, `False` when confirmed absent, or `None` when the + checkpoint could not be read. + """ + try: + current = await client.threads.get_state(thread_id) + except BaseException: + logger.exception( + "Could not verify archive-path update for thread %s", thread_id + ) + return None + values = current.get("values") + event = values.get("_summarization_event") if isinstance(values, Mapping) else None + return isinstance(event, Mapping) and event.get("file_path") == path + + +async def _commit_deferred_archive( + client: Any, # noqa: ANN401 # untyped LangGraph SDK client + thread_id: str, + checkpoint_id: str, + execution: OffloadExecution, + update: dict[str, Any], + prepared: PreparedOperationCost, +) -> None: + """Reserve summary state, then append and link its archive transactionally. + + Raises: + _OffloadIndeterminateError: If the archive was written but its + checkpoint link cannot be read back. + """ + archive = execution.archive + if archive is None: + await _commit_state_update(client, thread_id, checkpoint_id, update, prepared) + return + async with _archive_lock(archive.session_id): + await _commit_state_update(client, thread_id, checkpoint_id, update, prepared) + try: + append = await archive.write() + except Exception: + logger.exception( + "/offload reserved its summary but the archive append failed" + ) + return + if append is None: + logger.error("/offload reserved its summary but the archive append failed") + return + event = archive.update(append.path)["_summarization_event"] + try: + await client.threads.update_state( + thread_id, {"_summarization_event": event} + ) + except BaseException as exc: + landed = await _archive_path_landed(client, thread_id, append.path) + if landed is True: + execution.result["archive_path"] = append.path + if isinstance(exc, asyncio.CancelledError): + raise + return + if landed is False: + await append.rollback() + if isinstance(exc, asyncio.CancelledError): + raise + logger.exception( + "Archive link failed for thread %s; restored prior archive", + thread_id, + ) + return + msg = "Offload wrote its archive but could not confirm the archive link." + raise _OffloadIndeterminateError(msg) from None + execution.result["archive_path"] = append.path + + +async def _join_task_deferring_cancellation[T]( + task: asyncio.Task[T], +) -> asyncio.CancelledError | None: + """Join a settlement task while retaining the first cancellation edge. + + Returns: + The cancellation to re-raise after settlement, or `None`. + """ + cancellation: asyncio.CancelledError | None = None + while not task.done(): + try: + await asyncio.wait((task,)) + except asyncio.CancelledError as exc: + cancellation = cancellation or exc + return cancellation + + +async def _execute_offload( + thread_id: str, + *, + operation_id: str, + context: dict[str, Any], + hook_responses: dict[str, object], +) -> OffloadResponse: + """Execute and commit one server-owned offload attempt. + + Args: + thread_id: LangGraph thread to compact. + operation_id: Opaque client-generated attempt identity. + context: Runtime model and hooks context. + hook_responses: Accumulated hook replies keyed by invocation id. + + Returns: + A complete result or a hook request that must be answered. + + Raises: + TypeError: If `thread_id` is empty. + _OffloadConflictError: If the thread is active or changes before commit. + _OffloadUnavailableError: If the server runtime cannot be built, so no + operation can run. + RuntimeError: If the operation attempts to write conversation messages. + """ + if not thread_id: + msg = "thread_id path parameter must be non-empty." + raise TypeError(msg) + client = _thread_client() + async with _thread_lock(thread_id): + await _require_idle_thread(client, thread_id) + before = await client.threads.get_state(thread_id) + if before.get("next") or before.get("tasks") or before.get("interrupts"): + msg = "Cannot offload a thread with pending graph work." + raise _OffloadConflictError(msg) + + state = _hydrate_state(before.get("values")) + if not state.get("messages"): + # An empty thread is "nothing to offload", not a failure. Answer it + # here: `_checkpoint_id` below rejects a thread with no checkpoint, + # so without this the graceful `empty` branch in + # `OffloadOperation.execute` is unreachable over HTTP and the user + # is told the operation failed. + return { + "status": "complete", + "result": unchanged_offload_result("empty", messages=0, tokens=0), + } + + checkpoint_id = _checkpoint_id(before) + context = _checkpoint_model_context(context, state) + context["thread_id"] = thread_id + namespace = f"dcode_offload:{operation_id}" + info = ExecutionInfo( + checkpoint_id=checkpoint_id, + checkpoint_ns=namespace, + task_id=operation_id, + thread_id=thread_id, + run_id=operation_id, + ) + try: + server = await get_server_runtime() + except SystemExit as exc: + msg = ( + "The server could not build its agent runtime, so /offload is " + "unavailable. Check the server log for the startup failure." + ) + raise _OffloadUnavailableError(msg) from exc + runtime = Runtime[CLIContextSchema]( + context=cast("CLIContextSchema", context), + store=getattr(server.agent, "store", None), + execution_info=info, + ) + config = cast( + "RunnableConfig", + { + "configurable": { + "thread_id": thread_id, + "checkpoint_id": checkpoint_id, + "checkpoint_ns": namespace, + "run_id": operation_id, + } + }, + ) + token = var_child_runnable_config.set(config) + try: + with operation_hook_responses(hook_responses): + execution = await server.offload.execute(state, runtime) + except HookTransportInterruptError as interrupt: + return { + "status": "interrupt", + "request": build_hook_interrupt_payload(interrupt.request), + } + finally: + var_child_runnable_config.reset(token) + + await _require_idle_thread(client, thread_id) + current = await client.threads.get_state(thread_id) + if _checkpoint_id(current) != checkpoint_id: + # Compaction already ran, so the summarizer model call has been made + # and paid for. Name the discarded work: the records stay in the + # recorder and would otherwise be swept into an unrelated later turn + # with no trace of where they came from. + logger.warning( + "Discarding a completed offload for thread %s: the thread " + "advanced past checkpoint %s while compaction was running, so " + "the summary (and its model spend) cannot be committed", + thread_id, + checkpoint_id, + ) + msg = ( + "The thread changed while offload was running; no state was committed." + ) + raise _OffloadConflictError(msg) + + prepared = prepare_operation_cost(state, thread_id) + update: dict[str, Any] = {**execution.update, **prepared.update} + if forbidden := set(update) - _WRITABLE_STATE_CHANNELS: + # A security boundary, not a defensive assertion: this route commits + # to the latest checkpoint rather than the one it read, so a + # `messages` write here would be unattributed to any run and could + # clobber messages a concurrent run appended in that window. See + # THREAT_MODEL.md (TB10/DF27) before relaxing this. + # + # Checked as an allowlist against `OffloadStateUpdate` rather than + # for `messages` alone, so the runtime guard enforces the same + # invariant the type states instead of a subset of it: a future + # merge that adds any other channel is refused here too. + msg = ( + "Server offload operations may not write " + f"{sorted(forbidden)} to the checkpoint." + ) + prepared.rollback() + raise RuntimeError(msg) + if not update: + # Nothing to persist, but `prepare_operation_cost` already drained + # the recorder. Returning without rolling back would delete that + # spend from the thread's lifetime total (the drain is destructive). + prepared.rollback() + return {"status": "complete", "result": execution.result} + commit = asyncio.create_task( + _commit_deferred_archive( + client, + thread_id, + checkpoint_id, + execution, + update, + prepared, + ) + ) + cancellation = await _join_task_deferring_cancellation(commit) + commit.result() + if cancellation is not None: + raise cancellation + return {"status": "complete", "result": execution.result} + + +async def offload(request: Request) -> JSONResponse: + """Handle one thread offload or hook-resume round. + + Request body: `operation_id` (non-empty string, stable across the rounds of + one attempt), `context` (runtime model and Hooks v2 context), and + `hook_responses` (replies accumulated so far, keyed by invocation id). + The thread comes from the path, never the body. + + A round either completes or returns a hook request to answer. There is no + suspended coroutine server-side: a resume round **re-executes the operation + from the top**, and `_invoke_hook` replays already-answered invocations from + `hook_responses` instead of raising again. That is what makes the loop + terminate, and it is why the dispatched call's id must be stable across + rounds (`_forced_offload_call_id`). + + Status codes, and what each means for whether state committed: + + - 200 -- completed, or a resumable hook request; no state written in the + latter case. + - 422 -- malformed request, named by field. Nothing ran. + - 409 -- thread conflict: active, interrupted, holding pending graph work, + unregistered, carrying no checkpoint to offload, or advanced past the + checkpoint read. Nothing committed. + - 503 -- the server runtime could not be built. Nothing ran. + - 500 -- either an indeterminate write (compaction happened and the commit + cannot be confirmed; the detail says so and is user-actionable) or an + unexpected server fault. + + Invariants this boundary owns: it reads and hydrates checkpoint state itself, + it commits only the channels `OffloadStateUpdate` permits and refuses any + `messages` write outright, and it settles the drained cost records on every + exit path. See THREAT_MODEL.md C18 (TB10/DF27). + + Returns: + JSON operation response. + + Raises: + asyncio.CancelledError: When the cancellation route stops this operation. + """ + # Request validation is scoped to its own block so that a `TypeError` or + # `ValueError` raised *inside* the operation (a server-side fault) is not + # misreported to the client as a 4xx and, worse, swallowed without a log. + try: + thread_id = request.path_params["thread_id"] + operation_id, context, hook_responses = _operation_payload(await request.json()) + except (TypeError, ValueError) as exc: + return JSONResponse({"detail": str(exc)}, status_code=422) + + key = (thread_id, operation_id) + refusal = _register_operation(key) + if refusal is not None: + return JSONResponse({"detail": refusal}, status_code=409) + outcome: _OperationOutcome | None = "finished" + try: + try: + response = await _execute_offload( + thread_id, + operation_id=operation_id, + context=context, + hook_responses=hook_responses, + ) + except asyncio.CancelledError: + outcome = "cancelled" + raise + except _OffloadConflictError as exc: + return JSONResponse({"detail": str(exc)}, status_code=409) + except _OffloadUnavailableError as exc: + logger.exception("Offload unavailable: server runtime build failed") + return JSONResponse({"detail": str(exc)}, status_code=503) + except _OffloadIndeterminateError as exc: + return JSONResponse({"detail": str(exc)}, status_code=500) + except Exception: + logger.exception("Server-owned /offload failed") + return JSONResponse( + { + "detail": ( + "Offload failed on the server; see the server log for details." + ) + }, + status_code=500, + ) + if response["status"] == "interrupt": + outcome = None + return JSONResponse(response) + finally: + _finish_operation(key, outcome) + + +async def cancel_offload(request: Request) -> JSONResponse: + """Cancel one operation id and wait until its server task is terminal. + + Returns: + JSON containing `cancelled` when cancellation won, or `finished` when + the operation had already reached a terminal result. + """ + key = ( + request.path_params["thread_id"], + request.path_params["operation_id"], + ) + outcome = _operation_outcomes.get(key) + if outcome is not None: + return JSONResponse({"status": outcome}) + task = _active_operations.get(key) + if task is None: + _remember_operation(key, "cancelled") + return JSONResponse({"status": "cancelled"}) + task.cancel() + await asyncio.wait((task,)) + return JSONResponse({"status": _operation_outcomes.get(key, "finished")}) + + +app = Starlette( + routes=[ + Route( + "/dcode/threads/{thread_id:str}/offload", + offload, + methods=["POST"], + ), + Route( + "/dcode/threads/{thread_id:str}/offload/{operation_id:str}/cancel", + cancel_offload, + methods=["POST"], + ), + ] +) diff --git a/libs/code/deepagents_code/offload_middleware.py b/libs/code/deepagents_code/offload_middleware.py index ad5401c731..5593ad80db 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, Any, Literal, NamedTuple, Protocol, cast +from uuid import NAMESPACE_URL, uuid4, uuid5 +from weakref import WeakValueDictionary from deepagents.backends.protocol import FILE_NOT_FOUND from deepagents.middleware.summarization import ( + SummarizationState, SummarizationToolMiddleware, create_summarization_middleware, create_summarization_tool_middleware, @@ -19,11 +22,14 @@ ToolRuntime, # noqa: TC002 # inspected for runtime injection ) from langchain_core.exceptions import ContextOverflowError -from langchain_core.messages import ToolMessage -from langchain_core.tools import InjectedToolArg, StructuredTool -from langgraph.types import Command +from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.tools import StructuredTool +from langgraph.config import get_config +from langgraph.types import Command # noqa: TC002 # inspected for tool schema +from typing_extensions import TypedDict from deepagents_code._cli_context import CLIContextSchema +from deepagents_code.cost_tracking import CostState from deepagents_code.hooks.models.domain import ( CompactTrigger, HookEvent, @@ -32,6 +38,8 @@ ) from deepagents_code.hooks.server_middleware import ( _DEFAULT_DEADLINE, + _PRE_TOOL_STATE_KEY, + HookTransportInterruptError, _event_enabled, _hook_context, _invoke_hook, @@ -42,6 +50,7 @@ if TYPE_CHECKING: from collections.abc import Awaitable, Callable + from deepagents.backends.composite import CompositeBackend from deepagents.backends.protocol import ( BackendProtocol, EditResult, @@ -55,74 +64,320 @@ ModelResponse, ) from langchain.chat_models import BaseChatModel - from langgraph.prebuilt.tool_node import ToolCallRequest + from langchain_core.messages import AnyMessage + from langgraph.runtime import Runtime + + from deepagents_code.hooks.server_middleware import ServerHooksMiddleware logger = logging.getLogger(__name__) -COMPACTION_FAILURE_PREFIX = "Compaction failed" -"""Stable prefix for forced-compaction failure tool messages. +class _OffloadState(CostState, SummarizationState, total=False): + """Checkpoint channels server-owned forced compaction reads and writes. -`/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. + `_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 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 -is recognized. Because the scan is bounded to messages produced by the current -`/offload` attempt, a stale failure from an unrelated prior turn is not matched. -Only the *prefix position* is load-bearing; wording after it is free to change. -""" + `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. + """ -_OFFLOAD_SEED_ID_PREFIX = "offload-seed-" +type OffloadStatus = Literal["compacted", "empty", "noop", "denied", "failed"] +"""Outcome of one offload attempt. Aliased so the result type and the private +`_result` factory cannot drift apart.""" -class _AutoCompactionBlockedError(Exception): - """Carry a blocked provider overflow past the SDK fallback handler.""" - def __init__(self, overflow: ContextOverflowError) -> None: - super().__init__(str(overflow)) - self.overflow = overflow +class OffloadResult(TypedDict): + """Typed result emitted by the server-owned offload operation.""" + + status: OffloadStatus + messages_offloaded: int + messages_kept: int + tokens_before: int + tokens_after: int + archive_path: str | None + archive_ephemeral: bool + error: str | None + + +class OffloadStateUpdate(TypedDict, total=False): + """The only checkpoint channels a server-owned operation may write. + + Naming the permitted channels makes the load-bearing invariant -- that this + route can never write `messages` -- a property of the type rather than a + single string check performed after the summarizer has already been billed. + The runtime check in `offload_api` stays as a backstop for the `Any`-typed + values the summarization SDK hands back. + """ + + _summarization_event: dict[str, Any] + _summarization_session_id: str + _session_cost_usd: float + + +class OffloadExecution(NamedTuple): + """State update and typed result produced by one server operation.""" + + update: OffloadStateUpdate + result: OffloadResult + archive: _PendingArchive | None = None + + +_archive_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary() + + +def _archive_lock(session_id: str) -> asyncio.Lock: + """Return the process-local lock serializing one archive's read/write cycle.""" + lock = _archive_locks.get(session_id) + if lock is None: + lock = asyncio.Lock() + _archive_locks[session_id] = lock + return lock + + +class _PendingArchive(NamedTuple): + """Archive append deferred until the checkpoint summary is reserved.""" + + summarization: SummarizationMiddleware + backend: BackendProtocol + messages: list[AnyMessage] + session_id: str + summary: str + state_cutoff: int + + def update(self, file_path: str | None) -> dict[str, Any]: + """Build the summary update with the archive's settled path. + + Returns: + State update containing the summary, cutoff, and archive path. + """ + return CLICompactionMiddleware._forced_compaction_update( + self.summarization, + self.summary, + file_path, + self.state_cutoff, + self.session_id, + ) + + async def _previous_content(self, path: str) -> tuple[bool, str]: + """Read the archive snapshot needed to undo an uncommitted append. + + Returns: + Whether the archive existed and its prior UTF-8 content. + + Raises: + RuntimeError: If the backend cannot return the archive snapshot. + """ + responses = await self.backend.adownload_files([path]) + if not responses: + msg = f"archive backend returned no response for {path}" + raise RuntimeError(msg) + response = responses[0] + if response.error == FILE_NOT_FOUND: + return False, "" + if response.error is not None: + msg = f"archive read failed for {path}: {response.error}" + raise RuntimeError(msg) + content = response.content or b"" + return True, content.decode("utf-8") + + async def write(self) -> _ArchiveAppend | None: + """Append staged messages and retain enough state for rollback. + + Returns: + The reversible append, or `None` when the SDK could not write it. + """ + path = self.summarization._get_history_path(self.session_id) + existed, previous = await self._previous_content(path) + guard = cast("BackendProtocol", _ArchiveReadGuard(self.backend)) + written_path = await self.summarization._aoffload_to_backend( + guard, self.messages, self.session_id + ) + append = _ArchiveAppend(self.backend, path, existed, previous) + if written_path is None: + await append.rollback() + return None + return append + + +class _ArchiveAppend(NamedTuple): + """Completed archive append that can be restored until checkpointed.""" + + backend: BackendProtocol + path: str + existed: bool + previous: str + + async def rollback(self) -> None: + """Restore the exact archive snapshot from before the append. + + Raises: + RuntimeError: If the backend cannot restore the snapshot. + """ + result = ( + await self.backend.awrite(self.path, self.previous) + if self.existed + else await self.backend.adelete(self.path) + ) + if result.error is not None: + msg = f"archive rollback failed for {self.path}: {result.error}" + raise RuntimeError(msg) + + +class _ForcedCompactionPlan(NamedTuple): + """Checkpoint update plus its not-yet-written archive append.""" + summarization: SummarizationMiddleware + summary: str + state_cutoff: int + archive: _PendingArchive -def _offload_seed_message_id(tool_call_id: str) -> str: - """Return the stable message ID for a forced `/offload` tool call. + def update(self, file_path: str | None) -> dict[str, Any]: + """Build the summary update with the archive's settled path. + + Returns: + State update containing the summary, cutoff, and archive path. + """ + return self.archive.update(file_path) + + +def unchanged_offload_result( + status: OffloadStatus, + *, + messages: int, + tokens: int, + error: str | None = None, +) -> OffloadResult: + """Build a result for an operation that did not compact state. + + Module-level so the HTTP boundary can report an unchanged outcome without + resolving the server runtime first -- an empty thread has nothing to + compact, and building the agent only to describe that is both wasteful and + a way for a construction failure to turn "nothing to do" into a 500. Args: - tool_call_id: The seeded `compact_conversation` tool call ID. + status: A non-compacting outcome. + messages: Messages left in the conversation. + tokens: Context estimate, unchanged by definition. + error: Reason, required for `denied` and `failed`. Returns: - The synthetic assistant message ID associated with the tool call. + Typed result containing unchanged context statistics. + + Raises: + ValueError: If a refusal carries no reason. """ - return f"{_OFFLOAD_SEED_ID_PREFIX}{tool_call_id}" + if status in {"denied", "failed"} and not error: + # `error` is `str | None` on every status because the wire shape is one + # flat object, so the checker cannot make "a refusal has a reason" a + # compile-time fact. Enforce it at the single construction point + # instead: a reasonless refusal renders as the client's generic "the + # server rejected the operation", which tells the user nothing. + msg = f"An offload {status!r} result must carry a reason." + raise ValueError(msg) + return { + "status": status, + "messages_offloaded": 0, + "messages_kept": messages, + "tokens_before": tokens, + "tokens_after": tokens, + "archive_path": None, + "archive_ephemeral": False, + "error": error, + } + + +class OffloadCompleteResponse(TypedDict): + """Wire response for an attempt that finished without needing the client.""" + + status: Literal["complete"] + result: OffloadResult + + +class OffloadInterruptResponse(TypedDict): + """Wire response carrying a hook request the client must fulfill. + + `request` stays a plain mapping on purpose: the client transports it back + without inspecting it, and only `hooks.interrupt` owns its shape. + """ + + status: Literal["interrupt"] + request: dict[str, Any] + + +type OffloadResponse = OffloadCompleteResponse | OffloadInterruptResponse +"""One round of the offload operation protocol. + +Tagged on `status` so the producer (`offload_api._execute_offload`) and the +consumer (`RemoteAgent.aoffload`) are checked against one definition instead of +two independently hand-written `isinstance` ladders. The client still validates +at runtime -- this crosses HTTP, so the type is a contract, not a guarantee. +""" + +_OFFLOAD_OPERATION_ATTR = "_dcode_offload_operation" -def _without_offload_seed(messages: list[Any], tool_call_id: str) -> list[Any]: - """Exclude the synthetic `/offload` seed from retention calculations. + +def attach_offload_operation( + backend: CompositeBackend, + operation: OffloadOperation, +) -> None: + """Publish the operation on the backend shared with the server runtime. Args: - messages: Effective conversation messages including the forced tool call. - tool_call_id: The seeded `compact_conversation` tool call ID. + backend: Composite backend owned by the agent server. + operation: Offload implementation bound to that backend. + + Raises: + ValueError: If compaction writes through a different backend. + """ + # The SDK requires `backend` in its constructor, so a real summarization + # middleware always has one; `None` here means a test double, which is + # allowed through rather than asserted against. + bound = getattr(operation._compaction._summarization, "_backend", None) + if bound is not None and bound is not backend: + msg = "Offload operation must use the agent's composite backend" + raise ValueError(msg) + setattr(backend, _OFFLOAD_OPERATION_ATTR, operation) + + +def offload_operation_from(backend: CompositeBackend) -> OffloadOperation | None: + """Return the server operation published on `backend`, when available.""" + operation = getattr(backend, _OFFLOAD_OPERATION_ATTR, None) + return operation if isinstance(operation, OffloadOperation) 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: - Conversation messages without the matching synthetic assistant message. + The `cutoff_index`, or `0` when the event is missing or malformed. """ - if not tool_call_id: - return messages - seed_id = _offload_seed_message_id(tool_call_id) - return [ - message - for message in messages - if ( - message.get("id") - if isinstance(message, dict) - else getattr(message, "id", None) - ) - != seed_id - ] + if isinstance(event, dict): + cutoff = event.get("cutoff_index") + # `bool` is excluded explicitly: it passes `isinstance(_, int)`, so a + # malformed `cutoff_index: true` would otherwise read as cutoff 1 and + # silently shift the offloaded/kept counts by one message. The HTTP + # boundary rejects bools for `model_context_limit` for the same reason. + if isinstance(cutoff, int) and not isinstance(cutoff, bool): + return cutoff + return 0 + + +class _AutoCompactionBlockedError(Exception): + """Carry a blocked provider overflow past the SDK fallback handler.""" + + def __init__(self, overflow: ContextOverflowError) -> None: + super().__init__(str(overflow)) + self.overflow = overflow class RuntimeModelConfig(NamedTuple): @@ -142,11 +397,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 server operation 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 server operation + 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 +438,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 {}, @@ -178,25 +458,6 @@ def _runtime_model_config(runtime: ToolRuntime) -> RuntimeModelConfig: ) -def _offload_tool_call_id(context: object) -> str | None: - """Read the sole tool-call ID authorized for an `/offload` run. - - Args: - context: Runtime context supplied to the agent graph. - - Returns: - The authorized tool-call ID, or `None` during an ordinary agent run. - """ - value = ( - context.offload_tool_call_id - if isinstance(context, CLIContextSchema) - else context.get("offload_tool_call_id") - if isinstance(context, dict) - else None - ) - return value if isinstance(value, str) and value else None - - class _ArchiveReadGuard: """Prevent an archive write after its prerequisite read fails. @@ -343,12 +604,11 @@ async def aedit( class CLICompactionMiddleware(SummarizationToolMiddleware): - """Add hook-aware automatic and explicit forced compaction for dcode. + """Add hook-aware automatic and server-owned compaction for dcode. - The SDK tool's normal, model-initiated behavior remains unchanged. The - 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. + The SDK tool's normal, model-initiated behavior remains unchanged. + `_aplan_forced_compaction_update` is the state-only planning entry point + used by the server-owned `/offload` operation. """ @property @@ -460,6 +720,20 @@ def gated_handler(next_request: ModelRequest) -> ModelResponse: except _AutoCompactionBlockedError as blocked: raise blocked.overflow from None + async def _awrap_with_archive_lock( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], Awaitable[ModelResponse]], + ) -> ModelResponse | ExtendedModelResponse: + """Serialize an automatic archive append with other compactions. + + Returns: + The wrapped model response. + """ + session_id = self._summarization._get_session_id(request.state) + async with _archive_lock(session_id): + return await self._summarization.awrap_model_call(request, handler) + async def awrap_model_call( # ty: ignore[invalid-method-override] # delegates auto summarizer self, request: ModelRequest, @@ -475,7 +749,7 @@ async def awrap_model_call( # ty: ignore[invalid-method-override] # delegates if prepared is not None: if not self._pre_auto_compact(prepared): return await call_model(prepared) - return await self._summarization.awrap_model_call(request, call_model) + return await self._awrap_with_archive_lock(request, call_model) overflow_gated = False @@ -490,119 +764,33 @@ async def gated_handler(next_request: ModelRequest) -> ModelResponse: raise try: - return await self._summarization.awrap_model_call(request, gated_handler) + return await self._awrap_with_archive_lock(request, gated_handler) except _AutoCompactionBlockedError as blocked: raise blocked.overflow from None - @staticmethod - def _offload_rejection(request: ToolCallRequest) -> ToolMessage | None: - """Reject every tool except the exact call seeded by `/offload`. - - Args: - request: Tool call about to be executed by the graph's tool node. - - Returns: - An error result for an unauthorized `/offload` tool call, otherwise - `None` for an ordinary run or the exact seeded compaction call. - """ - expected_id = _offload_tool_call_id(request.runtime.context) - if expected_id is None: - return None - - tool_call = request.tool_call - args = tool_call.get("args") - messages = request.state.get("messages", []) - last_message = messages[-1] if messages else None - last_message_id = ( - last_message.get("id") - if isinstance(last_message, dict) - else getattr(last_message, "id", None) - ) - is_seeded_compaction = ( - tool_call.get("id") == expected_id - and tool_call.get("name") == "compact_conversation" - and isinstance(args, dict) - and args.get("force") is True - and last_message_id == _offload_seed_message_id(expected_id) - ) - if is_seeded_compaction: - return None - - return ToolMessage( - content=( - "Not executed: /offload only authorizes its seeded " - "conversation compaction call." - ), - name=tool_call.get("name"), - tool_call_id=tool_call["id"], - status="error", - ) - - def wrap_tool_call( - self, - request: ToolCallRequest, - handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]], - ) -> ToolMessage | Command[Any]: - """Apply the `/offload` per-run tool guard before synchronous tools. - - Args: - request: Tool call about to be executed. - handler: The remaining middleware/tool execution chain. - - Returns: - The guarded rejection or the downstream tool result. - """ - if (rejection := self._offload_rejection(request)) is not None: - return rejection - return handler(request) - - async def awrap_tool_call( - self, - request: ToolCallRequest, - handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]], - ) -> ToolMessage | Command[Any]: - """Apply the `/offload` per-run tool guard before asynchronous tools. - - Args: - request: Tool call about to be executed. - handler: The remaining middleware/tool execution chain. + async def _arun_compact(self, runtime: ToolRuntime[Any, Any]) -> Command: + """Serialize a model-initiated archive append with other compactions. Returns: - The guarded rejection or the downstream tool result. + The compact tool's state command. """ - if (rejection := self._offload_rejection(request)) is not None: - return rejection - return await handler(request) + session_id = self._summarization._get_session_id(runtime.state) + async with _archive_lock(session_id): + return await super()._arun_compact(runtime) def _create_compact_tool(self) -> StructuredTool: """Create the CLI variant of `compact_conversation`. Returns: - A tool that accepts the `/offload`-only `force` flag. + The model-initiated compaction tool. """ middleware = self - # `force` is annotated `InjectedToolArg` so it is stripped from the - # schema the model sees. ToolNode also strips the seeded value before - # invocation, so forced mode is selected from the trusted runtime - # context after `_offload_rejection` validates the raw tool call. - def sync_compact( - runtime: ToolRuntime[Any, Any], - force: Annotated[bool, InjectedToolArg] = False, - ) -> Command: - del force - if _offload_tool_call_id(runtime.context) != runtime.tool_call_id: - return middleware._run_compact(runtime) - return middleware._run_forced_compact(runtime) - - async def async_compact( - runtime: ToolRuntime[Any, Any], - force: Annotated[bool, InjectedToolArg] = False, - ) -> Command: - del force - if _offload_tool_call_id(runtime.context) != runtime.tool_call_id: - return await middleware._arun_compact(runtime) - return await middleware._arun_forced_compact(runtime) + def sync_compact(runtime: ToolRuntime[Any, Any]) -> Command: + return middleware._run_compact(runtime) + + async def async_compact(runtime: ToolRuntime[Any, Any]) -> Command: + return await middleware._arun_compact(runtime) return StructuredTool.from_function( name="compact_conversation", @@ -615,16 +803,8 @@ async def async_compact( coroutine=async_compact, ) - def _guarded_backend(self) -> BackendProtocol: - """Wrap the configured backend with fail-closed archive append behavior. - - Returns: - A backend adapter that refuses writes after raised archive reads. - """ - 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,137 +847,186 @@ 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 - - def _run_forced_compact(self, runtime: ToolRuntime) -> Command: - """Synchronously compact without the SDK eligibility gate. - - This deliberately mirrors the SDK's own `_run_compact` step sequence - (apply prior event, determine cutoff, partition, summarize, offload, - build result) minus the eligibility gate. Because it is a fork rather - than an override, it must be kept in parity when the SDK's compaction - flow changes; the closest-fitting SDK-side fix (a `force=` seam on - `_run_compact`) is out of scope for this PR, which is confined to - Deep Agents Code. `test_forced_compact_matches_sdk_summarizer_calls` - guards the summarizer-method call set against drift, but only by - *existence*: it catches a renamed or removed dependency, not a changed - signature nor a new step added to `_run_compact` (e.g. if the SDK later - moved inline-media offload into the gated path). Two known consequences - of that today: this fork does not call `_offload_inline_media` (only the - auto `wrap_model_call` path does), so inline base64 media in compacted - messages is not offloaded to referenceable paths and is dropped from the - XML archive -- pre-existing SDK tool-path behavior, not introduced here. + # 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. `_PendingArchive` + # applies the guard separately when writing. + return create_summarization_middleware(model, self._summarization._backend) + + async def _aplan_forced_compaction_update( + self, state: _OffloadState, runtime: _HasRunContext + ) -> _ForcedCompactionPlan | None: + """Summarize forced-compaction history without writing its archive. + + Unlike the tool paths, this raises on failure instead of returning a + `ToolMessage`: the server operation has no tool node to carry one, so it + converts the exception into a client-visible error. Any summarizer, + or planning 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 compaction state update or an error tool message. + The checkpoint/archive plan, or `None` when nothing can be compacted. + + Raises: + ValueError: If called directly with no messages. The owning server + operation handles an empty thread before reaching this helper. """ - 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) - cutoff = summarization._determine_cutoff_index(effective) - if cutoff == 0: - return self._nothing_to_compact(tool_call_id) - - session_id = summarization._get_session_id(runtime.state) - 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, session_id - ) - # 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. - return self._build_compact_result( - runtime, to_summarize, summary, file_path, event, cutoff, session_id - ) - 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) + summarization = await asyncio.to_thread( + self._summarization_for_runtime, runtime + ) + messages = state.get("messages", []) + event = state.get("_summarization_event") + if not messages: + msg = "Offload compaction requires checkpointed conversation messages." + 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) + session_id = summarization._get_session_id(state) + archive = _PendingArchive( + summarization, + self._summarization._backend, + to_summarize, + session_id, + summary, + state_cutoff, + ) + return _ForcedCompactionPlan(summarization, summary, state_cutoff, archive) + + async def arun_forced_compaction_update( + self, state: _OffloadState, runtime: _HasRunContext + ) -> dict[str, Any] | None: + """Run forced compaction and persist its archive immediately. - async def _arun_forced_compact(self, runtime: ToolRuntime) -> Command: - """Asynchronously compact without the SDK eligibility gate. + The HTTP operation uses `_aplan_forced_compaction_update` directly so + it can reserve the checkpoint before this side effect. Direct callers + retain the historical all-in-one behavior. Returns: - The compaction state update or an error tool message. + The completed state update, or `None` when nothing can be compacted. """ - tool_call_id = runtime.tool_call_id or "" + plan = await self._aplan_forced_compaction_update(state, runtime) + if plan is None: + return None try: - summarization = await asyncio.to_thread( - 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) - cutoff = summarization._determine_cutoff_index(effective) - if cutoff == 0: - return self._nothing_to_compact(tool_call_id) - - session_id = summarization._get_session_id(runtime.state) - 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, session_id - ) - # See `_run_forced_compact` for why the inherited builder is reused - # and why it stays inside the `try`. - return self._build_compact_result( - runtime, to_summarize, summary, file_path, event, cutoff, session_id + async with _archive_lock(plan.archive.session_id): + append = await plan.archive.write() + except Exception: + logger.exception("/offload archive append failed") + append = None + if append 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(plan.archive.messages), ) - 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 plan.update(append.path if append is not None else None) @staticmethod - def _forced_compact_error(tool_call_id: str, exc: Exception) -> Command: - """Build a forced-compaction failure result with a stable prefix. - - Owned by dcode so the `/offload` client can detect failures via - `COMPACTION_FAILURE_PREFIX`. The tool must return a `ToolMessage` rather - than raise, so the model (and the client) see the failure as ordinary - tool output. - - The message is intentionally generic about *where* the failure occurred: - the guarded body spans cutoff determination, summary generation, the - archive write, and result building, so it does not assert a specific - stage (and does not claim nothing was written — an archive may have been - persisted before a later step failed). It states only what is always - true on this path: the summarization event was not committed, so the - effective conversation is unchanged. + def _forced_compaction_update( + summarization: SummarizationMiddleware, + summary: str, + file_path: str | None, + state_cutoff: int, + session_id: str, + ) -> dict[str, Any]: + """Build the state-only result used by the server `/offload` operation. + + The returned dict carries the `_summarization_event` payload plus the + session id, so it is not itself a `SummarizationEvent` and cannot be + annotated as one. That the event's `summary_message` really is the + `HumanMessage` the channel expects is therefore enforced at runtime by + the `isinstance` check below, not by the type checker. Args: - tool_call_id: The originating tool call ID. - exc: The exception raised while compacting. + 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. + session_id: The id that named the history file. Persisted under + `_summarization_session_id` (mirroring the SDK's compact and + auto-summarize paths) so a later offload appends to the same + archive instead of minting a fresh file. Returns: - A `Command` whose `ToolMessage` content starts with - `COMPACTION_FAILURE_PREFIX`. + The summarization state update. + + Raises: + TypeError: If the summarizer's first message is not the + `HumanMessage` the event schema declares. """ - return Command( - update={ - "messages": [ - ToolMessage( - content=( - f"{COMPACTION_FAILURE_PREFIX}: an error occurred " - f"during compaction ({type(exc).__name__}: {exc}). " - "Your conversation is unchanged." - ), - tool_call_id=tool_call_id, - ) - ], - } - ) + 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, + }, + "_summarization_session_id": session_id, + } def _create_cli_compaction_middleware( @@ -818,3 +1047,246 @@ def _create_cli_compaction_middleware( sdk_middleware._summarization, system_prompt=sdk_middleware.system_prompt, ) + + +_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 request re-executes the operation + **from the top** rather than resuming mid-coroutine. An id minted fresh 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") -- which would break `/offload` for + exactly those users who have a `PreCompact`/`PreToolUse` hook configured, + and make the client's whole fulfill/resume loop unreachable. + + *Distinct across attempts.* 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 attempt's decision -- including + a denial -- instead of re-running the user's hook. + + `configurable.checkpoint_ns` satisfies both, and + `offload_api._execute_offload` is the invariant's owner: it derives the + namespace as `dcode_offload:{operation_id}` from the client's per-attempt + `operation_id`, which `RemoteAgent.aoffload` mints once and reuses across + every resume round of that attempt. Changing either that namespace format or + the client's reuse of `operation_id` breaks hook resume, silently, for hook + users only. + + Returns: + An id stable across this attempt's resume rounds and distinct from every + other attempt'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)}" + + +class OffloadOperation: + """Compact checkpoint state behind dcode's server-owned HTTP boundary.""" + + def __init__( + self, + compaction: CLICompactionMiddleware, + hooks: ServerHooksMiddleware, + ) -> None: + """Initialize the operation with the agent's own policy and hooks. + + Args: + compaction: Compaction middleware bound to the agent backend. + hooks: Server hook middleware used by the interactive graph. + """ + self._compaction = compaction + self._hooks = hooks + + _result = staticmethod(unchanged_offload_result) + + async def _run_hooks( + self, runtime: Runtime[CLIContextSchema] + ) -> tuple[Literal["denied", "failed"], str] | None: + """Dispatch the forced call through `PreCompact` and `PreToolUse`. + + Returns: + Failure status and detail, or `None` when hooks allow compaction. + + Raises: + HookTransportInterruptError: If the client must fulfill a hook request. + """ + try: + forced_call_id = _forced_offload_call_id() + hook_update = await self._hooks.aafter_model( + cast( + "Any", + { + "messages": [ + AIMessage( + content="", + tool_calls=[ + { + "name": "compact_conversation", + "args": {"force": True}, + "id": forced_call_id, + } + ], + ) + ] + }, + ), + cast("Runtime[Any]", runtime), + ) + except HookTransportInterruptError: + raise + except Exception as exc: + logger.exception("/offload hook dispatch failed") + return "failed", f"Offload hooks failed: {type(exc).__name__}: {exc}" + + # Fail closed on a missing channel rather than defaulting to allow. + # `_after_model` always returns this key, so its absence means the + # channel, the id derivation, or the outcome shape drifted -- and a + # `.get(..., {})` chain would read a user's *denial* as "no outcome" and + # compact straight through it, with no log. + outcomes = hook_update.get(_PRE_TOOL_STATE_KEY) + if not isinstance(outcomes, dict): + logger.error( + "Compaction hooks returned no %s channel for /offload; refusing " + "rather than treating a possible denial as an allow", + _PRE_TOOL_STATE_KEY, + ) + return "failed", ( + "Could not read the compaction hook decision; offload refused." + ) + outcome = outcomes.get(forced_call_id) or {} + if outcome.get("behavior") == "deny": + reason = outcome.get("reason") or "Blocked by a compaction hook" + return "denied", str(reason) + if outcome.get("context"): + logger.warning( + "Discarding PreToolUse additionalContext for the /offload " + "operation; no tool result or model turn exists to carry it" + ) + return None + + async def execute( + self, + state: _OffloadState, + runtime: Runtime[CLIContextSchema], + ) -> OffloadExecution: + """Run one offload against server-read checkpoint state. + + Returns: + State update for the server to persist and the typed client result. + + Raises: + HookTransportInterruptError: If the client must fulfill a hook request. + """ + from langchain_core.messages.utils import count_tokens_approximately + + messages = list(state.get("messages", [])) + event = state.get("_summarization_event") + effective = self._compaction._summarization._apply_event_to_messages( + messages, event + ) + tokens_before = count_tokens_approximately(effective) + if not messages: + result = self._result("empty", messages=0, tokens=0) + return OffloadExecution({}, result) + + hook_failure = await self._run_hooks(runtime) + if hook_failure is not None: + status, error = hook_failure + result = self._result( + status, + messages=max(0, len(messages) - _event_cutoff(event)), + tokens=tokens_before, + error=error, + ) + return OffloadExecution({}, result) + + try: + plan = await self._compaction._aplan_forced_compaction_update( + state, runtime + ) + except HookTransportInterruptError: + raise + except Exception as exc: + logger.exception("forced /offload compaction failed") + result = self._result( + "failed", + messages=max(0, len(messages) - _event_cutoff(event)), + tokens=tokens_before, + error=f"Compaction failed: {type(exc).__name__}: {exc}", + ) + return OffloadExecution({}, result) + + if plan is None: + result = self._result( + "noop", + messages=max(0, len(messages) - _event_cutoff(event)), + tokens=tokens_before, + ) + return OffloadExecution({}, result) + + update = plan.update(None) + new_event = update["_summarization_event"] + new_cutoff = _event_cutoff(new_event) + prior_cutoff = _event_cutoff(event) + effective_after = self._compaction._summarization._apply_event_to_messages( + messages, new_event + ) + file_path = new_event.get("file_path") + from deepagents_code.offload import offload_storage_is_ephemeral + + result: OffloadResult = { + "status": "compacted", + "messages_offloaded": max(0, new_cutoff - prior_cutoff), + "messages_kept": max(0, len(messages) - new_cutoff), + "tokens_before": tokens_before, + "tokens_after": count_tokens_approximately(effective_after), + "archive_path": file_path if isinstance(file_path, str) else None, + "archive_ephemeral": offload_storage_is_ephemeral(), + "error": None, + } + # Forward only the permitted channels rather than the SDK's whole update + # dict, so a future SDK change that adds keys (`messages` above all) + # cannot reach the checkpoint write through this operation. + return OffloadExecution( + { + "_summarization_event": new_event, + "_summarization_session_id": update["_summarization_session_id"], + }, + result, + plan.archive, + ) diff --git a/libs/code/deepagents_code/server_graph.py b/libs/code/deepagents_code/server_graph.py index ad42f39d97..1ab15ff502 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,10 @@ if TYPE_CHECKING: from collections.abc import Awaitable, Callable + from deepagents.backends.composite import CompositeBackend + + from deepagents_code.offload_middleware import OffloadOperation + logger = logging.getLogger(__name__) _sandbox_cm: Any = None @@ -185,15 +189,35 @@ 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 tuple so the three slots are addressed by + name: `agent` is structurally opaque to the type checker (the SDK exposes no + usable compiled-graph type here), so a positional transposition would hand + LangGraph the backend as its compiled graph with no complaint. + """ + + agent: Any + """Compiled LangGraph agent graph served as `agent`.""" + + backend: CompositeBackend + """Composite backend the agent and its operations were built with.""" + + offload: OffloadOperation + """Server-owned thread offload operation bound to `backend`.""" + + +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 agent graph, its configured composite backend, and the server-owned + offload operation bound to that backend. """ config = ServerConfig.from_env() @@ -319,7 +343,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 +356,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,65 +387,119 @@ def _create_cli_agent_sync() -> Any: # noqa: ANN401 goal_criteria_tools=read_only_context_tools, rubric_grader_tools=read_only_context_tools, ) - return agent + from deepagents_code.offload_middleware import offload_operation_from - return await asyncio.to_thread(_create_cli_agent_sync) + offload = offload_operation_from(composite_backend) + if offload is None: + msg = ( + "Agent backend did not publish its offload operation; " + "/offload has no server implementation." + ) + raise RuntimeError(msg) + return ServerRuntime( + agent=agent, + backend=composite_backend, + offload=offload, + ) + + return await asyncio.to_thread(_create_cli_graphs_sync) + + +def _build_runtime_factory( + builder: Callable[[], Awaitable[ServerRuntime]] | None = None, +) -> Callable[[], Awaitable[ServerRuntime]]: + """Build the cached factory for all server-owned runtime resources. + + The cache is load-bearing, not an optimization: MCP discovery, sandbox + creation, and `atexit` registration each must happen exactly once. Building + per request would re-discover MCP servers, leak sandbox sessions, and stack + duplicate `atexit` handlers. Two consumers now share this cache -- the + interactive graph and the offload HTTP route -- so both must resolve the + *same* agent, backend, and compaction policy for a server-side archive to be + readable by the agent. + + The cache and its lock live in this closure rather than in module-level + globals, so importing this module introduces no shared mutable state; the + single process-wide instance is created explicitly at the bottom of the + module. + + Args: + builder: Optional alternate builder used by unit tests. + + Returns: + Async runtime factory shared by the graph and custom operation API. + """ + runtime: ServerRuntime | None = None + lock = asyncio.Lock() + + async def get_runtime() -> ServerRuntime: + """Return the cached interactive graph and operation resources.""" + nonlocal runtime + if runtime is None: + async with lock: + if runtime is None: + try: + from deepagents_code.configuration.service import ( + require_healthy_managed_config, + ) + + require_healthy_managed_config(refresh=True) + runtime = await (builder or _make_graphs)() + except Exception as exc: # noqa: BLE001 # startup barrier + emit_startup_failure(exc) + sys.exit(1) + return runtime + + return get_runtime def _build_graph_factory( - builder: Callable[[], Awaitable[Any]] | None = None, + builder: Callable[[], Awaitable[ServerRuntime]] | None = None, ) -> Callable[[], Awaitable[Any]]: - """Build the cached async graph factory exposed to `langgraph dev`. + """Build a cached graph factory, for tests. - 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 - mutable state. + `langgraph.json` references the module-level `make_graph`, which delegates to + `get_server_runtime`; nothing in production calls this. It survives so unit + tests can inject a builder. Args: - builder: Optional alternate graph builder. + builder: Optional alternate runtime builder used by unit tests. Returns: - A zero-arg async factory that builds the graph once and returns the - cached instance on every subsequent call. + Async graph factory for the interactive `agent` graph. """ - missing = object() - graph: Any = missing - lock = asyncio.Lock() + get_runtime = _build_runtime_factory(builder) async def make_graph() -> Any: # noqa: ANN401 - """Create (or return the cached) agent graph for `langgraph dev`. - - 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 - 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. + """Create or return the cached agent graph for `langgraph dev`. Returns: Compiled LangGraph agent graph. """ - nonlocal graph - if graph is not missing: - return graph - async with lock: - if graph is missing: - try: - from deepagents_code.configuration.service import ( - require_healthy_managed_config, - ) - - require_healthy_managed_config(refresh=True) - 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 (await get_runtime()).agent return make_graph -make_graph = _build_graph_factory() +_get_runtime = _build_runtime_factory() + + +async def get_server_runtime() -> ServerRuntime: + """Return resources shared by the graph and dcode operation routes. + + Builds once and caches. A construction failure is converted into a + startup-error marker (scraped by the parent app process) before + `sys.exit(1)`, which is right for the `langgraph.json` graph factory at + startup. Callers in request scope must contain that exit -- `SystemExit` is a + `BaseException` -- as `offload_api._execute_offload` does, mapping it to a 503 + rather than killing the server mid-request. + + Returns: + The cached server runtime. + """ + return await _get_runtime() + + +async def make_graph() -> Any: # noqa: ANN401 + """Return the cached interactive graph for `langgraph.json`.""" + return (await get_server_runtime()).agent diff --git a/libs/code/deepagents_code/tui/textual_adapter.py b/libs/code/deepagents_code/tui/textual_adapter.py index 830b2af53d..b72f89e80c 100644 --- a/libs/code/deepagents_code/tui/textual_adapter.py +++ b/libs/code/deepagents_code/tui/textual_adapter.py @@ -3809,8 +3809,8 @@ def _report_tokens( """Refresh the token-count UI display. Persistence into graph state is owned by `ResumeStateMiddleware.after_model` - (normal turns), `_handle_offload` (offload turns), and the interrupt-cleanup - `aupdate_state` write (partial turns) — never this helper. + (normal turns), the server-side offload route (offload turns), and the + interrupt-cleanup `aupdate_state` write (partial turns) — never this helper. Args: adapter: UI adapter with token callbacks. 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..70d41da888 100644 --- a/libs/code/tests/integration_tests/test_offload_server_side.py +++ b/libs/code/tests/integration_tests/test_offload_server_side.py @@ -1,15 +1,16 @@ """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 through dcode's server HTTP operation, +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 +import re import uuid from typing import TYPE_CHECKING @@ -20,7 +21,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 +38,28 @@ 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" + ) + (config_dir / "prices.json").write_text( + """ +[ + { + "id": "itest", + "name": "Integration Test", + "api_pattern": "itest", + "models": [ + { + "id": "fake", + "match": {"equals": "fake"}, + "prices": {"input_mtok": 1.0, "output_mtok": 2.0} + } + ] + } +] """.strip() + "\n" ) @@ -44,15 +76,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 +124,17 @@ 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 never changes the thread's graph association, so the same client + # can immediately seed a read through the interactive graph. + agent_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,8 +175,10 @@ async def test_offload_runs_server_side_and_is_agent_readable( enough content, runs `/offload`, and asserts: - no `ErrorMessage` and an "Offloaded " success message, - - a persisted `_summarization_event` with `cutoff > 0` and - `file_path == /conversation_history/.md`, + - the operation succeeds through the custom server route, + - a persisted `_summarization_event` with `cutoff > 0` and a + `file_path` of `/conversation_history/session_.md` -- the SDK + names the archive from the summarization session id, not the thread id, - the archive is readable THROUGH THE AGENT (via its own `read_file` tool), proving the bytes live in the agent's backend server-side, and - local archives land in the persistent per-user history directory. @@ -189,6 +230,19 @@ async def test_offload_runs_server_side_and_is_agent_readable( config = {"configurable": {"thread_id": thread_id}} + # Captured before the operation to prove its state-only commit does + # not replace or otherwise rewrite conversation messages. + before_state = await agent.aget_state(config) + messages_before = list( + (getattr(before_state, "values", None) or {}).get("messages", []) + ) + cost_before = float( + (getattr(before_state, "values", None) or {}).get( + "_session_cost_usd", 0.0 + ) + ) + assert messages_before + # Production construction: no client-owned backend. app = DeepAgentsApp( agent=agent, # ty: ignore @@ -230,6 +284,17 @@ 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 {} + assert float(values.get("_session_cost_usd", 0.0)) > cost_before + + # `/offload` frees context by advancing the summarization cutoff, + # not by deleting messages: raw history stays checkpointed. Assert + # identity to prove the operation never supplied message input. + 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") @@ -237,9 +302,15 @@ async def test_offload_runs_server_side_and_is_agent_readable( assert cutoff > 0 # In local mode the history prefix lives under a per-session # `artifacts_root`, so assert the suffix rather than a fixed prefix. + # The leaf is the summarization *session* id (`_get_history_path`), + # which is not the thread id: asserting `{thread_id}.md` here made + # this test claim a naming scheme the SDK never produces. archive_path = _event_field(summarization_event, "file_path") assert isinstance(archive_path, str) - assert archive_path.endswith(f"/conversation_history/{thread_id}.md") + assert re.fullmatch( + r".*/conversation_history/session_[0-9a-f]{32}\.md", archive_path + ), archive_path + archive_name = archive_path.rsplit("/", 1)[1] # CRUCIAL: the archive must be readable THROUGH THE AGENT, proving # the bytes exist in the agent's own backend server-side. @@ -251,9 +322,323 @@ async def test_offload_runs_server_side_and_is_agent_readable( assert "Summarized at" in read_back persistent_archive = ( - home_dir / ".deepagents" / "conversation_history" / f"{thread_id}.md" + home_dir / ".deepagents" / "conversation_history" / archive_name ) assert persistent_archive.exists() assert "keeps enough unique detail" in persistent_archive.read_text() finally: model_config.clear_caches() + + +async def _reject_any_hook( # noqa: RUF029 # must satisfy the async fulfill_hook signature + request: object, +) -> dict[str, object]: + """Fail loudly if the offload unexpectedly routes a hook to the client. + + No hooks are configured in this test, so a well-formed operation never + interrupts. Returning a deny would mask a protocol bug as a hook denial. + + Raises: + AssertionError: Always — no hook request is expected here. + """ + msg = f"Unexpected hook request during offload: {request!r}" + raise AssertionError(msg) + + +async def _wait_for_file(path: Path) -> None: + """Poll until `path` exists, so the test can sync with the server process. + + The gate files are written by the server subprocess, whose clock and event + loop are independent of the test's; polling (with a generous ceiling) is + the only synchronization primitive available across that boundary. + + Raises: + TimeoutError: If the file does not appear within 60 seconds. + """ + import asyncio + + loop = asyncio.get_running_loop() + deadline = loop.time() + 60.0 + while not path.exists(): # noqa: ASYNC240 # cheap stat per poll; the gate protocol is file-based by design + if loop.time() > deadline: + msg = f"Timed out waiting for the server to create {path}" + raise TimeoutError(msg) + await asyncio.sleep(0.05) + + +@pytest.mark.timeout(240) +async def test_concurrent_run_during_offload_preserves_messages( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A run starting mid-offload must not be clobbered by the offload commit. + + The server operation reads state, runs compaction (a model call), then + commits a state-only update. Between its final idle check and the + `update_state` write there is a window in which a user run can start. This + test holds the compaction model call open (via the + `DCA_TEST_OFFLOAD_GATE_DIR` summary gate) and starts a run inside that + window, then asserts the invariant the design relies on: LangGraph either + rejects the offload's write or serializes it before the run — in both + cases the run's message must be present in the final thread state, and the + offload either committed cleanly or reported a conflict (it must never + silently branch from the stale checkpoint). + """ + import asyncio + + home_dir = tmp_path / "home" + project_dir = tmp_path / "project" + gate_dir = tmp_path / "gate" + assistant_id = "itest-offload-race" + + home_dir.mkdir() + project_dir.mkdir() + gate_dir.mkdir() + + monkeypatch.setenv("HOME", str(home_dir)) + monkeypatch.setenv("DEEPAGENTS_CODE_NO_UPDATE_CHECK", "1") + # Reaches the server subprocess through `_build_server_env`'s + # `os.environ.copy()`; gates only summary-generation model calls. + monkeypatch.setenv("DCA_TEST_OFFLOAD_GATE_DIR", str(gate_dir)) + monkeypatch.chdir(project_dir) + + _write_model_config(home_dir) + + from deepagents_code import model_config + from deepagents_code.client.launch.server_manager import server_session + from deepagents_code.config import create_model + from deepagents_code.sessions import generate_thread_id + + config_path = home_dir / ".deepagents" / "config.toml" + monkeypatch.setattr(model_config, "DEFAULT_CONFIG_DIR", config_path.parent) + monkeypatch.setattr(model_config, "DEFAULT_CONFIG_PATH", config_path) + + model_config.clear_caches() + try: + create_model("itest:fake").apply_to_settings() + thread_id = generate_thread_id() + + async with server_session( + assistant_id=assistant_id, + model_name="itest:fake", + no_mcp=True, + enable_shell=False, + interactive=True, + sandbox_type="none", + ) as (agent, _server_proc): + for turn in range(1, 5): + await _run_turn( + agent, + thread_id=thread_id, + assistant_id=assistant_id, + prompt=_build_long_prompt(turn), + ) + + config = {"configurable": {"thread_id": thread_id}} + messages_before = list( + (getattr(await agent.aget_state(config), "values", None) or {}).get( + "messages", [] + ) + ) + assert messages_before + + # An offload whose summary call blocks at the gate. Errors are + # captured rather than raised so the gate release and the invariant + # check run regardless of how the operation resolves. + offload_error: list[BaseException] = [] + + async def _offload() -> None: + try: + await agent.aoffload( + config=config, + context={"model": "itest:fake"}, + fulfill_hook=_reject_any_hook, + ) + except BaseException as exc: # noqa: BLE001 # asserted below + offload_error.append(exc) + + offload_task = asyncio.create_task(_offload()) + + # Wait until the server is provably mid-summary, i.e. past its idle + # checks and inside the window the final commit must be safe in. + await _wait_for_file(gate_dir / "entered") + + # Launch a real run on the same thread while offload is blocked. + # Its model call is not a summary request, so it passes the gate. + run_task = asyncio.create_task( + _run_turn( + agent, + thread_id=thread_id, + assistant_id=assistant_id, + prompt="concurrent turn: the message that must survive", + ) + ) + # Give the run a beat to register server-side before releasing the + # offload, so the commit and the run genuinely overlap. + await asyncio.sleep(1.0) + (gate_dir / "release").write_text("1") + + await asyncio.wait_for(run_task, timeout=120) + await asyncio.wait_for(offload_task, timeout=120) + + # The offload either committed or failed with a conflict; both are + # acceptable outcomes of a genuine race. A hang or an unexpected + # exception type is not. + for exc in offload_error: + text = f"{type(exc).__name__}: {exc}" + assert "changed" in text or "active" in text or "409" in text, text + + # The invariant: whatever the offload did, the concurrent run's + # message survived. If LangGraph ever lets the state-only write + # branch from the stale checkpoint, this fails because the run's + # appended messages would be missing. + final_values = getattr(await agent.aget_state(config), "values", None) or {} + final_contents = [ + str(getattr(m, "content", m.get("content", ""))) + for m in final_values.get("messages", []) + ] + assert any( + "the message that must survive" in content for content in final_contents + ), final_contents + # The pre-offload history was not truncated either: the event only + # advances a cutoff; raw messages stay checkpointed. + assert len(final_values.get("messages", [])) >= len(messages_before) + finally: + # Never leave the server subprocess blocked on the gate. + (gate_dir / "release").write_text("1") + model_config.clear_caches() + + +_TEST_AUTH_MODULE = '''\ +"""Minimal token auth backend for the custom-route-auth integration test.""" + +from langgraph_sdk import Auth + +auth = Auth() + + +@auth.authenticate +async def authenticate(authorization: str | None) -> str: + """Accept only the fixed test bearer token; reject everything else. + + Returns: + A user id for the one credential this test server trusts. + + Raises: + Auth.exceptions.HTTPException: On a missing or wrong token. + """ + if authorization != "Bearer itest-token": + raise Auth.exceptions.HTTPException(status_code=401, detail="nope") + return "itest-user" +''' + + +@pytest.mark.timeout(240) +async def test_offload_route_respects_configured_auth( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The custom offload routes are gated exactly like the graph routes. + + Production dcode servers run `LANGGRAPH_AUTH_TYPE=noop` (localhost trust), + but the generated `langgraph.json` sets `enable_custom_route_auth: True` + so that a deployment which *does* configure an auth backend gets the + `/dcode/*` operation routes behind the same middleware as `/threads`. + The threat model asserts that; this test proves it end to end: a real + server with a token-rejecting auth backend must reject an unauthenticated + POST to the offload route with the same 401 it gives a protected graph + route — not with a 404/422 that would mean the route bypassed auth — and + must accept the request once the credential is supplied. + """ + import httpx + + home_dir = tmp_path / "home" + project_dir = tmp_path / "project" + work_dir = tmp_path / "server_work" + home_dir.mkdir() + project_dir.mkdir() + work_dir.mkdir() + + monkeypatch.setenv("HOME", str(home_dir)) + monkeypatch.setenv("DEEPAGENTS_CODE_NO_UPDATE_CHECK", "1") + monkeypatch.chdir(project_dir) + + _write_model_config(home_dir) + + from deepagents_code import model_config + from deepagents_code.client.launch.server import ( + ServerProcess, + generate_langgraph_json, + ) + from deepagents_code.config import create_model + + config_path = home_dir / ".deepagents" / "config.toml" + monkeypatch.setattr(model_config, "DEFAULT_CONFIG_DIR", config_path.parent) + monkeypatch.setattr(model_config, "DEFAULT_CONFIG_PATH", config_path) + + model_config.clear_caches() + server: ServerProcess | None = None + try: + create_model("itest:fake").apply_to_settings() + + # The auth module lives in the server work dir (the subprocess's cwd, + # which `langgraph dev` puts on `sys.path`) so its import path stays + # relative to the deployment, exactly like a real deployment's + # `auth.py` next to its `langgraph.json`. + (work_dir / "itest_auth.py").write_text(_TEST_AUTH_MODULE) + generate_langgraph_json( + work_dir, + auth_path="./itest_auth.py:auth", + ) + + # No scaffold: the workspace is fully prepared above, and a missing + # langgraph.json here would be a test bug worth failing on. + server = ServerProcess(config_dir=work_dir, scaffold=None) + await server.start() + + async with httpx.AsyncClient(base_url=server.url) as http: + unauthenticated = await http.post( + "/dcode/threads/thread-1/offload", + json={"operation_id": "op-1", "context": {}, "hook_responses": {}}, + ) + protected_graph_route = await http.post("/threads", json={}) + assert unauthenticated.status_code == 401, ( + unauthenticated.status_code, + unauthenticated.text, + ) + assert protected_graph_route.status_code == 401, ( + protected_graph_route.status_code, + protected_graph_route.text, + ) + + headers = {"Authorization": "Bearer itest-token"} + authenticated = await http.post( + "/dcode/threads/thread-1/offload", + json={"operation_id": "op-1", "context": {}, "hook_responses": {}}, + headers=headers, + ) + # 404/409/500 all pass auth and fail inside the operation (the + # thread does not exist); only 401/403 would mean auth still + # rejected a credentialed request. + assert authenticated.status_code not in (401, 403), ( + authenticated.status_code, + authenticated.text, + ) + # A malformed context fails at the boundary with a field-naming + # 422, not a 500 from deep in model resolution. + malformed = await http.post( + "/dcode/threads/thread-1/offload", + json={ + "operation_id": "op-1", + "context": {"model": 123}, + "hook_responses": {}, + }, + headers=headers, + ) + assert malformed.status_code == 422, ( + malformed.status_code, + malformed.text, + ) + assert "context.model" in malformed.text + finally: + if server is not None: + server.stop() + model_config.clear_caches() diff --git a/libs/code/tests/unit_tests/hooks/test_server_lifecycle.py b/libs/code/tests/unit_tests/hooks/test_server_lifecycle.py index 8348398298..e32dbd6d98 100644 --- a/libs/code/tests/unit_tests/hooks/test_server_lifecycle.py +++ b/libs/code/tests/unit_tests/hooks/test_server_lifecycle.py @@ -8,7 +8,7 @@ from datetime import UTC, datetime, timedelta from pathlib import Path from typing import TYPE_CHECKING, Any, NotRequired -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from uuid import UUID, uuid4 import pytest @@ -67,6 +67,7 @@ from deepagents_code.hooks.presenter import HookPresenter from deepagents_code.hooks.runtime import HooksRuntime from deepagents_code.hooks.server_middleware import ( + HookTransportInterruptError, ServerHooksMiddleware, ServerHooksState, _append_message_text, @@ -81,6 +82,7 @@ _session_gate, _tool_result_error, _tool_result_text, + operation_hook_responses, ) from deepagents_code.hooks.snapshot import HooksSnapshot from deepagents_code.hooks.transcript import SUBAGENT_TRANSCRIPT_ID_METADATA_KEY @@ -569,6 +571,56 @@ def test_hook_resume_value_validates_identity() -> None: ) +def test_operation_hook_transport_requests_then_consumes_response() -> None: + """HTTP operations replay a deterministic hook from supplied responses.""" + request = _request() + event = request.invocation.event + assert isinstance(event, PreToolUseEvent) + gate = _session_gate( + { + "hooks_snapshot_id": request.snapshot_id, + "hooks_server_events": [HookEvent.PRE_TOOL_USE.value], + } + ) + assert gate is not None + + with ( + operation_hook_responses({}), + pytest.raises(HookTransportInterruptError) as exc_info, + ): + _invoke_hook( + request.invocation.context, + event, + gate=gate, + config={"configurable": {"thread_id": "thread-1"}}, + deadline=timedelta(seconds=1), + ) + + pending = exc_info.value.request + resume = build_hook_resume_value( + HookInvocationResponse( + protocol_version=1, + invocation_id=pending.invocation_id, + snapshot_id=pending.snapshot_id, + decision=PreToolUseDecision( + event=HookEvent.PRE_TOOL_USE, + permission=PermissionEffect(behavior="allow"), + ), + ) + ) + with operation_hook_responses({str(pending.invocation_id): resume}): + decision = _invoke_hook( + request.invocation.context, + event, + gate=gate, + config={"configurable": {"thread_id": "thread-1"}}, + deadline=timedelta(seconds=1), + ) + + assert isinstance(decision, PreToolUseDecision) + assert decision.permission.behavior == "allow" + + def _invoke_pre_tool_hook( monkeypatch: pytest.MonkeyPatch, request: HookInvocationRequest, @@ -1458,3 +1510,51 @@ def test_snapshot_configured_server_events() -> None: HookEvent.PRE_COMPACT, HookEvent.PRE_TOOL_USE, } + + +class TestAskDecisionInServerOperation: + """`ask` cannot prompt on the server-operation path, so it fails closed.""" + + @staticmethod + def _ask_call() -> tuple[ToolCallData, PermissionEffect]: + """Build a compaction call and an `ask` permission for it.""" + call = ToolCallData( + id="call-1", name="compact_conversation", args={"force": True} + ) + return call, PermissionEffect(behavior="ask", reason="please confirm") + + def test_ask_denies_instead_of_raising_a_scratchpad_keyerror(self) -> None: + """In operation mode there is no Pregel task for `interrupt()` to use. + + Without this branch `interrupt()` raises `KeyError` on LangGraph's + internal scratchpad config key, which the compaction chain's broad + handler turns into "Offload hooks failed: KeyError: ...". + """ + from deepagents_code.hooks.server_middleware import ( + _ask_permission_via_hitl, + operation_hook_responses, + ) + + call, permission = self._ask_call() + with operation_hook_responses({}): + blocked = _ask_permission_via_hitl(call, permission) + + assert blocked is not None + assert blocked.status == "error" + assert "cannot prompt for approval" in str(blocked.content) + assert "compact_conversation" in str(blocked.content) + + def test_graph_mode_still_escalates_through_hitl(self) -> None: + """Outside an operation the `ask` path must still reach `interrupt()`.""" + from deepagents_code.hooks import server_middleware + + call, permission = self._ask_call() + with patch.object( + server_middleware, + "interrupt", + return_value={"decisions": [{"type": "approve"}]}, + ) as interrupt_mock: + blocked = server_middleware._ask_permission_via_hitl(call, permission) + + interrupt_mock.assert_called_once() + assert blocked is None diff --git a/libs/code/tests/unit_tests/test_agent.py b/libs/code/tests/unit_tests/test_agent.py index 8fb97e5903..40b4eba46a 100644 --- a/libs/code/tests/unit_tests/test_agent.py +++ b/libs/code/tests/unit_tests/test_agent.py @@ -153,6 +153,25 @@ def test_add_interrupt_on_attaches_auto_approve_predicate() -> None: assert config.get("when") is _should_interrupt_tool_call +def test_agent_publishes_server_offload_operation(tmp_path: Path) -> None: + """The backend exposes offload without adding graph input fields.""" + 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, + ) + + from deepagents_code.offload_middleware import offload_operation_from + + schema = agent.get_input_jsonschema() + assert "dcode_operation" not in schema["properties"] + assert offload_operation_from(backend) is not None + + 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_compact_tool.py b/libs/code/tests/unit_tests/test_compact_tool.py index f59c2aa238..a023950d9c 100644 --- a/libs/code/tests/unit_tests/test_compact_tool.py +++ b/libs/code/tests/unit_tests/test_compact_tool.py @@ -6,21 +6,19 @@ from __future__ import annotations -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 from deepagents.backends.protocol import FileDownloadResponse, WriteResult from langchain.agents.middleware.types import ModelRequest from langchain_core.exceptions import ContextOverflowError -from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.messages import HumanMessage from langgraph.runtime import Runtime from deepagents_code._cli_context import CLIContextSchema from deepagents_code.offload_middleware import ( - COMPACTION_FAILURE_PREFIX, CLICompactionMiddleware, _ArchiveReadGuard, _runtime_model_config, @@ -120,7 +118,17 @@ class TestCLICompactionMiddleware: @staticmethod def _summarization() -> MagicMock: summarization = MagicMock() - summarization._backend = object() + backend = MagicMock() + backend.adownload_files = AsyncMock( + return_value=[ + FileDownloadResponse( + path="/conversation_history/thread.md", + error="file_not_found", + ) + ] + ) + summarization._backend = backend + summarization._get_history_path.return_value = "/conversation_history/thread.md" summarization._apply_event_to_messages.side_effect = lambda messages, _event: ( messages ) @@ -228,21 +236,209 @@ async def run_middleware() -> None: summarization._offload_to_backend.assert_not_called() summarization._create_summary.assert_not_called() - async def test_force_bypasses_sdk_eligibility_gate(self) -> None: - """Forced compaction partitions directly even below the proactive gate.""" + async def test_operation_path_writes_through_the_archive_guard(self) -> None: + """The server `/offload` operation's write path has the same invariant. + + The guard is applied per write site rather than by the backend's type, so + the server operation 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_plan_defers_archive_until_checkpoint_reservation( + self, + ) -> None: + """Planning may spend on a summary but cannot mutate archive storage.""" + summarization = self._summarization() + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + + plan = await middleware._aplan_forced_compaction_update( + {"messages": [HumanMessage("one"), HumanMessage("two")]}, runtime + ) + + assert plan is not None + assert plan.update(None)["_summarization_event"]["file_path"] is None + summarization._aoffload_to_backend.assert_not_awaited() + + 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_threads_and_persists_the_session_id(self) -> None: + """`/offload` must reuse and re-commit the SDK's archive-file id. + + The SDK's `_offload_to_backend` names the archive by `session_id`, and + the committed `_summarization_session_id` is what makes a later + compaction append to the same file instead of starting a new one. The + server operation bypasses the SDK's own state update, so it has to + thread the id through and write it back itself. + """ summarization = self._summarization() + summarization._get_session_id.return_value = "session_abc" middleware = CLICompactionMiddleware(summarization) runtime = MagicMock() runtime.context = None - runtime.state = {"messages": [HumanMessage("one"), HumanMessage("two")]} - runtime.tool_call_id = "tool-call" - result = await middleware._arun_forced_compact(runtime) + result = await middleware.arun_forced_compaction_update( + {"messages": [HumanMessage("one"), HumanMessage("two")]}, runtime + ) - summarization._is_eligible_for_compaction.assert_not_called() - summarization._acreate_summary.assert_awaited_once() - assert result.update is not None - assert result.update["_summarization_event"]["cutoff_index"] == 2 + assert result is not None + assert result["_summarization_session_id"] == "session_abc" + assert summarization._aoffload_to_backend.await_args.args[2] == "session_abc" + + 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. + + The server operation normally handles an empty thread before invoking + compaction. A direct caller that bypasses that service check still gets + an explicit error instead of a misleading successful no-op. + """ + summarization = self._summarization() + middleware = CLICompactionMiddleware(summarization) + runtime = MagicMock() + runtime.context = None + + with pytest.raises(ValueError, match="checkpointed 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`.""" @@ -274,8 +470,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 server + # operation applies the guard separately at the write site. + 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.""" @@ -310,246 +511,38 @@ 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).""" - summarization = self._summarization() - summarization._determine_cutoff_index.return_value = 0 - middleware = CLICompactionMiddleware(summarization) - runtime = MagicMock() - runtime.context = None - runtime.state = {"messages": [HumanMessage("one")]} - runtime.tool_call_id = "tool-call" - - result = await middleware._arun_forced_compact(runtime) - - assert result.update is not None - assert "_summarization_event" not in result.update - summarization._acreate_summary.assert_not_awaited() - assert "Nothing to compact" in result.update["messages"][0].content - - async def test_async_force_excludes_seed_from_retention_cutoff(self) -> None: - """The async cutoff is calculated from the pre-seed conversation.""" - summarization = self._summarization() - summarization._determine_cutoff_index.side_effect = lambda messages: ( - 0 if len(messages) == 6 else 1 - ) - middleware = CLICompactionMiddleware(summarization) - conversation = [HumanMessage(str(index)) for index in range(6)] - seed = AIMessage( - content="", - id="offload-seed-tool-call", - tool_calls=[ - { - "name": "compact_conversation", - "args": {"force": True}, - "id": "tool-call", - } - ], - ) - runtime = MagicMock() - runtime.context = None - runtime.state = {"messages": [*conversation, seed]} - runtime.tool_call_id = "tool-call" - - result = await middleware._arun_forced_compact(runtime) - - assert result.update is not None - assert "_summarization_event" not in result.update - summarization._determine_cutoff_index.assert_called_once_with(conversation) - summarization._partition_messages.assert_not_called() - - def test_sync_force_excludes_serialized_seed_from_retention_cutoff(self) -> None: - """The sync cutoff also ignores a serialized synthetic seed.""" - summarization = self._summarization() - summarization._determine_cutoff_index.side_effect = lambda messages: ( - 0 if len(messages) == 6 else 1 - ) - middleware = CLICompactionMiddleware(summarization) - conversation = [HumanMessage(str(index)) for index in range(6)] - seed = {"id": "offload-seed-tool-call", "type": "ai", "content": ""} - runtime = MagicMock() - runtime.context = None - runtime.state = {"messages": [*conversation, seed]} - runtime.tool_call_id = "tool-call" - - result = middleware._run_forced_compact(runtime) - - assert result.update is not None - assert "_summarization_event" not in result.update - summarization._determine_cutoff_index.assert_called_once_with(conversation) - summarization._partition_messages.assert_not_called() - - async def test_forced_compact_error_when_summary_fails(self) -> None: - """A summary failure returns the failure prefix and does not compact.""" - summarization = self._summarization() - summarization._acreate_summary = AsyncMock(side_effect=RuntimeError("boom")) - middleware = CLICompactionMiddleware(summarization) - runtime = MagicMock() - runtime.context = None - runtime.state = {"messages": [HumanMessage("one"), HumanMessage("two")]} - runtime.tool_call_id = "tool-call" - - result = await middleware._arun_forced_compact(runtime) - - # The failure must NOT persist an event, and must carry the stable - # prefix the `/offload` client keys on. - assert result.update is not None - assert "_summarization_event" not in result.update - content = result.update["messages"][0].content - assert content.startswith(COMPACTION_FAILURE_PREFIX) - assert "RuntimeError" in content - - def test_sync_forced_compact_compacts(self) -> None: - """The synchronous forced path mirrors the async one.""" - summarization = self._summarization() - summarization._create_summary.return_value = "Summary" - summarization._offload_to_backend.return_value = ( - "/conversation_history/thread.md" - ) - middleware = CLICompactionMiddleware(summarization) - runtime = MagicMock() - runtime.context = None - runtime.state = {"messages": [HumanMessage("one"), HumanMessage("two")]} - runtime.tool_call_id = "tool-call" - - result = middleware._run_forced_compact(runtime) - - summarization._create_summary.assert_called_once() - assert result.update is not None - assert result.update["_summarization_event"]["cutoff_index"] == 2 - - def test_force_is_hidden_from_model_schema(self) -> None: - """`force` must not appear in the schema the model sees.""" - middleware = CLICompactionMiddleware(self._summarization()) - tool = middleware.tools[0] - # `tool_call_schema` is a pydantic model (or, rarely, a dict); either - # way the model-facing property set must not expose `force`. - schema: Any = tool.tool_call_schema - props = ( - schema.get("properties", {}) - if isinstance(schema, dict) - else schema.model_json_schema().get("properties", {}) - ) - assert "force" not in props - - def test_ordinary_context_delegates_to_gated_path(self) -> None: - """Caller-supplied `force` cannot bypass the trusted runtime context.""" + def test_model_initiated_tool_delegates_to_gated_path(self) -> None: + """The public tool keeps using the SDK's eligibility-gated sync path.""" middleware = CLICompactionMiddleware(self._summarization()) tool: Any = middleware.tools[0] runtime = MagicMock() - runtime.context = {} - runtime.tool_call_id = "model-call" - with ( - patch.object(middleware, "_run_compact", return_value="gated") as gated, - patch.object( - middleware, "_run_forced_compact", return_value="forced" - ) as forced, - ): - assert tool.func(runtime, force=False) == "gated" - assert tool.func(runtime, force=True) == "gated" - assert gated.call_count == 2 - forced.assert_not_called() - async def test_offload_context_delegates_to_forced_path_async(self) -> None: - """The authorized call ID in runtime context selects forced mode.""" + with patch.object(middleware, "_run_compact", return_value="gated") as gated: + assert tool.func(runtime) == "gated" + + gated.assert_called_once_with(runtime) + + async def test_model_initiated_tool_delegates_to_gated_path_async(self) -> None: + """The public tool keeps using the SDK's eligibility-gated async path.""" middleware = CLICompactionMiddleware(self._summarization()) tool: Any = middleware.tools[0] runtime = MagicMock() - runtime.context = {"offload_tool_call_id": "offload-call"} - runtime.tool_call_id = "offload-call" - with ( - patch.object( - middleware, - "_arun_compact", - new_callable=AsyncMock, - return_value="gated", - ) as gated, - patch.object( - middleware, - "_arun_forced_compact", - new_callable=AsyncMock, - return_value="forced", - ) as forced, - ): - # ToolNode replaces the seeded `force=True` with this default. - assert await tool.coroutine(runtime, force=False) == "forced" - gated.assert_not_awaited() - forced.assert_awaited_once_with(runtime) - - async def test_tool_node_preserves_forced_mode_via_runtime_context(self) -> None: - """A real ToolNode strips `force` but still reaches forced compaction.""" - from langchain_core.messages import ToolMessage - from langgraph.graph import END, START, StateGraph - from langgraph.prebuilt import ToolNode - from langgraph.types import Command - from typing_extensions import TypedDict - - class ToolState(TypedDict): - messages: list[object] - middleware = CLICompactionMiddleware(self._summarization()) - # LangGraph accepts these runtime schemas, but its generic bound is not - # recognized by ty on Python 3.14. - builder = StateGraph( - ToolState, # ty: ignore[invalid-argument-type] - context_schema=CLIContextSchema, - ) - builder.add_node("tools", ToolNode(middleware.tools)) - builder.add_edge(START, "tools") - builder.add_edge("tools", END) - graph = builder.compile() - tool_call_id = "offload-call" - seed = AIMessage( - content="", - id=f"offload-seed-{tool_call_id}", - tool_calls=[ - { - "name": "compact_conversation", - "args": {"force": True}, - "id": tool_call_id, - } - ], - ) - command = Command( - update={ - "messages": [ - ToolMessage(content="compacted", tool_call_id=tool_call_id) - ] - } - ) - - with ( - patch.object( - middleware, - "_arun_compact", - new_callable=AsyncMock, - return_value=command, - ) as gated, - patch.object( - middleware, - "_arun_forced_compact", - new_callable=AsyncMock, - return_value=command, - ) as forced, - warnings.catch_warnings(), - ): - warnings.filterwarnings( - "error", message="Pydantic serializer warnings", category=UserWarning - ) - await graph.ainvoke( - ToolState(messages=[seed]), # ty: ignore[invalid-argument-type] - context=CLIContextSchema( # ty: ignore[invalid-argument-type] - offload_tool_call_id=tool_call_id - ), - ) + with patch.object( + middleware, + "_arun_compact", + new=AsyncMock(return_value="gated"), + ) as gated: + assert await tool.coroutine(runtime) == "gated" - gated.assert_not_awaited() - forced.assert_awaited_once() + gated.assert_awaited_once_with(runtime) - async def test_read_failure_never_reaches_truncating_archive_write(self) -> None: - """A transient archive read failure aborts the SDK write fallback.""" + async def test_operation_read_failure_never_truncates_archive(self) -> None: + """A transient archive read failure blocks the server operation's write.""" from deepagents.middleware.summarization import SummarizationMiddleware summarization = self._summarization() @@ -571,62 +564,16 @@ async def sdk_offload( summarization._aoffload_to_backend = AsyncMock(side_effect=sdk_offload) middleware = CLICompactionMiddleware(summarization) runtime = MagicMock() - runtime.context = {"offload_tool_call_id": "tool-call"} - runtime.state = {"messages": [HumanMessage("one"), HumanMessage("two")]} - runtime.tool_call_id = "tool-call" + runtime.context = None - result = await middleware._arun_forced_compact(runtime) + result = await middleware.arun_forced_compaction_update( + {"messages": [HumanMessage("one"), HumanMessage("two")]}, runtime + ) backend.awrite.assert_not_awaited() backend.aedit.assert_not_awaited() - assert result.update is not None - assert result.update["_summarization_event"]["file_path"] is None - - def test_sync_forced_compact_noops_when_nothing_old_enough(self) -> None: - """The sync forced path also no-ops at cutoff 0 (mirrors the async one).""" - summarization = self._summarization() - summarization._determine_cutoff_index.return_value = 0 - middleware = CLICompactionMiddleware(summarization) - runtime = MagicMock() - runtime.context = None - runtime.state = {"messages": [HumanMessage("one")]} - runtime.tool_call_id = "tool-call" - - result = middleware._run_forced_compact(runtime) - - assert result.update is not None - assert "_summarization_event" not in result.update - summarization._create_summary.assert_not_called() - assert "Nothing to compact" in result.update["messages"][0].content - - def test_sync_forced_compact_error_when_summary_fails(self) -> None: - """A sync summary failure returns the failure prefix and does not compact.""" - summarization = self._summarization() - summarization._create_summary = MagicMock(side_effect=RuntimeError("boom")) - middleware = CLICompactionMiddleware(summarization) - runtime = MagicMock() - runtime.context = None - runtime.state = {"messages": [HumanMessage("one"), HumanMessage("two")]} - runtime.tool_call_id = "tool-call" - - result = middleware._run_forced_compact(runtime) - - assert result.update is not None - assert "_summarization_event" not in result.update - content = result.update["messages"][0].content - assert content.startswith(COMPACTION_FAILURE_PREFIX) - assert "RuntimeError" in content - - def test_forced_compact_error_starts_with_prefix(self) -> None: - """The prefix position is the load-bearing failure-detection contract.""" - command = CLICompactionMiddleware._forced_compact_error( - "call-1", RuntimeError("boom") - ) - assert command.update is not None - (message,) = command.update["messages"] - assert message.content.startswith(COMPACTION_FAILURE_PREFIX) - assert message.tool_call_id == "call-1" - assert "RuntimeError" in message.content + assert result is not None + assert result["_summarization_event"]["file_path"] is None def test_factory_builds_cli_middleware_threading_system_prompt(self) -> None: """The factory returns a CLI middleware carrying the SDK's config.""" @@ -703,40 +650,7 @@ def test_named_fields_disambiguate_the_two_dict_slots(self) -> None: class TestSdkContractGuards: - """Guard the SDK seams the forced-compaction fork depends on. - - `CLICompactionMiddleware` forks the SDK's gated compaction flow and keys - failure detection on a shared message prefix. These tests fail loudly in CI - if a coordinated SDK bump renames a depended-on private method or changes - the failure wording, instead of the fork silently drifting out of parity. - """ - - def test_forced_compact_matches_sdk_summarizer_calls(self) -> None: - """Every SDK method the fork invokes must still exist.""" - from deepagents.middleware.summarization import ( - SummarizationMiddleware, - SummarizationToolMiddleware, - ) - - # Called on `self._summarization` (a SummarizationMiddleware). - for name in ( - "_apply_event_to_messages", - "_determine_cutoff_index", - "_partition_messages", - "_create_summary", - "_acreate_summary", - "_get_session_id", - "_offload_to_backend", - "_aoffload_to_backend", - ): - assert callable(getattr(SummarizationMiddleware, name, None)), name - - # Inherited SDK helpers called on the tool-middleware subclass. - for name in ( - "_build_compact_result", - "_nothing_to_compact", - ): - assert callable(getattr(SummarizationToolMiddleware, name, None)), name + """Guard summarization-event assumptions shared with the SDK.""" def test_summarization_cutoff_is_an_absolute_index(self) -> None: """`cutoff_index` must index unfiltered persisted messages. @@ -760,18 +674,3 @@ def test_summarization_cutoff_is_an_absolute_index(self) -> None: ) assert applied == ["S", "m2", "m3"] - - def test_failure_prefix_matches_sdk_failure_message(self) -> None: - """Dcode's prefix must match the SDK's own compaction-failure wording. - - `/offload` detects failures from either path by this prefix, so the - SDK's `_compact_error` message must keep starting with it. - """ - from deepagents.middleware.summarization import SummarizationToolMiddleware - - command = SummarizationToolMiddleware._compact_error( - "call-1", RuntimeError("boom") - ) - assert command.update is not None - (message,) = command.update["messages"] - assert message.content.startswith(COMPACTION_FAILURE_PREFIX) diff --git a/libs/code/tests/unit_tests/test_configurable_model.py b/libs/code/tests/unit_tests/test_configurable_model.py index 56db57400c..fdc926ac00 100644 --- a/libs/code/tests/unit_tests/test_configurable_model.py +++ b/libs/code/tests/unit_tests/test_configurable_model.py @@ -27,6 +27,7 @@ _is_anthropic_model, _is_fireworks_model, _is_openai_model, + _model_spec_from_model, _ResolvedModelRequest, ) @@ -113,6 +114,20 @@ def _make_model_result( class TestCheckpointPersistence: """Tests for private resume-state checkpoint updates.""" + def test_startup_custom_provider_uses_configured_spec(self) -> None: + """Custom classes must checkpoint their configured provider alias.""" + from deepagents_code.config import settings + + model = _make_model("fake") + model._get_ls_params.return_value = { + "ls_provider": "deterministicintegrationchatmodel" + } + with ( + patch.object(settings, "model_provider", "itest"), + patch.object(settings, "model_name", "fake"), + ): + assert _model_spec_from_model(model) == "itest:fake" + def test_records_request_start_only_after_success(self) -> None: middleware = ConfigurableModelMiddleware(openai_prompt_cache_key=True) request = _make_request(_make_model("gpt-5.6")) diff --git a/libs/code/tests/unit_tests/test_cost_tracking.py b/libs/code/tests/unit_tests/test_cost_tracking.py index e2b0424189..8ecfacc71f 100644 --- a/libs/code/tests/unit_tests/test_cost_tracking.py +++ b/libs/code/tests/unit_tests/test_cost_tracking.py @@ -4,6 +4,7 @@ import asyncio import builtins +import gc import inspect import json import logging @@ -2259,6 +2260,89 @@ def test_a_refusal_leaves_the_installed_catalog_alone( class TestCostTrackingMiddleware: """Tests for cumulative cost writes on the model checkpoint path.""" + def test_prepared_operation_cost_can_commit_or_rollback( + self, + recorder: _SessionCostRecorder, + ) -> None: + """Operation pricing is additive and restores records after failure.""" + _collect( + recorder, + _record(message_id="offload-summary"), + checkpoint_ns="dcode_offload:operation-1", + ) + state = cast( + "CostState", + { + "messages": [], + "_model_spec": f"{KNOWN_PROVIDER}:{KNOWN_MODEL}", + }, + ) + + prepared = cost_tracking.prepare_operation_cost(state, THREAD_ID) + one_call = estimate_cost(_usage(), KNOWN_MODEL, KNOWN_PROVIDER) + + assert one_call is not None + assert prepared.update == {"_session_cost_usd": pytest.approx(one_call)} + assert recorder.drain(THREAD_ID) == [] + + prepared.rollback() + retried = cost_tracking.prepare_operation_cost(state, THREAD_ID) + assert retried.update == {"_session_cost_usd": pytest.approx(one_call)} + + def test_committed_prepare_does_not_restore_records( + self, + recorder: _SessionCostRecorder, + ) -> None: + """A committed prepare keeps its records drained. + + Restoring them would let the next drain price the same spend a second + time, so `commit` must settle the instance without touching the + recorder. + """ + _collect( + recorder, + _record(message_id="offload-summary"), + checkpoint_ns="dcode_offload:operation-1", + ) + state = cast( + "CostState", + {"messages": [], "_model_spec": f"{KNOWN_PROVIDER}:{KNOWN_MODEL}"}, + ) + + prepared = cost_tracking.prepare_operation_cost(state, THREAD_ID) + prepared.commit() + prepared.rollback() + + assert cost_tracking.prepare_operation_cost(state, THREAD_ID).update == {} + + def test_abandoned_prepare_warns_that_spend_was_lost( + self, + recorder: _SessionCostRecorder, + caplog: pytest.LogCaptureFixture, + ) -> None: + """An unsettled prepare must not disappear silently. + + The drain is destructive, so a prepare that is neither committed nor + rolled back deletes its spend from the thread's lifetime total. That was + the one settlement outcome with no observable trace. + """ + _collect( + recorder, + _record(message_id="offload-summary"), + checkpoint_ns="dcode_offload:operation-1", + ) + state = cast( + "CostState", + {"messages": [], "_model_spec": f"{KNOWN_PROVIDER}:{KNOWN_MODEL}"}, + ) + + prepared = cost_tracking.prepare_operation_cost(state, THREAD_ID) + with caplog.at_level(logging.WARNING): + del prepared + gc.collect() + + assert "abandoned without commit or rollback" in caplog.text + def test_cost_channel_is_private_and_additive(self) -> None: """The channel must compile to a summing reducer, not a `LastValue`. diff --git a/libs/code/tests/unit_tests/test_offload.py b/libs/code/tests/unit_tests/test_offload.py index 3732758b55..b76d1aa19a 100644 --- a/libs/code/tests/unit_tests/test_offload.py +++ b/libs/code/tests/unit_tests/test_offload.py @@ -10,20 +10,22 @@ import time from contextlib import nullcontext from pathlib import Path, PureWindowsPath -from typing import TYPE_CHECKING, Any +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, cast from unittest.mock import AsyncMock, MagicMock, patch if TYPE_CHECKING: - from collections.abc import Coroutine + from collections.abc import Callable, Coroutine import pytest from deepagents.backends.utils import validate_path +from langgraph.runtime import Runtime from textual.worker import WorkerCancelled from deepagents_code import offload +from deepagents_code._cli_context import CLIContextSchema from deepagents_code._session_stats import format_token_count -from deepagents_code._tracing import RESUME_TRACE_TAG -from deepagents_code.app import DeepAgentsApp, QueuedMessage +from deepagents_code.app import DeepAgentsApp from deepagents_code.command_registry import get_slash_commands from deepagents_code.configuration.types import TomlSnapshot from deepagents_code.hooks.manager import HooksManager @@ -80,23 +82,39 @@ def _summary_event( } -def _state_values( - messages: list[Any], event: dict[str, Any] | None = None -) -> dict[str, Any]: - """Build a thread state-values dict (as returned by _get_thread_state_values).""" - values: dict[str, Any] = {"messages": messages} - if event is not None: - values["_summarization_event"] = event - return values +def _compacted_result() -> dict[str, Any]: + """Build a successful server-owned offload result.""" + return { + "status": "compacted", + "messages_offloaded": 6, + "messages_kept": 4, + "tokens_before": 1000, + "tokens_after": 250, + "archive_path": "/conversation_history/test-thread.md", + "archive_ephemeral": False, + "error": None, + } 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 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. + 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() + agent.aoffload = AsyncMock() + app._agent = agent + app._backend = None + app._lc_thread_id = "test-thread" + app._agent_running = False + return agent + + +def _setup_local_offload_app(app: DeepAgentsApp) -> MagicMock: + """Configure a `DeepAgentsApp` with a local in-process agent.""" agent = MagicMock() agent.aupdate_state = AsyncMock() app._agent = agent @@ -123,68 +141,45 @@ def test_offload_sorted_alphabetically(self) -> None: assert model_idx < offload_idx < quit_idx -class TestOffloadGuards: - """Test guard conditions that prevent offloading.""" +class TestOffloadCommand: + """The TUI requests a typed operation and does not manage server state.""" async def test_no_agent_shows_error(self) -> None: - """Should show error when there is no active agent.""" app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - app._agent = None - app._lc_thread_id = None - await app._handle_offload() - await pilot.pause() - - msgs = app.query(AppMessage) - assert any("Nothing to offload" in str(w._content) for w in msgs) + assert any( + "Nothing to offload" in str(w._content) for w in app.query(AppMessage) + ) async def test_offload_while_busy_queues_instead_of_overlapping(self) -> None: - """A `/offload` submitted while busy must queue, not start a second run. - - `_handle_offload` reserves the turn (`_agent_running = True`) before its - first await, so by the time a same-tick duplicate submission is - processed `_submit_input` already sees the app busy and queues it. The - in-flight offload runs exactly once and the queued duplicate drains - afterward. - """ app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - - before = _state_values(_make_dict_messages(6)) - after = _state_values(_make_dict_messages(8), _summary_event(4)) + remote = _setup_server_offload_app(app) drive_started = asyncio.Event() release_drive = asyncio.Event() drive_calls = 0 - async def block_drive(_config: object, _seed_id: object = None) -> None: + async def block_offload(**_kwargs: Any) -> dict[str, Any]: nonlocal drive_calls drive_calls += 1 drive_started.set() - await release_drive.wait() + if drive_calls == 1: + await release_drive.wait() + return _compacted_result() + remote.aoffload = AsyncMock(side_effect=block_offload) with ( patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after, before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - side_effect=block_drive, + app, "_sync_session_cost_from_checkpoint", new=AsyncMock() ), + patch.object(app, "_run_session_start_hook", new=AsyncMock()), ): app.post_message(ChatInput.Submitted("/offload", "command")) await asyncio.wait_for(drive_started.wait(), timeout=1) - # The reservation is already in effect: a duplicate submission - # queues instead of entering `_handle_offload` concurrently. assert app._agent_running is True app.post_message(ChatInput.Submitted("/offload", "command")) await pilot.pause() @@ -197,986 +192,410 @@ async def block_drive(_config: object, _seed_id: object = None) -> None: await worker.wait() await pilot.pause() - # The queued duplicate drained after the first offload completed; - # both ran sequentially, never concurrently. assert drive_calls == 2 assert app._agent_running is False assert app._offload_worker is None assert not app._pending_messages - async def test_nothing_to_compact_noop(self) -> None: - """Show a no-op message when server-side compaction changed nothing. - - With `force=True` the eligibility gate is bypassed, so the only no-op - left is "cutoff == 0" — the persisted event is unchanged. - """ + async def test_server_result_is_rendered_without_reading_checkpoint_state( + self, + ) -> None: app = DeepAgentsApp() + result = { + "status": "compacted", + "messages_offloaded": 6, + "messages_kept": 4, + "tokens_before": 1000, + "tokens_after": 250, + "archive_path": "/conversation_history/test-thread.md", + "archive_ephemeral": False, + "error": None, + } async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - - before = _state_values(_make_dict_messages(3)) - after = _state_values(_make_dict_messages(3)) - after["_session_cost_usd"] = 0.75 - app._set_session_cost(0.5) - app._add_provisional_cost(0.75) - + remote = _setup_server_offload_app(app) + remote.aoffload = AsyncMock(return_value=result) with ( patch.object( app, "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], + new=AsyncMock(side_effect=AssertionError("client state read")), ), patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, + app, "_sync_session_cost_from_checkpoint", new=AsyncMock() ), + patch.object(app, "_run_session_start_hook", new=AsyncMock()), ): 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 - ) - # The graph prices the compaction run's own model call, so the - # committed total replaces the client's provisional estimate. - assert app._session_cost_usd == pytest.approx(0.75) - assert app._displayed_cost_usd == pytest.approx(0.75) - - async def test_empty_state_shows_error(self) -> None: - """Should show error when state has no values.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - app._agent = MagicMock() - app._backend = MagicMock() - app._lc_thread_id = "test-thread" - app._agent_running = False - - mock_state = MagicMock() - mock_state.values = {} - app._agent.aget_state = AsyncMock(return_value=mock_state) - - await app._handle_offload() - await pilot.pause() - - msgs = app.query(AppMessage) - assert any("Nothing to offload" in str(w._content) for w in msgs) + text = "\n".join(str(w._content) for w in app.query(AppMessage)) + assert "Offloaded 6 older messages" in text + assert "4 messages kept" in text + + remote.aoffload.assert_awaited_once() + await_args = remote.aoffload.await_args + assert await_args is not None + kwargs = await_args.kwargs + assert kwargs["config"] == {"configurable": {"thread_id": "test-thread"}} + assert "messages" not in kwargs["context"] + + async def test_context_carries_the_session_approval_mode(self) -> None: + """Hooks must see the session's real mode during `/offload`. + + The server defaults a missing `approval_mode` to `manual`, so omitting + it would show a configured `PreCompact`/`PreToolUse` hook Manual even in + YOLO -- a different mode than the same hook sees on every interactive + turn. + """ + from deepagents_code.approval_mode import ApprovalMode - async def test_state_read_failure_shows_error(self) -> None: - """Should show error when reading state raises an exception.""" app = DeepAgentsApp() + result = { + "status": "noop", + "messages_offloaded": 0, + "messages_kept": 1, + "tokens_before": 10, + "tokens_after": 10, + "archive_path": None, + "archive_ephemeral": False, + "error": None, + } async with app.run_test() as pilot: await pilot.pause() - app._agent = MagicMock() - app._backend = MagicMock() - app._lc_thread_id = "test-thread" - app._agent_running = False - - app._agent.aget_state = AsyncMock( - side_effect=RuntimeError("connection lost") - ) - - await app._handle_offload() - await pilot.pause() - - msgs = app.query(ErrorMessage) - assert any("Failed to read state" in str(w._content) for w in msgs) + remote = _setup_server_offload_app(app) + remote.aoffload = AsyncMock(return_value=result) + app._approval_mode = ApprovalMode.YOLO + app._auto_approve = True + with patch.object( + app, "_sync_session_cost_from_checkpoint", new=AsyncMock() + ): + await app._handle_offload() + await_args = remote.aoffload.await_args + assert await_args is not None + context = await_args.kwargs["context"] + assert context["approval_mode"] == "yolo" + assert context["auto_approve"] is True -class TestOffloadSuccess: - """Test successful offload flow.""" + async def test_failing_session_start_hook_does_not_erase_the_result(self) -> None: + """A hook raising after a committed compaction must not hide the outcome. - async def test_successful_offload_drives_server_tool(self) -> None: - """Should trigger server-side compaction and render persisted state.""" + The compaction is already durable server-side by this point, so letting + the hook's exception reach the generic handler would leave the user with + only "Offload failed" while their conversation really was compacted and + the status bar kept pre-offload counts. + """ app = DeepAgentsApp() + result = { + "status": "compacted", + "messages_offloaded": 6, + "messages_kept": 4, + "tokens_before": 1000, + "tokens_after": 250, + "archive_path": "/conversation_history/test-thread.md", + "archive_ephemeral": False, + "error": None, + } 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), - ) - + remote = _setup_server_offload_app(app) + remote.aoffload = AsyncMock(return_value=result) + tokens = MagicMock() with ( patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], + app, "_sync_session_cost_from_checkpoint", new=AsyncMock() ), patch.object( app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ) as mock_drive, + "_run_session_start_hook", + new=AsyncMock(side_effect=RuntimeError("hook spawn failed")), + ), + patch.object(app, "_on_tokens_update", new=tokens), ): await app._handle_offload() - await pilot.pause() - # The client drives the server-side tool exactly once and never - # writes `_summarization_event` itself — the tool owns that write. - mock_drive.assert_awaited_once() + text = "\n".join(str(w._content) for w in app.query(AppMessage)) + assert "Offloaded 6 older messages" in text + errors = "\n".join(str(w._content) for w in app.query(ErrorMessage)) + assert "SessionStart hook failed" in errors + assert "Offload failed" not in errors + tokens.assert_called_once_with(250, approximate=True) - msgs = app.query(AppMessage) - # 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_failing_report_does_not_report_a_committed_offload_as_failed( + self, + ) -> None: + """A rendering failure after the commit must not say "Offload failed". - async def test_committed_offload_survives_stream_failure(self) -> None: - """A checkpointed tool update wins over a later stream failure.""" + Everything between `aoffload` returning and the SessionStart hook is + local reporting over a conversation the server has already compacted. + Routing a failure there into the generic handler would tell the user to + offload again, compacting an already-compacted conversation. + """ app = DeepAgentsApp() + result = { + "status": "compacted", + "messages_offloaded": 6, + "messages_kept": 4, + "tokens_before": 1000, + "tokens_after": 250, + "archive_path": "/conversation_history/test-thread.md", + "archive_ephemeral": False, + "error": None, + } 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(4)) - + remote = _setup_server_offload_app(app) + remote.aoffload = AsyncMock(return_value=result) with ( patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], + app, "_sync_session_cost_from_checkpoint", new=AsyncMock() ), patch.object( app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - side_effect=RuntimeError("stream unavailable"), + "_on_tokens_update", + new=MagicMock(side_effect=RuntimeError("status bar exploded")), ), ): await app._handle_offload() - await pilot.pause() - assert any( - "Offloaded 4 older messages" in str(widget._content) - for widget in app.query(AppMessage) - ) - assert not any( - "Offload failed" in str(widget._content) - for widget in app.query(ErrorMessage) - ) + errors = "\n".join(str(w._content) for w in app.query(ErrorMessage)) + assert "could not be displayed" in errors + assert "Offload failed" not in errors + + async def test_session_start_hook_fires_after_a_committed_offload(self) -> None: + """The `COMPACT` lifecycle event still reaches configured hooks.""" + from deepagents_code.hooks.models.domain import SessionStartCause - async def test_offload_shows_feedback_message(self) -> None: - """Should report message and turn counts for offloaded and kept slices.""" app = DeepAgentsApp() + result = { + "status": "compacted", + "messages_offloaded": 2, + "messages_kept": 1, + "tokens_before": 100, + "tokens_after": 50, + "archive_path": "/conversation_history/test-thread.md", + "archive_ephemeral": False, + "error": None, + } 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(4)) - + remote = _setup_server_offload_app(app) + remote.aoffload = AsyncMock(return_value=result) + hook = AsyncMock() with ( patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, + app, "_sync_session_cost_from_checkpoint", new=AsyncMock() ), + patch.object(app, "_run_session_start_hook", new=hook), ): await app._handle_offload() - await pilot.pause() - msgs = app.query(AppMessage) - # Cutoff 4 against a prior cutoff of 0 over `_make_dict_messages(10)`, - # which alternates human/ai from index 0: 4 offloaded (humans at 0 and - # 2 -> 2 turns), 6 kept of the 10 before-messages (humans at 4, 6 and - # 8 -> 3 turns). - assert any( - "Offloaded 4 older messages (2 conversation turns)" - in str(widget._content) - for widget in msgs - ) - assert any( - "6 messages (3 conversation turns) kept" in str(widget._content) - for widget in msgs - ) - # No provider total was persisted, so the report is conversation-only. - assert any("Conversation: ~" in str(w._content) for w in msgs) + hook.assert_awaited_once_with(SessionStartCause.COMPACT) - async def test_kept_turns_ignore_tools_and_internal_messages(self) -> None: - """Turn counts should skip AI/tool rows and internal humans on both sides.""" + async def test_server_failure_is_rendered_from_typed_result(self) -> None: app = DeepAgentsApp() + result = { + "status": "failed", + "messages_offloaded": 0, + "messages_kept": 4, + "tokens_before": 100, + "tokens_after": 100, + "archive_path": None, + "archive_ephemeral": False, + "error": "summary unavailable", + } async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - - before_messages = [ - {"type": "human", "content": "Old prompt", "id": "old-human"}, - {"type": "ai", "content": "", "id": "old-tool-call"}, - { - "type": "tool", - "content": "Old result", - "id": "old-tool", - "tool_call_id": "old-call", - }, - {"type": "ai", "content": "Old answer", "id": "old-ai"}, - {"role": "user", "content": "Kept prompt", "id": "kept-human"}, - {"role": "assistant", "content": "", "id": "kept-tool-call"}, - { - "type": "tool", - "content": "Kept result", - "id": "kept-tool", - "tool_call_id": "kept-call", - }, - {"role": "assistant", "content": "Kept answer", "id": "kept-ai"}, - { - "type": "human", - "content": "Internal state", - "id": "internal-human", - "additional_kwargs": {"lc_source": "goal_state"}, - }, - ] - before = _state_values(before_messages) - after = _state_values( - [*before_messages, *_make_dict_messages(2)], _summary_event(4) - ) - + remote = _setup_server_offload_app(app) + remote.aoffload = AsyncMock(return_value=result) with ( patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, + app, "_sync_session_cost_from_checkpoint", new=AsyncMock() ), ): await app._handle_offload() - await pilot.pause() + assert any( + "summary unavailable" in str(w._content) + for w in app.query(ErrorMessage) + ) - messages = app.query(AppMessage) - assert any( - "Offloaded 4 older messages (1 conversation turn)" - in str(widget._content) - for widget in messages - ) + async def test_local_agent_shows_unsupported_message(self) -> None: + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + _setup_local_offload_app(app) + await app._handle_offload() assert any( - "5 messages (1 conversation turn) kept" in str(widget._content) - for widget in messages + "not supported for local agents" in str(widget._content) + for widget in app.query(AppMessage) ) - async def test_offloaded_turns_ignore_internal_messages(self) -> None: - """Offloaded turns apply the same internal filter as kept turns. - Goal-state notices and `[SYSTEM]`-prefixed humans accumulate over a long - thread, so they are likelier to fall in the offloaded slice than the kept - one. Both `is_internal_message` criteria are exercised here. - """ - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - _setup_server_offload_app(app) +def test_a_reasonless_refusal_cannot_be_built() -> None: + """A `denied`/`failed` result with no reason must not be constructible. - before_messages = [ - {"type": "human", "content": "Old prompt", "id": "old-human"}, - { - "type": "human", - "content": "Internal state", - "id": "offloaded-internal", - "additional_kwargs": {"lc_source": "goal_state"}, - }, - { - "type": "human", - "content": "[SYSTEM] Task interrupted by user.", - "id": "offloaded-system-prefix", - }, - {"type": "ai", "content": "Old answer", "id": "old-ai"}, - {"type": "human", "content": "Kept prompt", "id": "kept-human"}, - {"type": "ai", "content": "Kept answer", "id": "kept-ai"}, - ] - before = _state_values(before_messages) - after = _state_values( - [*before_messages, *_make_dict_messages(2)], _summary_event(4) - ) + The wire shape is one flat object, so `error` is `str | None` on every + status and the checker cannot make "a refusal has a reason" a compile-time + fact. A reasonless refusal reaches the user as the client's generic "the + server rejected the operation", which says nothing, so the single + construction point enforces it. + """ + from deepagents_code.offload_middleware import unchanged_offload_result - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() + for status in ("denied", "failed"): + with pytest.raises(ValueError, match="must carry a reason"): + unchanged_offload_result(status, messages=1, tokens=2) # ty: ignore[invalid-argument-type] - messages = app.query(AppMessage) - # Four messages offloaded, but only "Old prompt" is a real turn. - assert any( - "Offloaded 4 older messages (1 conversation turn)" - in str(widget._content) - for widget in messages - ) + # Unchanged, non-refusal outcomes legitimately carry no reason. + assert unchanged_offload_result("empty", messages=0, tokens=0)["error"] is None + + +class TestServerOffloadReporting: + """The server path reports its estimates with explicit metric labels.""" + + @staticmethod + def _result(**overrides: object) -> dict[str, object]: + """Build a `compacted` server result.""" + return { + "status": "compacted", + "messages_offloaded": 6, + "messages_kept": 4, + "tokens_before": 1000, + "tokens_after": 250, + "archive_path": "/conversation_history/test-thread.md", + "archive_ephemeral": False, + "error": None, + } | overrides + + async def _render(self, app: DeepAgentsApp, result: dict[str, object]) -> str: + """Drive `/offload` against a server result and return the rendered text.""" + remote = _setup_server_offload_app(app) + remote.aoffload = AsyncMock(return_value=result) + with ( + patch.object(app, "_sync_session_cost_from_checkpoint", new=AsyncMock()), + patch.object(app, "_run_session_start_hook", new=AsyncMock()), + ): + await app._handle_offload() + return "\n".join(str(w._content) for w in app.query(AppMessage)) + "\n".join( + str(w._content) for w in app.query(ErrorMessage) + ) - async def test_singular_kept_message_and_zero_offloaded_turns(self) -> None: - """Kept singular and zero-turn plural render correctly. + async def test_estimates_are_labelled_conversation_and_marked(self) -> None: + """Server figures are conversation-scale estimates, not context totals. - A cutoff that leaves exactly one message covers `messages_kept == 1`, and - an offloaded slice of pure AI/tool traffic covers `turns_offloaded == 0` -- - the plural-zero branch, which no other test renders. + "Conversation" excludes the system/tool overhead that "Context" + includes, so labelling an estimate "Context" invites the user to compare + two percentages that are not comparable across offloads. """ app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - - before_messages = [ - {"type": "ai", "content": "", "id": "old-tool-call"}, - { - "type": "tool", - "content": "Old result", - "id": "old-tool", - "tool_call_id": "old-call", - }, - {"type": "ai", "content": "Old answer", "id": "old-ai"}, - {"type": "human", "content": "Kept prompt", "id": "kept-human"}, - ] - before = _state_values(before_messages) - after = _state_values( - [*before_messages, *_make_dict_messages(2)], _summary_event(3) - ) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() + app._context_tokens = 0 + app._tokens_approximate = True + text = await self._render(app, self._result()) - messages = app.query(AppMessage) - # Three offloaded rows, none of them a human turn. - assert any( - "Offloaded 3 older messages (0 conversation turns)" - in str(widget._content) - for widget in messages - ) - # Exactly one kept message, which is a human turn: both singulars. - assert any( - "1 message (1 conversation turn) kept" in str(widget._content) - for widget in messages - ) + assert "Conversation: ~1.0K → ~250 tokens (75% decrease)" in text + assert "Context:" not in text - async def test_zero_kept_turns_when_cutoff_reaches_end(self) -> None: - """A cutoff past the last human renders `0 conversation turns` kept.""" + async def test_a_larger_summary_never_reports_a_negative_decrease(self) -> None: + """A summary can exceed what it replaced; that is an increase, not -14%.""" app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - - before_messages = [ - {"type": "human", "content": "Old prompt", "id": "old-human"}, - {"type": "ai", "content": "Old answer", "id": "old-ai"}, - {"type": "ai", "content": "Trailing", "id": "trailing-ai"}, - ] - before = _state_values(before_messages) - after = _state_values( - [*before_messages, *_make_dict_messages(2)], _summary_event(2) + app._context_tokens = 0 + app._tokens_approximate = True + text = await self._render( + app, self._result(tokens_before=100, tokens_after=114) ) - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - messages = app.query(AppMessage) - assert any( - "1 message (0 conversation turns) kept" in str(widget._content) - for widget in messages - ) + assert "(increase)" in text + assert "-" not in text.split("tokens")[1].split(",")[0] + assert "summary was larger than the messages it replaced" in text - async def test_offload_updates_context_tokens(self) -> None: - """Should update `_context_tokens` to the post-compaction count. + async def test_a_real_provider_total_promotes_the_report_to_context(self) -> None: + """A cached count from a real turn is the provider's own total. - The count is taken from the pre-seed conversation plus the new event, so - it excludes the tool's own machinery (the seeded call, the tool result, - and the trailing model turn) that the post-run state carries. Using - distinct before/after message lists guards against regressing to the - post-run state, which would understate the reduction. + The delta is subtracted from that total rather than rebuilt as + `overhead + after`, so both figures stay on the provider's scale. """ - from langchain_core.messages.utils import count_tokens_approximately + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + app._context_tokens = 5000 + app._tokens_approximate = False + text = await self._render(app, self._result()) - from deepagents_code.app import _effective_conversation + # 5000 - (1000 - 250) = 4250; `before` is exact, only `after` estimated. + assert "Context: 5.0K → ~4.2K tokens (15% decrease)" in text + async def test_ephemeral_storage_is_disclosed(self) -> None: + """History in a temp fallback must not be presented as durable.""" app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) + app._context_tokens = 0 + app._tokens_approximate = True + text = await self._render(app, self._result(archive_ephemeral=True)) - before_messages = _make_dict_messages(10) - after_messages = _make_dict_messages(12) - after_event = _summary_event(4) - before = _state_values(before_messages) - after = _state_values(after_messages, after_event) + assert "may not survive a restart" in text - expected = count_tokens_approximately( - _effective_conversation(before_messages, after_event) - ) + async def test_durable_storage_adds_no_caveat(self) -> None: + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + app._context_tokens = 0 + app._tokens_approximate = True + text = await self._render(app, self._result(archive_ephemeral=False)) - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() + assert "may not survive" not in text - assert app._context_tokens == expected + async def test_a_failed_archive_write_reports_unrecoverable_messages(self) -> None: + """Context was freed but the history is gone; both facts must be said. - async def test_offload_preserves_fixed_overhead_in_context_report(self) -> None: - """Provider totals should keep fixed overhead in the post-offload estimate.""" - from langchain_core.messages.utils import count_tokens_approximately + This is data-loss messaging: reporting plain success here would tell the + user their conversation is archived when it is not. + """ + app = DeepAgentsApp() + async with app.run_test() as pilot: + await pilot.pause() + app._context_tokens = 0 + app._tokens_approximate = True + text = await self._render(app, self._result(archive_path=None)) + errors = [str(w._content) for w in app.query(ErrorMessage)] - from deepagents_code.app import _effective_conversation + assert "not recoverable" in text + # An error, not a success message: the offload did not fully succeed. + assert errors + assert "not recoverable" in "\n".join(errors) + async def test_singular_message_labels(self) -> None: app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - - before_messages = _make_dict_messages(10) - after_event = _summary_event(4) - conversation_before = count_tokens_approximately(before_messages) - conversation_after = count_tokens_approximately( - _effective_conversation(before_messages, after_event) + app._context_tokens = 0 + app._tokens_approximate = True + text = await self._render( + app, self._result(messages_offloaded=1, messages_kept=1) ) - fixed_tokens = 50_000 - reported_before = conversation_before + fixed_tokens - expected_after = conversation_after + fixed_tokens - before = _state_values(before_messages) - before["_context_tokens"] = reported_before - after = _state_values( - [*before_messages, *_make_dict_messages(2)], after_event - ) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - expected_report = ( - f"Context: {format_token_count(reported_before)} → " - f"~{format_token_count(expected_after)} tokens" - ) - assert any( - expected_report in str(widget._content) - for widget in app.query(AppMessage) - ) - assert app._context_tokens == expected_after - assert app._tokens_approximate is True + assert "1 older message," in text + assert "1 message kept" in text - async def test_offload_report_stays_on_provider_scale_when_total_is_low( - self, - ) -> None: - """A provider total below the local estimate must not free fixed overhead. - - When `_context_tokens` is stale (or the approximation overshoots), the - reported total can fall below `conversation_tokens_before`. The report must - still subtract only the conversation *delta*, keeping both figures on the - provider's scale -- rebuilding the after-figure as - `max(0, reported - conversation_before) + conversation_after` would collapse - the overhead to zero and credit the offload with freeing the whole system - prompt and tool schema. - """ - from langchain_core.messages.utils import count_tokens_approximately - from deepagents_code.app import _effective_conversation +class TestOffloadInterrupt: + """Test that Escape can cancel `/offload` through the real App dispatch.""" + async def test_command_reserves_turn_before_worker_starts(self) -> None: app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) + worker = MagicMock() + scheduled: list[Coroutine[Any, Any, None]] = [] - before_messages = _make_dict_messages(10) - after_event = _summary_event(4) - conversation_before = count_tokens_approximately(before_messages) - conversation_after = count_tokens_approximately( - _effective_conversation(before_messages, after_event) - ) - # Stale/low provider total: below the local conversation estimate. - reported_before = conversation_before // 2 - expected_after = reported_before - ( - conversation_before - conversation_after - ) - assert expected_after > 0, "fixture should not exercise the zero floor" - before = _state_values(before_messages) - before["_context_tokens"] = reported_before - after = _state_values( - [*before_messages, *_make_dict_messages(2)], after_event - ) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - contents = [str(widget._content) for widget in app.query(AppMessage)] - expected_report = ( - f"Context: {format_token_count(reported_before)} → " - f"~{format_token_count(expected_after)} tokens" - ) - assert any(expected_report in content for content in contents) - # The overhead was never treated as freed, so the reduction stays - # modest rather than approaching 100%. - assert not any("(100% decrease)" in content for content in contents) - assert app._context_tokens == expected_after - - async def test_offload_reports_oversized_summary_as_increase(self) -> None: - """A summary larger than what it replaced is reported as an increase.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - _setup_server_offload_app(app) - - before_messages = _make_dict_messages(10) - # A summary far longer than the four messages it replaces, so - # `tokens_after` exceeds `tokens_before`. - after_event = _summary_event(4) - after_event["summary_message"]["content"] = "verbose summary " * 500 - before = _state_values(before_messages) - after = _state_values( - [*before_messages, *_make_dict_messages(2)], after_event - ) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - contents = [str(widget._content) for widget in app.query(AppMessage)] - assert any("Offloaded " in content for content in contents) - assert any("(increase)" in content for content in contents) - assert not any( - "freeing up context window space" in content for content in contents - ) - assert any("context increased" in content for content in contents) - - async def test_no_ui_clear_reload(self) -> None: - """Should NOT clear/reload UI since messages stay in state.""" - 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(4)) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - patch.object( - app, "_clear_messages", new_callable=AsyncMock - ) as mock_clear, - patch.object( - app, "_load_thread_history", new_callable=AsyncMock - ) as mock_load, - ): - await app._handle_offload() - await pilot.pause() - - mock_clear.assert_not_called() - mock_load.assert_not_called() - - -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.""" - 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) - artifacts = [ - { - "type": "ai", - "content": "", - "id": "offload-seed-test", - "tool_calls": [ - { - "name": "compact_conversation", - "args": {"force": True}, - "id": "seed-call", - } - ], - }, - { - "type": "tool", - "content": "Nothing to compact yet.", - "id": "offload-result-test", - "tool_call_id": "seed-call", - }, - { - "type": "ai", - "content": "Trailing response", - "id": "offload-trailing-test", - "tool_calls": [], - }, - ] - before = _state_values(messages, event) - after = _state_values([*messages, *artifacts], event) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - 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_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", - ] - - async def test_cutoff_one_offloads_single_message(self) -> None: - """A cutoff of 1 reports a single offloaded message.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - _setup_server_offload_app(app) - - before = _state_values(_make_dict_messages(7)) - after = _state_values(_make_dict_messages(9), _summary_event(1)) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - msgs = app.query(AppMessage) - # Anchored on both parens: each singular is a prefix of its plural, so - # "(" rules out "1 older messages" and ")" rules out "1 conversation - # turns". An unanchored substring would pass against hardcoded plurals. - assert any( - "Offloaded 1 older message (1 conversation turn)" in str(w._content) - for w in msgs - ) - - -class TestReOffload: - """Test offload when a prior _summarization_event already exists.""" - - async def test_reoffload_uses_absolute_cutoff_delta(self) -> None: - """Re-offload counts only the newly offloaded messages. - - With a prior cutoff of 5 and a new absolute cutoff of 7, exactly two - additional messages were offloaded this run. - """ - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - _setup_server_offload_app(app) - - prior_event = _summary_event(5, file_path=None) - before = _state_values(_make_dict_messages(15), prior_event) - after = _state_values(_make_dict_messages(17), _summary_event(7)) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - msgs = app.query(AppMessage) - # Offloaded count is the new cutoff of seven minus a prior cutoff of - # five. The turn parenthetical is asserted too: `_make_dict_messages` - # puts humans on even indices, so the offloaded slice [5:7] holds - # exactly one. Dropping `prior_cutoff` (slicing [:7]) would report - # four, so this is the only assertion that pins the slice *start*. - assert any( - "Offloaded 2 older messages (1 conversation turn)" in str(w._content) - for w in msgs - ) - # Kept: 15 before-messages minus the cutoff of seven, with humans at - # indices 8, 10, 12 and 14. - assert any( - "8 messages (4 conversation turns) kept" in str(w._content) - for w in msgs - ) - - async def test_reoffload_noop_restores_prior_summary(self) -> None: - """A summary-only re-offload restores the prior summarization event.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - agent = _setup_server_offload_app(app) - - prior_event = _summary_event(5, file_path=None) - replacement_event = _summary_event(5) - replacement_event["summary_message"]["content"] = "Replacement summary." - before_messages = _make_dict_messages(11) - after_messages = [*before_messages, *_make_dict_messages(2)] - after_messages[-2]["id"] = "offload-seed" - after_messages[-1]["id"] = "offload-result" - before = _state_values(before_messages, prior_event) - after = _state_values(after_messages, replacement_event) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=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", - ] - assert any( - "Nothing to offload" in str(widget._content) - for widget in app.query(AppMessage) - ) - - -class TestAgentRunningGuard: - """Test that _handle_offload sets _agent_running to prevent races.""" - - async def test_agent_running_set_during_offload(self) -> None: - """Should set _agent_running=True during offload and reset after.""" - 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(4)) - - running_during_offload: list[bool] = [] - quiescent_during_offload: list[bool] = [] - - def capture_running(_config: object, _seed_id: object = None) -> None: - running_during_offload.append(app._agent_running) - quiescent_during_offload.append(app._agent_quiescent.is_set()) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - side_effect=capture_running, - ), - ): - await app._handle_offload() - await pilot.pause() - - # _agent_running should have been True while the tool ran - assert running_during_offload == [True] - assert quiescent_during_offload == [False] - # And reset after completion - assert app._agent_running is False - assert app._agent_quiescent.is_set() - - async def test_agent_running_reset_after_failure(self) -> None: - """Should reset _agent_running=False even when offload fails.""" - 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_server_side_compaction", - new_callable=AsyncMock, - side_effect=RuntimeError("stream down"), - ), - ): - await app._handle_offload() - await pilot.pause() - - assert app._agent_running is False - - -class TestOffloadInterrupt: - """Test that Escape can cancel `/offload` through the real App dispatch.""" - - async def test_command_reserves_turn_before_worker_starts(self) -> None: - """Command dispatch should reserve busy state before scheduling offload.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - worker = MagicMock() - scheduled: list[Coroutine[Any, Any, None]] = [] - - def defer_worker( - work: Coroutine[Any, Any, None], **_kwargs: object - ) -> MagicMock: - scheduled.append(work) - return worker + def defer_worker( + work: Coroutine[Any, Any, None], **_kwargs: object + ) -> MagicMock: + scheduled.append(work) + return worker with patch.object(app, "run_worker", side_effect=defer_worker): await app._handle_command("/offload") @@ -1199,514 +618,96 @@ def defer_worker( assert app._offload_worker is None assert not app._pending_messages - async def test_escape_cancels_offload_worker(self) -> None: - """Escape should reach and cancel an offload blocked in its stream.""" + async def test_escape_cancels_server_owned_offload(self) -> None: 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)) + remote = _setup_server_offload_app(app) drive_started = asyncio.Event() drive_cancelled = asyncio.Event() - async def block_drive(_config: object, _seed_id: object = None) -> None: + async def block_offload(**_kwargs: Any) -> dict[str, Any]: drive_started.set() try: await asyncio.Future() finally: drive_cancelled.set() + return _compacted_result() - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, reconciled], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - side_effect=block_drive, - ), - patch.object( - app, - "_remove_unanswered_offload_seed", - new_callable=AsyncMock, - return_value=True, - ) as cleanup, - ): - app.post_message(ChatInput.Submitted("/offload", "command")) - await asyncio.wait_for(drive_started.wait(), timeout=1) + remote.aoffload = AsyncMock(side_effect=block_offload) + app.post_message(ChatInput.Submitted("/offload", "command")) + await asyncio.wait_for(drive_started.wait(), timeout=1) - worker = app._offload_worker - assert worker is not None - assert app._agent_running is True + worker = app._offload_worker + assert worker is not None + assert app._agent_running is True - await pilot.press("escape") - await asyncio.wait_for(drive_cancelled.wait(), timeout=1) - with pytest.raises(WorkerCancelled): - await worker.wait() + await pilot.press("escape") + await asyncio.wait_for(drive_cancelled.wait(), timeout=1) + with pytest.raises(WorkerCancelled): + await worker.wait() - cleanup.assert_awaited_once() assert worker.is_cancelled assert app._agent_running is False assert app._agent_quiescent.is_set() assert app._loading_widget is None async def test_offload_blocks_queued_prompt_until_done(self) -> None: - """A prompt submitted while offload is reserved must not overlap it. - - Command dispatch sets `_agent_running` before scheduling the worker, so - a prompt queued before the worker starts stays queued until the offload - finishes and `_run_offload_task` drains the queue. - """ app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) - - # The drive blocks the offload so the reservation window stays - # open. It runs on the real loop (not a Mock) because pilot.pause - # must be able to yield between message-processing steps below. + remote = _setup_server_offload_app(app) drive_started = asyncio.Event() release_drive = asyncio.Event() - async def block_drive(_config: object, _seed_id: object = None) -> None: + async def block_offload(**_kwargs: Any) -> dict[str, Any]: drive_started.set() await release_drive.wait() + return _compacted_result() + remote.aoffload = AsyncMock(side_effect=block_offload) + dispatch = AsyncMock() with ( patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[ - _state_values(_make_dict_messages(6)), - _state_values(_make_dict_messages(8), _summary_event(4)), - ], - ), - patch.object( - app, - "_drive_server_side_compaction", - side_effect=block_drive, + app, "_sync_session_cost_from_checkpoint", new=AsyncMock() ), - # Spy on the queue drain itself: patching `_dispatch_queued_message` - # would let the drain's awaited dispatch return instantly, so the - # worker could complete before the test's own `pilot.pause()`s - # observe the mid-offload state. - patch.object( - app, - "_process_next_from_queue", - wraps=app._process_next_from_queue, - ) as drain, + patch.object(app, "_run_session_start_hook", new=AsyncMock()), ): app.post_message(ChatInput.Submitted("/offload", "command")) - await pilot.pause() await asyncio.wait_for(drive_started.wait(), timeout=1) - # The reservation is already in effect: a prompt submitted now - # queues instead of starting an agent turn concurrently. - assert app._agent_running is True - app.post_message(ChatInput.Submitted("hello", "prompt")) - await pilot.pause() - assert len(app._pending_messages) == 1 + with patch.object(app, "_dispatch_queued_message", new=dispatch): + app.post_message(ChatInput.Submitted("hello", "prompt")) + await pilot.pause() + assert app._agent_running is True + assert len(app._pending_messages) == 1 + dispatch.assert_not_awaited() - release_drive.set() - worker = app._offload_worker - assert worker is not None - await worker.wait() - await pilot.pause() + release_drive.set() + worker = app._offload_worker + assert worker is not None + await worker.wait() + await pilot.pause() - # The worker's teardown drained the queued prompt only after the - # offload completed. - drain.assert_called() + dispatch.assert_awaited_once() assert app._agent_running is False assert app._offload_worker is None assert not app._pending_messages - async def test_escape_before_offload_worker_step_recovers(self) -> None: - """Cancelling the worker mid-run releases it for later `Esc` presses. - - `Esc` routes to `_cancel_worker(self._offload_worker)`. If the offload - worker then ends without clearing `_offload_worker` (the same wedge - `_recover_unstarted_agent_worker` covers for the agent worker when a - cancel lands before the worker's first step), every later `Esc` would - be consumed re-cancelling the dead worker instead of reaching the - double-`Esc` input-clear path. - """ + async def test_server_failure_releases_busy_state_and_spinner(self) -> None: app = DeepAgentsApp() async with app.run_test() as pilot: await pilot.pause() - _setup_server_offload_app(app) + remote = _setup_server_offload_app(app) + remote.aoffload = AsyncMock(side_effect=RuntimeError("server unavailable")) - # Hold the offload at its first await (the state read) so the - # cancel lands while the worker is in flight. - hold_state_read = asyncio.Event() + with patch.object(app, "_set_spinner", new_callable=AsyncMock) as spinner: + await app._handle_offload() - async def block_state_read(_thread_id: object) -> dict[str, Any]: - await hold_state_read.wait() - return _state_values(_make_dict_messages(6)) - - with patch.object( - app, - "_get_thread_state_values", - side_effect=block_state_read, - ): - app.post_message(ChatInput.Submitted("/offload", "command")) - await pilot.pause() - worker = app._offload_worker - assert worker is not None - # Command dispatch reserves the turn before the worker starts, - # so it is set even though compaction never got to run. - assert app._agent_running is True - - await pilot.press("escape") - await pilot.pause() - await pilot.pause() - - assert worker.is_cancelled - assert app._offload_worker is None - assert app._agent_running is False - - # The next `Esc` is no longer consumed by a stale offload - # worker: it reaches the double-`Esc` clear-input path. - app._chat_input.value = "draft" - await pilot.press("escape") - assert app._clear_input_pending is True - - -class TestOffloadErrorHandling: - """Test error handling during offload.""" - - async def test_missing_archive_path_warns_about_unrecoverable_history( - self, - ) -> None: - """A failed backend write surfaces in a single, non-contradictory message. - - The reduction and the unrecoverable-archive warning are combined into one - `ErrorMessage` rather than a warning immediately followed by a separate - success line. - """ - 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(4, file_path=None) - ) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - # Both the reduction and the archive-failure warning land in one - # ErrorMessage. The turn parenthetical and the trailing stats line are - # asserted too, so the error path cannot silently keep older wording. - assert any( - "Offloaded 4 older messages (2 conversation turns), " - "freeing up context window space." - in str(widget._content) - and "could not be saved to storage" in str(widget._content) - and "conversation turns) kept." in str(widget._content) - for widget in app.query(ErrorMessage) - ) - # No separate success line is emitted alongside the warning. - assert not any( - "Offloaded" in str(widget._content) for widget in app.query(AppMessage) - ) - - async def test_tool_reported_compaction_failure_shows_error(self) -> None: - """A "Compaction failed" ToolMessage surfaces as an `ErrorMessage`.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - _setup_server_offload_app(app) - - before = _state_values(_make_dict_messages(10)) - tool_error = ( - "Compaction failed: an error occurred while generating the " - "summary (RuntimeError: model unavailable)." - ) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=tool_error, - ), - ): - await app._handle_offload() - await pilot.pause() - - error_msgs = app.query(ErrorMessage) - assert any("Compaction failed" in str(w._content) for w in error_msgs) - # A no-success guarantee: the offloaded feedback is not shown. - assert not any( - "Offloaded " in str(w._content) for w in app.query(AppMessage) - ) - - async def test_stale_compaction_failure_is_not_reported(self) -> None: - """A no-op ignores failure messages committed by an earlier run.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - _setup_server_offload_app(app) - - messages = [ - *_make_dict_messages(3), - { - "type": "tool", - "content": "Compaction failed: old failure", - "tool_call_id": "old-call", - }, - ] - before = _state_values(messages) - after = _state_values([*messages, *_make_dict_messages(1)]) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - assert not any( - "old failure" in str(widget._content) - for widget in app.query(ErrorMessage) - ) - assert any( - "the conversation is already compact" in str(widget._content) - for widget in app.query(AppMessage) - ) - - async def test_current_durable_compaction_failure_is_reported(self) -> None: - """A failure appended by this invocation survives a missed stream event.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - _setup_server_offload_app(app) - - messages = _make_dict_messages(3) - before = _state_values(messages) - after = _state_values( - [ - *messages, - { - "type": "tool", - "content": "Compaction failed: current failure", - "tool_call_id": "current-call", - }, - ] - ) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - assert any( - "current failure" in str(widget._content) - 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. - - 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. - """ - 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_server_side_compaction", - 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_awaited_once() - assert any( - "Offload failed" in str(widget._content) - 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. - - 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. - """ - 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_server_side_compaction", - 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_awaited_once() - error_text = " ".join( - str(widget._content) for widget in app.query(ErrorMessage) - ) - assert "inconsistent state" in error_text - assert "Offload failed" in error_text - - async def test_compaction_run_failure_shows_error(self) -> None: - """Should show error and leave state untouched when the run raises.""" - 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_server_side_compaction", - new_callable=AsyncMock, - side_effect=RuntimeError("stream unavailable"), - ), - ): - await app._handle_offload() - await pilot.pause() - - error_msgs = app.query(ErrorMessage) - assert any("Offload failed" in str(w._content) for w in error_msgs) - - async def test_spinner_hidden_after_failure(self) -> None: - """Should hide spinner even when offload fails.""" - 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_server_side_compaction", - new_callable=AsyncMock, - side_effect=RuntimeError("backend down"), - ), - patch.object( - app, "_set_spinner", new_callable=AsyncMock - ) as mock_spinner, - ): - await app._handle_offload() - await pilot.pause() - - # Spinner should be shown then hidden - assert mock_spinner.call_count == 2 - mock_spinner.assert_any_call("Offloading") - mock_spinner.assert_any_call(None) + spinner.assert_any_await("Offloading") + spinner.assert_awaited_with(None) + assert app._agent_running is False + assert app._agent_quiescent.is_set() class TestOffloadFallbackRoot: @@ -1799,8 +800,6 @@ def test_fallback_root_rejects_foreign_owned_per_user_dir( report a foreign owner for the predictable per-user dir only, so it is rejected while the freshly-created unique dir (real ownership) passes. """ - from types import SimpleNamespace - getuid = getattr(os, "getuid", None) if getuid is None: pytest.skip("uid ownership check requires os.getuid") @@ -2187,1245 +1186,149 @@ def test_unlink_failure_is_swallowed( ) assert sweep_offloaded_history() == 0 - assert archive.exists() - - def test_archive_refreshed_between_iterdir_and_unlink_is_kept( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ) -> None: - """An archive rewritten after the sweep lists it must not be deleted. - - Simulates a second `dcode` process refreshing an expired archive after - this process's `iterdir()` has already enumerated it: the pre-unlink - `fstat` observes the refreshed mtime and keeps the file, so the rewrite - is not orphaned by a stale expiry decision. - """ - archive_dir = self._setup(tmp_path, monkeypatch) - archive = archive_dir / "old.md" - archive.write_text("old") - old_time = time.time() - 31 * 86_400 - os.utime(archive, (old_time, old_time)) - - real_fstat = os.fstat - refreshed = False - - def fstat_with_refresh(fd: int) -> os.stat_result: - nonlocal refreshed - if not refreshed: - refreshed = True - # The racing writer rewrites the archive before the sweep's - # fstat lands, making it fresh again. - archive.write_text("refreshed") - fresh_time = time.time() - os.utime(archive, (fresh_time, fresh_time)) - return real_fstat(fd) - - monkeypatch.setattr(os, "fstat", fstat_with_refresh) - - assert sweep_offloaded_history() == 0 - assert archive.exists() - assert archive.read_text() == "refreshed" - - -class TestArtifactsRoot: - """Cover the real-filesystem artifacts root for offloaded tool results.""" - - def test_artifacts_root_is_stable_and_hardened( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ) -> None: - """The per-user artifacts dir is predictable, private, and reused.""" - temp_dir = tmp_path / "tmp" - temp_dir.mkdir() - getuid = getattr(os, "getuid", None) - uid = getuid() if getuid is not None else os.getpid() - - monkeypatch.setattr(tempfile, "gettempdir", lambda: str(temp_dir)) - - storage = _artifacts_root() - root_path = Path(storage.root) - - assert storage.large_results_dir is None - assert root_path.samefile(temp_dir / f"dcode-artifacts-{uid}") - assert stat.S_IMODE(root_path.stat().st_mode) == 0o700 - # Stable across calls (paths embedded in resumed threads stay resolvable). - assert _artifacts_root() == storage - - def test_windows_artifacts_root_is_accepted_by_filesystem_tools(self) -> None: - """A Windows temp path retains its drive without a rejected drive prefix.""" - disk_root = PureWindowsPath( - "C:/Users/test/AppData/Local/Temp/dcode-artifacts-123" - ) - - root = _filesystem_tool_path(disk_root) - result_path = f"{root}/large_tool_results/tool-call-id" - - assert root == "//?/C:/Users/test/AppData/Local/Temp/dcode-artifacts-123" - assert PureWindowsPath(root).is_absolute() - assert validate_path(result_path) == result_path - - def test_artifacts_root_falls_back_when_predictable_path_foreign_owned( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ) -> None: - """A predictable dir owned by another user is rejected for a unique one.""" - from types import SimpleNamespace - - getuid = getattr(os, "getuid", None) - if getuid is None: - pytest.skip("uid ownership check requires os.getuid") - - temp_dir = tmp_path / "tmp" - temp_dir.mkdir() - uid = getuid() - reserved = temp_dir / f"dcode-artifacts-{uid}" - reserved.mkdir() # a real, us-owned directory; lstat is faked below - - real_lstat = Path.lstat - - def fake_lstat(self: Path) -> Any: # noqa: ANN401 - info = real_lstat(self) - if self == reserved: - return SimpleNamespace(st_mode=info.st_mode, st_uid=info.st_uid + 1) - return info - - monkeypatch.setattr(tempfile, "gettempdir", lambda: str(temp_dir)) - monkeypatch.setattr(Path, "lstat", fake_lstat) - - storage = _artifacts_root() - next_storage = _artifacts_root() - - assert storage.root == "/dcode-artifacts-fallback" - assert next_storage.root == storage.root - assert storage.large_results_dir is not None - assert next_storage.large_results_dir is not None - assert not storage.large_results_dir.samefile(reserved) - assert storage.large_results_dir.name.startswith(f"dcode-artifacts-{uid}-") - assert stat.S_IMODE(storage.large_results_dir.stat().st_mode) == 0o700 - assert next_storage.large_results_dir != storage.large_results_dir - - -class TestOffloadStorageCaveat: - """Surface the persistence caveat when offload uses ephemeral storage.""" - - async def test_ephemeral_storage_appends_caveat_to_success(self) -> None: - """A successful offload into temp storage warns it may not persist.""" - 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)) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - patch( - "deepagents_code.offload.offload_storage_is_ephemeral", - return_value=True, - ), - ): - await app._handle_offload() - await pilot.pause() - - msgs = app.query(AppMessage) - assert any("Offloaded 6 older messages" in str(w._content) for w in msgs) - assert any("may not survive a restart" in str(w._content) for w in msgs) - - async def test_persistent_storage_omits_caveat(self) -> None: - """A successful offload into persistent storage adds no caveat.""" - 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)) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - patch( - "deepagents_code.offload.offload_storage_is_ephemeral", - return_value=False, - ), - ): - await app._handle_offload() - await pilot.pause() - - msgs = app.query(AppMessage) - assert any("Offloaded 6 older messages" in str(w._content) for w in msgs) - assert not any("may not survive a restart" in str(w._content) for w in msgs) - - -class TestNoopArtifactCleanup: - """A failed no-op restoration must not be reported as an offload failure.""" - - async def test_cleanup_failure_keeps_noop_report(self) -> None: - """When restoration fails, still report the no-op — not "Offload failed".""" - 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. - agent.aupdate_state = AsyncMock(side_effect=RuntimeError("write failed")) - - before = _state_values(_make_dict_messages(4)) - after = _state_values(_make_dict_messages(6)) - - with ( - patch.object( - app, - "_get_thread_state_values", - new_callable=AsyncMock, - side_effect=[before, after], - ), - patch.object( - app, - "_drive_server_side_compaction", - new_callable=AsyncMock, - return_value=None, - ), - ): - await app._handle_offload() - await pilot.pause() - - assert any( - "the conversation is already compact" in str(w._content) - for w in app.query(AppMessage) - ) - assert not any( - "Offload failed" in str(w._content) for w in app.query(ErrorMessage) - ) - - -class TestOffloadRouting: - """Test that /offload is routed through _handle_command.""" - - async def test_offload_routed_from_handle_command(self) -> None: - """'/offload' should be correctly routed through _handle_command.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - app._agent = None - app._lc_thread_id = None - - await app._handle_command("/offload") - await pilot.pause() - - msgs = app.query(AppMessage) - assert any("Nothing to offload" in str(w._content) for w in msgs) - - async def test_compact_alias_routed_from_handle_command(self) -> None: - """'/compact' should still route through _handle_command for backward compat.""" - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - app._agent = None - app._lc_thread_id = None - - await app._handle_command("/compact") - await pilot.pause() - - msgs = app.query(AppMessage) - assert any("Nothing to offload" in str(w._content) for w in msgs) - - -class TestOffloadToolGuard: - """Server-side tool execution guard for hidden `/offload` turns.""" - - @pytest.mark.parametrize( - "tool_call", - [ - {"name": "write_file", "args": {"path": "x"}, "id": "model-call"}, - { - "name": "compact_conversation", - "args": {"force": True}, - # Even reusing the authorized ID cannot turn a later model - # message into the one server-seeded call. - "id": "seed-call", - }, - ], - ) - async def test_blocks_every_call_except_seed( - self, tool_call: dict[str, Any] - ) -> None: - """Unrelated and repeated tools never reach their execution handler.""" - from langchain_core.messages import ToolMessage - - from deepagents_code.offload_middleware import CLICompactionMiddleware - - middleware = object.__new__(CLICompactionMiddleware) - request = MagicMock() - request.runtime.context = {"offload_tool_call_id": "seed-call"} - request.tool_call = tool_call - request.state = {"messages": [{"id": "model-generated-message"}]} - handler = AsyncMock() - - result = await middleware.awrap_tool_call(request, handler) - - assert isinstance(result, ToolMessage) - assert result.status == "error" - handler.assert_not_awaited() - - async def test_allows_exact_seeded_compaction(self) -> None: - """The one forced call seeded by `/offload` reaches the tool handler.""" - from langchain_core.messages import ToolMessage - - from deepagents_code.offload_middleware import CLICompactionMiddleware - - middleware = object.__new__(CLICompactionMiddleware) - request = MagicMock() - request.runtime.context = {"offload_tool_call_id": "seed-call"} - request.tool_call = { - "name": "compact_conversation", - "args": {"force": True}, - "id": "seed-call", - } - request.state = {"messages": [{"id": "offload-seed-seed-call"}]} - expected = ToolMessage(content="done", tool_call_id="seed-call") - handler = AsyncMock(return_value=expected) - - result = await middleware.awrap_tool_call(request, handler) - - assert result is expected - handler.assert_awaited_once_with(request) - - async def test_ordinary_runs_are_unchanged(self) -> None: - """Without `/offload` context, normal tools pass through the guard.""" - from langchain_core.messages import ToolMessage - - from deepagents_code.offload_middleware import CLICompactionMiddleware - - middleware = object.__new__(CLICompactionMiddleware) - request = MagicMock() - request.runtime.context = {} - request.tool_call = {"name": "write_file", "args": {}, "id": "normal-call"} - expected = ToolMessage(content="done", tool_call_id="normal-call") - handler = AsyncMock(return_value=expected) - - result = await middleware.awrap_tool_call(request, handler) - - assert result is expected - handler.assert_awaited_once_with(request) - - -class TestDriveServerSideCompaction: - """Unit-test the server-side `compact_conversation` trigger mechanism.""" - - @staticmethod - def _fake_remote_agent( - tool_content: str, - ) -> tuple[Any, list[Any], list[object], list[Any]]: - """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. - - Args: - tool_content: Body of the `ToolMessage` the resume stream yields. - - Returns: - The agent plus one list per recorded `astream` keyword -- inputs, - contexts, and configs -- each appended to in call order. - """ - from langchain_core.messages import ToolMessage - - from deepagents_code.client.remote_client import RemoteAgent - - astream_inputs: list[Any] = [] - astream_contexts: list[object] = [] - astream_configs: list[Any] = [] - - class _Interrupt: - id = "interrupt-1" - value = { # noqa: RUF012 # test stub; immutability irrelevant - "action_requests": [ - {"name": "compact_conversation", "args": {"force": True}} - ] - } - - async def _astream(stream_input: object, **kwargs: object): # noqa: RUF029, ANN202 - astream_inputs.append(stream_input) - astream_contexts.append(kwargs.get("context")) - astream_configs.append(kwargs.get("config")) - if stream_input is None: - yield ((), "updates", {"__interrupt__": [_Interrupt()]}) - else: - yield ( - (), - "messages", - (ToolMessage(content=tool_content, tool_call_id="x"), {}), - ) - - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream - return agent, astream_inputs, astream_contexts, astream_configs - - 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, astream_configs = ( - 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} - - config = {"configurable": {"thread_id": "test-thread"}} - with patch.object(settings, "model_context_limit", 4096): - result = await app._drive_server_side_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 - - # Only the resume round is tagged, so LangSmith can tell the - # continuation apart from the run that opened the turn. - initial_config, resume_config = astream_configs - assert RESUME_TRACE_TAG not in initial_config.get("tags", []) - assert RESUME_TRACE_TAG in resume_config["tags"] - assert initial_config["configurable"] == resume_config["configurable"] - - 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_server_side_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_server_side_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) - 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_stream_failure_keeps_usage_recorded_once(self) -> None: - """Usage completed before a stream failure is still merged once.""" - from langchain_core.messages import AIMessage - - from deepagents_code.client.remote_client import RemoteAgent - - usage_message = AIMessage( - content="summary", - id="failed-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 - ): - 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_server_side_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_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 - - 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, _configs = 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_server_side_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, _configs = 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_server_side_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_server_side_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 + assert archive.exists() - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - app._agent = agent - app._lc_thread_id = "test-thread" + def test_archive_refreshed_between_iterdir_and_unlink_is_kept( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """An archive rewritten after the sweep lists it must not be deleted. - config = {"configurable": {"thread_id": "test-thread"}} - result = await app._drive_server_side_compaction(config) # ty: ignore - await pilot.pause() + Simulates a second `dcode` process refreshing an expired archive after + this process's `iterdir()` has already enumerated it: the pre-unlink + `fstat` observes the refreshed mtime and keeps the file, so the rewrite + is not orphaned by a stale expiry decision. + """ + archive_dir = self._setup(tmp_path, monkeypatch) + archive = archive_dir / "old.md" + archive.write_text("old") + old_time = time.time() - 31 * 86_400 + os.utime(archive, (old_time, old_time)) - 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", - ), - {}, - ), - ) + real_fstat = os.fstat + refreshed = False - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream + def fstat_with_refresh(fd: int) -> os.stat_result: + nonlocal refreshed + if not refreshed: + refreshed = True + # The racing writer rewrites the archive before the sweep's + # fstat lands, making it fresh again. + archive.write_text("refreshed") + fresh_time = time.time() + os.utime(archive, (fresh_time, fresh_time)) + return real_fstat(fd) - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - app._agent = agent - app._lc_thread_id = "test-thread" + monkeypatch.setattr(os, "fstat", fstat_with_refresh) - config = {"configurable": {"thread_id": "test-thread"}} - result = await app._drive_server_side_compaction(config) # ty: ignore - await pilot.pause() + assert sweep_offloaded_history() == 0 + assert archive.exists() + assert archive.read_text() == "refreshed" - 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. +class TestArtifactsRoot: + """Cover the real-filesystem artifacts root for offloaded tool results.""" - 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 + def test_artifacts_root_is_stable_and_hardened( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The per-user artifacts dir is predictable, private, and reused.""" + temp_dir = tmp_path / "tmp" + temp_dir.mkdir() + getuid = getattr(os, "getuid", None) + uid = getuid() if getuid is not None else os.getpid() - astream_inputs: list[Any] = [] + monkeypatch.setattr(tempfile, "gettempdir", lambda: str(temp_dir)) - class _Interrupt: - def __init__(self, iid: str) -> None: - self.id = iid - self.value = {"action_requests": [{"name": "write_file", "args": {}}]} + storage = _artifacts_root() + root_path = Path(storage.root) - 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}")]}) + assert storage.large_results_dir is None + assert root_path.samefile(temp_dir / f"dcode-artifacts-{uid}") + assert stat.S_IMODE(root_path.stat().st_mode) == 0o700 + # Stable across calls (paths embedded in resumed threads stay resolvable). + assert _artifacts_root() == storage - agent = MagicMock(spec=RemoteAgent) - agent.aensure_thread = AsyncMock() - agent.aupdate_state = AsyncMock() - agent.astream = _astream + def test_windows_artifacts_root_is_accepted_by_filesystem_tools(self) -> None: + """A Windows temp path retains its drive without a rejected drive prefix.""" + disk_root = PureWindowsPath( + "C:/Users/test/AppData/Local/Temp/dcode-artifacts-123" + ) - app = DeepAgentsApp() - async with app.run_test() as pilot: - await pilot.pause() - app._agent = agent - app._lc_thread_id = "test-thread" + root = _filesystem_tool_path(disk_root) + result_path = f"{root}/large_tool_results/tool-call-id" - config = {"configurable": {"thread_id": "test-thread"}} - result = await app._drive_server_side_compaction(config) # ty: ignore - await pilot.pause() + assert root == "//?/C:/Users/test/AppData/Local/Temp/dcode-artifacts-123" + assert PureWindowsPath(root).is_absolute() + assert validate_path(result_path) == result_path - # 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) - ) + def test_artifacts_root_falls_back_when_predictable_path_foreign_owned( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A predictable dir owned by another user is rejected for a unique one.""" + getuid = getattr(os, "getuid", None) + if getuid is None: + pytest.skip("uid ownership check requires os.getuid") + temp_dir = tmp_path / "tmp" + temp_dir.mkdir() + uid = getuid() + reserved = temp_dir / f"dcode-artifacts-{uid}" + reserved.mkdir() # a real, us-owned directory; lstat is faked below -class TestRemoveUnansweredOffloadSeed: - """Cleanup of a committed-but-unanswered `/offload` seed after a failure.""" + real_lstat = Path.lstat - @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 fake_lstat(self: Path) -> Any: # noqa: ANN401 + info = real_lstat(self) + if self == reserved: + return SimpleNamespace(st_mode=info.st_mode, st_uid=info.st_uid + 1) + return info - 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 + monkeypatch.setattr(tempfile, "gettempdir", lambda: str(temp_dir)) + monkeypatch.setattr(Path, "lstat", fake_lstat) - 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" - ) + storage = _artifacts_root() + next_storage = _artifacts_root() - 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 storage.root == "/dcode-artifacts-fallback" + assert next_storage.root == storage.root + assert storage.large_results_dir is not None + assert next_storage.large_results_dir is not None + assert not storage.large_results_dir.samefile(reserved) + assert storage.large_results_dir.name.startswith(f"dcode-artifacts-{uid}-") + assert stat.S_IMODE(storage.large_results_dir.stat().st_mode) == 0o700 + assert next_storage.large_results_dir != storage.large_results_dir - 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() +class TestOffloadRouting: + """Test that /offload is routed through _handle_command.""" - async def test_noop_when_seed_absent(self) -> None: - """Nothing is removed when no seed with the id is present.""" + async def test_offload_routed_from_handle_command(self) -> None: + """'/offload' should be correctly routed through _handle_command.""" 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() + app._agent = None + app._lc_thread_id = None - 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 app._handle_command("/offload") 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 + msgs = app.query(AppMessage) + assert any("Nothing to offload" in str(w._content) for w in msgs) - async def test_returns_false_when_state_read_fails(self) -> None: - """A failed state read cannot confirm cleanup, so it reports unclean.""" + async def test_compact_alias_routed_from_handle_command(self) -> None: + """'/compact' should still route through _handle_command for backward compat.""" 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() + app._agent = None + app._lc_thread_id = None - 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 app._handle_command("/compact") 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 + msgs = app.query(AppMessage) + assert any("Nothing to offload" in str(w._content) for w in msgs) class TestFormatTokenCount: @@ -3453,21 +1356,41 @@ 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.""" +class TestEventCutoff: + """`_event_cutoff` feeds the offloaded/kept counts, so it must not guess. + + A wrong cutoff shifts `messages_offloaded`/`messages_kept` and the + already-compacted short circuit, so every malformed shape has to read as + zero rather than as a plausible index. + """ - def test_summarization_cutoff_reads_int(self) -> None: - from deepagents_code.app import _summarization_cutoff + @pytest.mark.parametrize( + ("event", "expected"), + [ + ({"cutoff_index": 3}, 3), + ({"cutoff_index": 0}, 0), + (None, 0), + ("not-a-dict", 0), + ({}, 0), + ({"cutoff_index": None}, 0), + ({"cutoff_index": "3"}, 0), + ({"cutoff_index": 3.5}, 0), + # `bool` is an `int` subclass, so an unguarded isinstance check + # would read this as cutoff 1. + ({"cutoff_index": True}, 0), + ({"cutoff_index": False}, 0), + ], + ) + def test_only_a_real_int_cutoff_is_honoured( + self, event: object, expected: int + ) -> None: + from deepagents_code.offload_middleware import _event_cutoff - assert _summarization_cutoff({"cutoff_index": 4}) == 4 + assert _event_cutoff(event) == expected - 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 +class TestOffloadHelpers: + """Pure helpers for effective-conversation reconstruction.""" def test_effective_conversation_applies_event(self) -> None: from deepagents_code.app import _effective_conversation @@ -3494,32 +1417,6 @@ def test_effective_conversation_degrades_on_malformed(self) -> None: == messages ) - def test_offload_accounting_bounds_the_cutoff_like_the_window_does(self) -> None: - """The report and the window must read one cutoff the same way. - - `/offload` computes `messages_offloaded` from the cutoff and its token - counts from `_effective_conversation`, which bounds-checks internally. - Called without `message_count`, an out-of-bounds cutoff is trusted by the - first and rejected by the second, so the user is told a large offload - happened next to roughly zero token savings, with nothing logged. - """ - from deepagents_code.app import ( - _effective_conversation, - _summarization_cutoff, - ) - - messages = ["m0", "m1"] - event = {"summary_message": "S", "cutoff_index": 99} - - cutoff = _summarization_cutoff(event, message_count=len(messages)) - window = _effective_conversation(messages, event) - - # Both degrade: the cutoff to 0, the window to the full list. An - # unbounded read would give cutoff 99 against a 2-message window. - assert cutoff == 0 - assert window == messages - assert max(0, len(messages) - cutoff) == len(messages) - def test_effective_conversation_logs_a_discarded_event( self, caplog: pytest.LogCaptureFixture ) -> None: @@ -3587,39 +1484,285 @@ def test_effective_conversation_cutoff_past_end(self) -> None: # Not the SDK's reading, which would be `["S"]`. 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 _deny_dispatched_call( + reason: str | None, +) -> Callable[[Any, Any], dict[str, Any]]: + """Build an `aafter_model` stub that denies the dispatched compact call.""" + 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}} - def test_find_compaction_failure_scans_durable_state(self) -> None: - from langchain_core.messages import HumanMessage, ToolMessage + return deny + + +class TestOffloadOperation: + """The server service owns checkpoint state and compaction policy.""" + + @staticmethod + def _runtime() -> Runtime[CLIContextSchema]: + return Runtime(context=CLIContextSchema()) - from deepagents_code.app import _find_compaction_failure - from deepagents_code.offload_middleware import COMPACTION_FAILURE_PREFIX + @staticmethod + def _middleware( + *, hook_update: dict[str, object] | None = None + ) -> tuple[Any, MagicMock, MagicMock]: + from deepagents_code.offload_middleware import OffloadOperation + + compaction = MagicMock() + compaction._aplan_forced_compaction_update = AsyncMock() + compaction._summarization._apply_event_to_messages.side_effect = ( + lambda messages, _event: messages + ) + hooks = MagicMock() + # Default to the shape `ServerHooksMiddleware._after_model` really + # returns: every one of its return paths carries the pre-tool channel, + # including the "no hook events enabled" path. The operation fails closed + # when the channel is absent, so a mock returning a bare `{}` would + # assert a contract the middleware never produces. + from deepagents_code.hooks.server_middleware import _PRE_TOOL_STATE_KEY + + hooks.aafter_model = AsyncMock( + return_value=hook_update + if hook_update is not None + else {_PRE_TOOL_STATE_KEY: {}} + ) + return OffloadOperation(compaction, hooks), compaction, hooks - failing = ToolMessage( - content=f"{COMPACTION_FAILURE_PREFIX}: boom", - tool_call_id="tc", + @staticmethod + def _plan(update: dict[str, object]) -> SimpleNamespace: + """Build the narrow compaction-plan shape consumed by the operation.""" + return SimpleNamespace(update=lambda _path: update, archive=MagicMock()) + + async def test_compacts_checkpoint_state_without_message_input(self) -> None: + event = _summary_event(2) + middleware, compaction, _hooks = self._middleware() + compaction._aplan_forced_compaction_update = AsyncMock( + return_value=self._plan( + { + "_summarization_event": event, + "_summarization_session_id": "archive-1", + } + ) ) - messages = [HumanMessage("hi"), failing] - assert ( - _find_compaction_failure(messages) == f"{COMPACTION_FAILURE_PREFIX}: boom" + state = { + "messages": _make_dict_messages(4), + } + + execution = await middleware.execute(state, self._runtime()) + + compaction._aplan_forced_compaction_update.assert_awaited_once() + await_args = compaction._aplan_forced_compaction_update.await_args + assert await_args is not None + state_arg = await_args.args[0] + assert state_arg is state + assert "messages" not in execution.update + assert execution.update["_summarization_session_id"] == "archive-1" + assert execution.result["status"] == "compacted" + assert execution.result["messages_offloaded"] == 2 + + async def test_a_hook_interrupt_propagates_instead_of_failing(self) -> None: + """A hook request must reach the client, not become a `failed` result. + + Two independent mechanisms protect this: the `BaseException` base, which + the compaction chain's broad `except Exception` handlers cannot catch, + and the explicit re-raise in `execute`. Either alone is sufficient, so + this asserts the outcome rather than a mechanism -- losing *both* turns + every interrupt into "Compaction failed: + HookTransportInterruptError", silently breaking `/offload` for hook + users only. Verified by mutating both. The boundary test mocks the whole + operation, so it cannot cover this. + """ + from uuid import uuid4 + + from deepagents_code.hooks.server_middleware import ( + HookTransportInterruptError, + ) + + middleware, compaction, _hooks = self._middleware() + request = SimpleNamespace(invocation_id=uuid4()) + compaction._aplan_forced_compaction_update = AsyncMock( + side_effect=HookTransportInterruptError(cast("Any", request)) + ) + + with pytest.raises(HookTransportInterruptError) as raised: + await middleware.execute( + {"messages": _make_dict_messages(4)}, self._runtime() + ) + + assert raised.value.request is request + + async def test_reoffload_reports_the_absolute_cutoff_delta(self) -> None: + """Counts are deltas against the prior event, not absolute cutoffs. + + With a prior cutoff of 0 the two are indistinguishable, so this drives a + chained offload: 6 messages, prior cutoff 2, new cutoff 5 must report 3 + offloaded and 1 kept. Reporting `new_cutoff` directly would say 5. + """ + middleware, compaction, _hooks = self._middleware() + compaction._aplan_forced_compaction_update = AsyncMock( + return_value=self._plan( + { + "_summarization_event": _summary_event(5), + "_summarization_session_id": "archive-1", + } + ) + ) + state = { + "messages": _make_dict_messages(6), + "_summarization_event": _summary_event(2), + } + + execution = await middleware.execute(state, self._runtime()) + + assert execution.result["status"] == "compacted" + assert execution.result["messages_offloaded"] == 3 + assert execution.result["messages_kept"] == 1 + + async def test_non_compacted_counts_never_go_negative(self) -> None: + """A stale cutoff beyond the message count must not report a negative.""" + middleware, compaction, _hooks = self._middleware() + compaction._aplan_forced_compaction_update = AsyncMock(return_value=None) + state = { + "messages": _make_dict_messages(2), + "_summarization_event": _summary_event(9), + } + + execution = await middleware.execute(state, self._runtime()) + + assert execution.result["status"] == "noop" + assert execution.result["messages_kept"] == 0 + assert execution.update == {} + + async def test_hook_denial_skips_compaction(self) -> None: + """A `PreToolUse` denial must stop the compaction. + + Keys the outcome on the id the node really generated and asserts the + dispatched call's `name`/`args`, so a re-spelled tool name or a dropped + `force` flag -- either of which silently exempts `/offload` from the + user's hook -- fails here instead of reading as "no outcome". + """ + middleware, compaction, hooks = self._middleware() + hooks.aafter_model = AsyncMock(side_effect=_deny_dispatched_call("policy")) + + execution = await middleware.execute( + {"messages": _make_dict_messages(4)}, self._runtime() + ) + + assert execution.result["status"] == "denied" + assert execution.result["error"] == "policy" + compaction._aplan_forced_compaction_update.assert_not_awaited() + + async def test_hook_denial_without_a_reason_still_stops_compaction(self) -> None: + """A denial carrying no reason must not read as an allow.""" + middleware, compaction, hooks = self._middleware() + hooks.aafter_model = AsyncMock(side_effect=_deny_dispatched_call(None)) + + execution = await middleware.execute( + {"messages": _make_dict_messages(4)}, self._runtime() ) - def test_find_compaction_failure_ignores_success(self) -> None: - from langchain_core.messages import ToolMessage + assert execution.result["status"] == "denied" + compaction._aplan_forced_compaction_update.assert_not_awaited() + + async def test_a_missing_hook_channel_refuses_instead_of_allowing(self) -> None: + """A hook decision that cannot be read must not be treated as an allow. + + Every `_after_model` return path carries the pre-tool channel, so its + absence means the channel, the id derivation, or the outcome shape + drifted. Reading that through a `.get(..., {})` chain would turn a user's + denial into "no outcome" and compact straight through it, with no log. + """ + middleware, compaction, hooks = self._middleware(hook_update={}) + + execution = await middleware.execute( + {"messages": _make_dict_messages(4)}, self._runtime() + ) + + assert execution.result["status"] == "failed" + assert "hook decision" in (execution.result["error"] or "") + compaction._aplan_forced_compaction_update.assert_not_awaited() + assert "messages" not in execution.update + hooks.aafter_model.assert_awaited_once() + + async def test_failure_returns_result_without_rewriting_messages(self) -> None: + middleware, compaction, _hooks = self._middleware() + compaction._aplan_forced_compaction_update = AsyncMock( + side_effect=OSError("archive unavailable") + ) + + execution = await middleware.execute( + {"messages": _make_dict_messages(4)}, self._runtime() + ) + + assert execution.result["status"] == "failed" + assert "archive unavailable" in (execution.result["error"] or "") + assert "messages" not in execution.update + + +class TestForcedOffloadCallId: + """The hook dispatch's call id must be stable across a run's resumes.""" + + def test_missing_checkpoint_namespace_is_logged_not_silent( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A run without a usable `checkpoint_ns` breaks hook resumes. + + 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() + + assert call_id.startswith("offload-precompact-") + assert "checkpoint_ns" in caplog.text + + 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. + + 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 + + with ( + patch.object( + offload_middleware, "get_config", side_effect=RuntimeError("no context") + ), + caplog.at_level("WARNING"), + ): + call_id = offload_middleware._forced_offload_call_id() + + assert call_id.startswith("offload-precompact-") + assert "checkpoint_ns" not in caplog.text + + 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 - from deepagents_code.app import _find_compaction_failure + 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() - 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 + assert first == second diff --git a/libs/code/tests/unit_tests/test_offload_api.py b/libs/code/tests/unit_tests/test_offload_api.py new file mode 100644 index 0000000000..dc26283251 --- /dev/null +++ b/libs/code/tests/unit_tests/test_offload_api.py @@ -0,0 +1,1703 @@ +"""Tests for the server-owned offload HTTP boundary.""" + +from __future__ import annotations + +import asyncio +import contextlib +from types import SimpleNamespace +from typing import TYPE_CHECKING, cast +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from deepagents_code.offload_middleware import OffloadExecution, OffloadResult + +if TYPE_CHECKING: + from collections.abc import Iterator + + from deepagents_code.offload_middleware import _PendingArchive + + +@pytest.fixture(autouse=True) +def _reset_offload_globals() -> Iterator[None]: + """Clear cached clients and operation state between tests. + + `offload_api` caches one `httpx`-backed client per process (building one per + request leaks a connection pool). Tests patch `get_client`, so the cache has + to be dropped or the first test's mock would serve every later one. + """ + from deepagents_code import offload_api + + offload_api._client = None + offload_api._active_operations.clear() + offload_api._operation_outcomes.clear() + try: + yield + finally: + offload_api._client = None + offload_api._active_operations.clear() + offload_api._operation_outcomes.clear() + + +class TestOperationPayload: + """Malformed client requests fail with a field-naming 422 at the boundary.""" + + def test_valid_payload_passes_with_unknown_keys(self) -> None: + from deepagents_code.offload_api import _operation_payload + + operation_id, context, responses = _operation_payload( + { + "operation_id": "op-1", + "context": { + "model": "openai:gpt-5", + "model_params": {"temperature": 0.2}, + "profile_overrides": {"max_input_tokens": 1000}, + "model_context_limit": 32000, + "auto_approve": True, + "hooks_server_events": ["PreCompact"], + "thread_id": "thread-1", + "some_future_field": {"ignored": True}, + }, + } + ) + + assert operation_id == "op-1" + assert context["model"] == "openai:gpt-5" + assert context["some_future_field"] == {"ignored": True} + assert responses == {} + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("model", 123), + ("classifier_model", ["openai:gpt-5"]), + ("approval_mode", 1), + ("thread_id", {"id": "t"}), + ("hooks_snapshot_id", 0.5), + ("prompt_id", True), + ("model_params", "temperature=0.2"), + ("profile_overrides", [("max_input_tokens", 1000)]), + ("model_context_limit", "32000"), + ("model_context_limit", True), + ("auto_approve", "yes"), + ("hooks_server_events", "PreCompact"), + ("hooks_server_events", ["PreCompact", 42]), + ], + ) + def test_bad_context_field_names_the_field(self, field: str, value: object) -> None: + from deepagents_code.offload_api import _operation_payload + + with pytest.raises(TypeError, match=f"context.{field}"): + _operation_payload( + { + "operation_id": "op-1", + "context": {field: value}, + } + ) + + def test_null_context_fields_pass(self) -> None: + from deepagents_code.offload_api import _operation_payload + + _, context, _ = _operation_payload( + { + "operation_id": "op-1", + "context": { + "model": None, + "model_params": None, + "model_context_limit": None, + "auto_approve": None, + "hooks_server_events": None, + }, + } + ) + assert context["model"] is None + + @pytest.mark.parametrize( + "key", + [ + "base_url", + "api_base", + "openai_api_base", + "anthropic_api_url", + "azure_endpoint", + "azure_openai_api_base", + "api_endpoint", + "openai_proxy", + "anthropic_proxy", + "proxy", + "proxies", + "http_client", + "http_async_client", + "transport", + "default_headers", + "custom_headers", + ], + ) + def test_transport_model_params_are_stripped(self, key: str) -> None: + """The boundary drops endpoint/transport keys from `model_params`. + + These keys would route the server's credentialed provider calls to a + client-chosen destination. + """ + from deepagents_code.offload_api import _operation_payload + + _, context, _ = _operation_payload( + { + "operation_id": "op-1", + "context": { + "model": "openai:gpt-5", + "model_params": { + key: "http://attacker.example/", + "temperature": 0.2, + }, + }, + } + ) + + assert context["model_params"] == {"temperature": 0.2} + + def test_stripping_is_logged_with_key_names_only( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A dropped transport key must leave a trace naming the key. + + Silently ignoring it leaves a user whose gateway config is being skipped + with nothing to find. The value is an endpoint or header, so only the + key name is logged. + """ + import logging + + from deepagents_code.offload_api import _operation_payload + + with caplog.at_level(logging.WARNING): + _operation_payload( + { + "operation_id": "op-1", + "context": { + "model_params": { + "base_url": "http://gateway.internal/v1", + "temperature": 0.2, + } + }, + } + ) + + assert "base_url" in caplog.text + assert "gateway.internal" not in caplog.text + + def test_clean_model_params_dict_is_untouched(self) -> None: + from deepagents_code.offload_api import _operation_payload + + _, context, _ = _operation_payload( + { + "operation_id": "op-1", + "context": {"model_params": {"temperature": 0.2, "max_tokens": 64}}, + } + ) + + assert context["model_params"] == {"temperature": 0.2, "max_tokens": 64} + + +def _thread_state(checkpoint_id: str = "checkpoint-1") -> dict[str, object]: + """Build an idle LangGraph thread-state response.""" + return { + "values": { + "messages": [ + { + "role": "user", + "content": "hello", + "id": "message-1", + } + ], + "_session_cost_usd": 1.0, + "_model_spec": "provider:checkpointed-model", + "_model_params": { + "base_url": "https://trusted.example/v1", + "temperature": 0.1, + }, + }, + "next": [], + "tasks": [], + "interrupts": [], + "checkpoint": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": checkpoint_id, + "checkpoint_map": {}, + }, + } + + +def _result( + archive_path: str | None = "/conversation_history/thread-1.md", +) -> OffloadResult: + """Build a complete operation result.""" + return { + "status": "compacted", + "messages_offloaded": 1, + "messages_kept": 1, + "tokens_before": 20, + "tokens_after": 10, + "archive_path": archive_path, + "archive_ephemeral": False, + "error": None, + } + + +class TestExecuteOffload: + """The route owns state hydration, validation, and atomic persistence.""" + + def test_hydrates_persisted_summary_message(self) -> None: + """A subsequent offload receives a message object in its prior event.""" + from langchain_core.messages import HumanMessage + + from deepagents_code.offload_api import _hydrate_state + + state = _hydrate_state( + { + "messages": [ + {"role": "user", "content": "new message", "id": "message-1"} + ], + "_summarization_event": { + "cutoff_index": 1, + "summary_message": { + "type": "human", + "content": "Prior summary.", + "id": "summary-1", + }, + }, + } + ) + + event = state["_summarization_event"] + assert isinstance(event, dict) + assert isinstance(event["summary_message"], HumanMessage) + assert event["summary_message"].content == "Prior summary." + + async def test_commits_event_and_cost_without_messages(self) -> None: + from deepagents_code import offload_api + + before = _thread_state() + calls: list[str] = [] + + async def update_state( # noqa: RUF029 # AsyncMock side effect contract + *_args: object, **_kwargs: object + ) -> None: + calls.append("checkpoint") + + append = SimpleNamespace( + path="/conversation_history/thread-1.md", rollback=AsyncMock() + ) + archive = SimpleNamespace( + session_id="archive-1", + write=AsyncMock(side_effect=lambda: calls.append("archive") or append), + update=lambda path: { + "_summarization_event": {"cutoff_index": 1, "file_path": path} + }, + ) + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before]), + update_state=AsyncMock(side_effect=update_state), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + { + "_summarization_event": { + "cutoff_index": 1, + "file_path": None, + }, + "_summarization_session_id": "archive-1", + }, + _result(archive_path=None), + cast("_PendingArchive", archive), + ) + ) + ) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + ) + + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object( + offload_api, + "get_server_runtime", + new=AsyncMock( + return_value=SimpleNamespace( + agent=SimpleNamespace(store=None), offload=operation + ) + ), + ), + patch.object(offload_api, "prepare_operation_cost", return_value=prepared), + ): + response = await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={"model": "test:model"}, + hook_responses={}, + ) + + assert response == {"status": "complete", "result": _result()} + state = operation.execute.await_args.args[0] + assert state["messages"][0].id == "message-1" + runtime = operation.execute.await_args.args[1] + assert runtime.context["thread_id"] == "thread-1" + assert runtime.context["model"] == "provider:checkpointed-model" + assert runtime.context["model_params"] == { + "base_url": "https://trusted.example/v1", + "temperature": 0.1, + } + assert threads.update_state.await_count == 2 + args = threads.update_state.await_args_list[0] + assert args.args[:2] == ( + "thread-1", + { + "_summarization_event": { + "cutoff_index": 1, + "file_path": None, + }, + "_summarization_session_id": "archive-1", + "_session_cost_usd": 0.25, + }, + ) + assert "messages" not in args.args[1] + assert "checkpoint" not in args.kwargs + assert threads.update_state.await_args_list[1].args == ( + "thread-1", + { + "_summarization_event": { + "cutoff_index": 1, + "file_path": "/conversation_history/thread-1.md", + } + }, + ) + assert calls == ["checkpoint", "archive", "checkpoint"] + prepared.rollback.assert_not_called() + + async def test_failed_archive_link_restores_the_append(self) -> None: + """A failed follow-up checkpoint cannot leave duplicate history.""" + from deepagents_code import offload_api + + before = _thread_state() + unlinked = _thread_state("reserved") + unlinked_values = cast("dict[str, object]", unlinked["values"]) + unlinked_values["_summarization_event"] = { + "cutoff_index": 1, + "file_path": None, + } + append = SimpleNamespace( + path="/conversation_history/thread-1.md", rollback=AsyncMock() + ) + archive = SimpleNamespace( + session_id="archive-1", + write=AsyncMock(return_value=append), + update=lambda path: { + "_summarization_event": {"cutoff_index": 1, "file_path": path} + }, + ) + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before, unlinked]), + update_state=AsyncMock( + side_effect=[None, RuntimeError("archive link unavailable")] + ), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + { + "_summarization_event": { + "cutoff_index": 1, + "file_path": None, + }, + "_summarization_session_id": "archive-1", + }, + _result(archive_path=None), + cast("_PendingArchive", archive), + ) + ) + ) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + records=[], + ) + + with self._patched(offload_api, threads, operation, prepared): + response = await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + assert response["status"] == "complete" + assert response["result"]["archive_path"] is None + append.rollback.assert_awaited_once() + prepared.rollback.assert_not_called() + + async def test_cancellation_waits_for_checkpoint_archive_settlement(self) -> None: + """A reserved commit must settle before cancellation becomes terminal.""" + from deepagents_code import offload_api + + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before]), + update_state=AsyncMock(), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {"_summarization_event": {"cutoff_index": 1}}, + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + ) + settlement_started = asyncio.Event() + finish_settlement = asyncio.Event() + + async def settle(*_args: object, **_kwargs: object) -> None: + settlement_started.set() + await finish_settlement.wait() + + with ( + self._patched(offload_api, threads, operation, prepared), + patch.object( + offload_api, + "_commit_deferred_archive", + new=AsyncMock(side_effect=settle), + ) as commit, + ): + task = asyncio.create_task( + offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + ) + await asyncio.wait_for(settlement_started.wait(), timeout=1) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + finish_settlement.set() + with pytest.raises(asyncio.CancelledError): + await task + + commit.assert_awaited_once() + + async def test_request_transport_cannot_replace_checkpointed_model(self) -> None: + """Offload uses the model settings from the target thread checkpoint.""" + from deepagents_code import offload_api + + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before]), + update_state=AsyncMock(), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {}, + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace(update={}, rollback=MagicMock(), commit=MagicMock()) + + with self._patched(offload_api, threads, operation, prepared): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={ + "model": "attacker:model", + "model_params": {"base_url": "https://attacker.example"}, + }, + hook_responses={}, + ) + + runtime = operation.execute.await_args.args[1] + assert runtime.context["model"] == "provider:checkpointed-model" + assert runtime.context["model_params"]["base_url"] == ( + "https://trusted.example/v1" + ) + + async def test_legacy_thread_reuses_startup_summarizer(self) -> None: + """A thread without model metadata ignores request model selection.""" + from deepagents_code import offload_api + + before = _thread_state() + values = cast("dict[str, object]", before["values"]) + assert isinstance(values, dict) + values.pop("_model_spec") + values.pop("_model_params") + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before]), + update_state=AsyncMock(), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {}, + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace(update={}, rollback=MagicMock(), commit=MagicMock()) + + with self._patched(offload_api, threads, operation, prepared): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={"model": "request:model", "model_params": {"x": 1}}, + hook_responses={}, + ) + + runtime = operation.execute.await_args.args[1] + assert "model" not in runtime.context + assert "model_params" not in runtime.context + + async def test_checkpoint_change_fails_without_state_commit(self) -> None: + from deepagents_code import offload_api + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock( + side_effect=[_thread_state("before"), _thread_state("changed")] + ), + update_state=AsyncMock(), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {"_summarization_event": {"cutoff_index": 1}}, + _result(), # ty: ignore[invalid-argument-type] + cast( + "_PendingArchive", + SimpleNamespace(session_id="archive-1", write=AsyncMock()), + ), + ) + ) + ) + + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object( + offload_api, + "get_server_runtime", + new=AsyncMock( + return_value=SimpleNamespace( + agent=SimpleNamespace(store=None), offload=operation + ) + ), + ), + patch.object(offload_api, "prepare_operation_cost") as prepare, + pytest.raises(offload_api._OffloadConflictError, match="thread changed"), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + threads.update_state.assert_not_awaited() + operation.execute.return_value.archive.write.assert_not_awaited() + prepare.assert_not_called() + + @pytest.mark.parametrize("status", ["busy", "interrupted"]) + async def test_thread_with_work_in_flight_is_rejected_before_operation( + self, status: str + ) -> None: + from deepagents_code import offload_api + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": status}), + get_state=AsyncMock(), + update_state=AsyncMock(), + ) + runtime = AsyncMock() + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object(offload_api, "get_server_runtime", new=runtime), + pytest.raises(offload_api._OffloadConflictError, match="active"), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + threads.get_state.assert_not_awaited() + runtime.assert_not_awaited() + + async def test_errored_thread_is_still_offloadable(self) -> None: + """A failed turn must not lock the user out of `/offload`. + + A run that raises leaves the thread row on `error` until the next run + completes, which is exactly when a user reaches for `/offload` to + recover. Reaching the state read proves the status gate let it past; + in-flight work is caught separately by the `next`/`tasks`/`interrupts` + check against the checkpoint. + """ + from deepagents_code import offload_api + + class _ReachedStateReadError(Exception): + """Sentinel proving control passed the thread-status gate.""" + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "error"}), + get_state=AsyncMock(side_effect=_ReachedStateReadError), + update_state=AsyncMock(), + ) + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object(offload_api, "get_server_runtime", new=AsyncMock()), + pytest.raises(_ReachedStateReadError), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + threads.update_state.assert_not_awaited() + + @staticmethod + @contextlib.contextmanager + def _patched( + offload_api: object, + threads: SimpleNamespace, + operation: SimpleNamespace, + prepared: object, + ) -> Iterator[None]: + """Patch the client, runtime, and cost seams `_execute_offload` uses.""" + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object( + offload_api, + "get_server_runtime", + new=AsyncMock( + return_value=SimpleNamespace( + agent=SimpleNamespace(store=None), offload=operation + ) + ), + ), + patch.object(offload_api, "prepare_operation_cost", return_value=prepared), + ): + yield + + @pytest.mark.parametrize("channel", ["messages", "todos"]) + async def test_unpermitted_update_is_refused_and_cost_rolled_back( + self, channel: str + ) -> None: + """A channel outside `OffloadStateUpdate` cannot reach the checkpoint. + + Unlike asserting on a mocked update that never contains the channel + (which passes whether the guard exists or not), this drives an update + that actually carries one. `todos` covers the allowlist itself: a + `messages`-only check would let every other channel through. + """ + from deepagents_code import offload_api + + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before]), + update_state=AsyncMock(), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + # Deliberately violates `OffloadStateUpdate` -- that is the + # point of the test: the runtime guard is the backstop for + # `Any`-typed values the SDK hands back. + {channel: ["smuggled"]}, # ty: ignore[invalid-key,invalid-argument-type] + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace(update={}, rollback=MagicMock(), commit=MagicMock()) + + with ( + self._patched(offload_api, threads, operation, prepared), + pytest.raises(RuntimeError, match=f"may not write .*{channel}"), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + threads.update_state.assert_not_awaited() + prepared.rollback.assert_called_once() + + async def test_empty_update_still_releases_claimed_cost_records(self) -> None: + """A prepare with nothing to write must not silently eat its records.""" + from deepagents_code import offload_api + + before = _thread_state() + noop_result = {**_result(), "status": "noop"} + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before]), + update_state=AsyncMock(), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution({}, noop_result) # ty: ignore[invalid-argument-type] + ) + ) + prepared = SimpleNamespace(update={}, rollback=MagicMock(), commit=MagicMock()) + + with self._patched(offload_api, threads, operation, prepared): + response = await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + assert response == {"status": "complete", "result": noop_result} + threads.update_state.assert_not_awaited() + prepared.rollback.assert_called_once() + + async def test_write_failure_without_advance_restores_cost(self) -> None: + """A write that provably did not land returns its records to the recorder.""" + from deepagents_code import offload_api + + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + # Third read is `_write_landed`: same checkpoint => did not land. + get_state=AsyncMock(side_effect=[before, before, before]), + update_state=AsyncMock(side_effect=RuntimeError("boom")), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {"_summarization_event": {"cutoff_index": 1}}, + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + records=[], + ) + + with ( + self._patched(offload_api, threads, operation, prepared), + pytest.raises(RuntimeError, match="boom"), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + prepared.rollback.assert_called_once() + + async def test_write_failure_after_advance_keeps_cost_claimed(self) -> None: + """An indeterminate write must not re-queue records and double-charge.""" + from deepagents_code import offload_api + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock( + side_effect=[ + _thread_state("before"), + _thread_state("before"), + # `_write_landed`: the thread advanced, so the write likely + # applied despite the transport error. + _thread_state("after"), + ] + ), + update_state=AsyncMock(side_effect=RuntimeError("connection reset")), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {"_summarization_event": {"cutoff_index": 1}}, + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + records=[], + ) + + with ( + self._patched(offload_api, threads, operation, prepared), + pytest.raises( + offload_api._OffloadIndeterminateError, match="could not confirm" + ), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + prepared.rollback.assert_not_called() + + async def test_unreadable_thread_does_not_claim_the_thread_advanced( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """An unreadable readback must not be logged as an observed advance. + + Both outcomes keep the records claimed, but only `advanced` has evidence + the write landed. Reporting a thread advance that was never observed + would tell anyone auditing a missing charge that the spend was + accounted for. + """ + import logging + + from deepagents_code import offload_api + + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock( + side_effect=[before, before, RuntimeError("thread store offline")] + ), + update_state=AsyncMock(side_effect=RuntimeError("connection reset")), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {"_summarization_event": {"cutoff_index": 1}}, + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + records=[], + ) + + with ( + self._patched(offload_api, threads, operation, prepared), + caplog.at_level(logging.ERROR), + pytest.raises(offload_api._OffloadIndeterminateError), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + assert "could not be read back" in caplog.text + assert "may be lost from the thread total" in caplog.text + assert "advanced past checkpoint" not in caplog.text + prepared.rollback.assert_not_called() + prepared.commit.assert_called_once() + + async def test_cancelled_probe_still_settles_the_cost_records(self) -> None: + """A cancel inside the write-landed probe must not skip settlement. + + `prepare_operation_cost` drains the recorder destructively, so a prepare + that is neither committed nor rolled back deletes that spend from the + thread's lifetime total permanently. The probe runs inside the + settlement handler, so an escape from it -- a disconnect or a shutdown + re-delivering cancellation while the handler unwinds -- would take both + branches off the table and lose the records with no trace. + """ + from deepagents_code import offload_api + + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before, asyncio.CancelledError()]), + update_state=AsyncMock(side_effect=RuntimeError("boom")), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {"_summarization_event": {"cutoff_index": 1}}, + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + records=[], + ) + + with ( + self._patched(offload_api, threads, operation, prepared), + pytest.raises(offload_api._OffloadIndeterminateError), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + # Unreadable means indeterminate, so the records stay claimed rather + # than being restored -- but a decision was reached either way. + prepared.rollback.assert_not_called() + + async def test_a_cancelled_write_is_not_converted_to_a_runtime_error( + self, + ) -> None: + """Cancellation must propagate so the task actually observes it.""" + from deepagents_code import offload_api + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock( + side_effect=[ + _thread_state("before"), + _thread_state("before"), + _thread_state("after"), + ] + ), + update_state=AsyncMock(side_effect=asyncio.CancelledError()), + ) + operation = SimpleNamespace( + execute=AsyncMock( + return_value=OffloadExecution( + {"_summarization_event": {"cutoff_index": 1}}, + _result(), # ty: ignore[invalid-argument-type] + ) + ) + ) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + records=[], + ) + + with ( + self._patched(offload_api, threads, operation, prepared), + pytest.raises(asyncio.CancelledError), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + prepared.rollback.assert_not_called() + + @staticmethod + def _hook_request() -> object: + """Build a real server-owned hook invocation request.""" + from datetime import UTC, datetime + from pathlib import Path + from uuid import uuid4 + + from deepagents_code.hooks.models.domain import ( + ApprovalMode, + HookContext, + HookEvent, + PreToolUseEvent, + ToolCallData, + ) + from deepagents_code.hooks.models.transport import ( + HookInvocation, + HookInvocationRequest, + ) + + return HookInvocationRequest( + protocol_version=1, + invocation_id=uuid4(), + snapshot_id="snapshot-1", + run_id="run-1", + invocation=HookInvocation( + context=HookContext( + thread_id="thread-1", + cwd=Path("/tmp"), + approval_mode=ApprovalMode.MANUAL, + ), + event=PreToolUseEvent( + event=HookEvent.PRE_TOOL_USE, + call=ToolCallData( + id="call-1", name="compact_conversation", args={"force": True} + ), + ), + ), + deadline=datetime(2026, 7, 23, tzinfo=UTC), + ) + + async def test_a_hook_request_becomes_an_interrupt_response(self) -> None: + """An unanswered hook must leave the route as a resumable interrupt. + + `HookTransportInterruptError` derives from `BaseException` so the + compaction chain cannot swallow it -- which also means the route's own + `except Exception` cannot catch it. Without the dedicated handler it + escapes to Starlette as a raw 500 for every user with a `PreCompact` or + `PreToolUse` hook configured, and nothing else in the suite notices. + """ + from deepagents_code import offload_api + from deepagents_code.hooks.interrupt import is_hook_interrupt_payload + from deepagents_code.hooks.server_middleware import ( + HookTransportInterruptError, + ) + + request = self._hook_request() + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(return_value=before), + update_state=AsyncMock(), + ) + operation = SimpleNamespace( + execute=AsyncMock(side_effect=HookTransportInterruptError(request)) # ty: ignore[invalid-argument-type] + ) + prepared = SimpleNamespace( + update={}, rollback=MagicMock(), commit=MagicMock(), records=[] + ) + + with self._patched(offload_api, threads, operation, prepared): + response = await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + assert response["status"] == "interrupt" + assert is_hook_interrupt_payload(response["request"]) + assert response["request"]["request"]["invocation_id"] == str( + request.invocation_id # ty: ignore[unresolved-attribute] + ) + # Nothing may be committed while a hook is still unanswered. + threads.update_state.assert_not_awaited() + + async def test_accumulated_hook_responses_reach_the_operation(self) -> None: + """The resume round must hand the replies back to the hook transport. + + `operation_hook_responses` is the single line that makes a multi-round + resume terminate: `_invoke_hook` replays an already-answered invocation + from that mapping instead of raising again. Passing an empty mapping + would re-raise the same invocation forever and the client would die at + its round limit, so assert the mapping is actually installed. + """ + from deepagents_code import offload_api + from deepagents_code.hooks.server_middleware import operation_hook_responses + + seen: list[object] = [] + + async def execute( # noqa: RUF029 -- must satisfy the async execute signature + *_args: object, **_kwargs: object + ) -> OffloadExecution: + # Read the context var the way the hook transport does. + from deepagents_code.hooks import server_middleware + + seen.append(server_middleware._HOOK_RESPONSES.get()) + return OffloadExecution( + {"_summarization_event": {"cutoff_index": 1}}, + _result(), # ty: ignore[invalid-argument-type] + ) + + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(side_effect=[before, before]), + update_state=AsyncMock(), + ) + operation = SimpleNamespace(execute=AsyncMock(side_effect=execute)) + prepared = SimpleNamespace( + update={"_session_cost_usd": 0.25}, + rollback=MagicMock(), + commit=MagicMock(), + delta_usd=0.25, + records=[], + ) + + replies: dict[str, object] = {"hook-1": {"decision": "allow"}} + with self._patched(offload_api, threads, operation, prepared): + response = await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses=replies, + ) + + assert response["status"] == "complete" + assert seen == [replies] + # Outside the operation the var is back to graph mode. + assert operation_hook_responses is not None + + async def test_pending_graph_work_is_rejected(self) -> None: + from deepagents_code import offload_api + + pending = {**_thread_state(), "next": ["tools"]} + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(return_value=pending), + update_state=AsyncMock(), + ) + runtime = AsyncMock() + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object(offload_api, "get_server_runtime", new=runtime), + pytest.raises( + offload_api._OffloadConflictError, match="pending graph work" + ), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + runtime.assert_not_awaited() + + async def test_unregistered_thread_is_rejected_with_an_actionable_conflict( + self, + ) -> None: + """A 404 from the live thread store must not become an opaque 500. + + The dev server keeps checkpoint persistence and thread registration + separate, so a resumed thread can 404 here while holding state on disk. + `NotFoundError` is an ordinary `Exception`, so without this mapping it + reaches the route's generic handler and the user is told to read the + server log. + """ + import httpx + from langgraph_sdk.errors import NotFoundError + + from deepagents_code import offload_api + + request = httpx.Request("GET", "http://localhost/threads/thread-1") + not_found = NotFoundError( + "missing", response=httpx.Response(404, request=request), body=None + ) + threads = SimpleNamespace( + get=AsyncMock(side_effect=not_found), + get_state=AsyncMock(), + update_state=AsyncMock(), + ) + runtime = AsyncMock() + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object(offload_api, "get_server_runtime", new=runtime), + pytest.raises( + offload_api._OffloadConflictError, match="not registered on the server" + ), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + threads.get_state.assert_not_awaited() + threads.update_state.assert_not_awaited() + runtime.assert_not_awaited() + + async def test_empty_thread_reports_nothing_to_offload(self) -> None: + """An empty thread is an unchanged outcome, not a failure. + + `_checkpoint_id` rejects a thread with no checkpoint, so answering the + empty case at the boundary is what keeps `OffloadOperation.execute`'s + graceful `empty` branch reachable over HTTP. Without it the user is told + the operation failed for a thread that simply has nothing to compact. + """ + from deepagents_code import offload_api + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(return_value={"values": {}, "checkpoint": {}}), + update_state=AsyncMock(), + ) + runtime = AsyncMock() + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object(offload_api, "get_server_runtime", new=runtime), + ): + response = await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + assert response["status"] == "complete" + assert response["result"]["status"] == "empty" + assert response["result"]["messages_kept"] == 0 + # Nothing is compacted, so nothing is written and no agent is built. + threads.update_state.assert_not_awaited() + runtime.assert_not_awaited() + + async def test_missing_checkpoint_with_messages_is_rejected(self) -> None: + """State with messages but no checkpoint cannot be written against.""" + from deepagents_code import offload_api + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock( + return_value={**_thread_state(), "checkpoint": {}}, + ), + update_state=AsyncMock(), + ) + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object(offload_api, "get_server_runtime", new=AsyncMock()), + pytest.raises(offload_api._OffloadConflictError, match="no checkpoint"), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + +def test_validated_context_fields_exist_on_the_schema() -> None: + """The validator's field lists must not drift from `CLIContextSchema`. + + The names are hand-written string tuples, so a rename in the dataclass would + leave this route validating a key nobody sends -- forever, with no test + failing. The protocol version is pinned the same way; this closes the other + hand-maintained list. + """ + from dataclasses import fields + + from deepagents_code import offload_api + from deepagents_code._cli_context import CLIContextSchema + + declared = {f.name for f in fields(CLIContextSchema)} + validated = { + *offload_api._CONTEXT_STR_OR_NONE_FIELDS, + *offload_api._CONTEXT_DICT_FIELDS, + } + + assert validated <= declared, validated - declared + + +class TestRouteRegistration: + """The Starlette app exposes the paths and methods the client calls. + + Every other route test fabricates a request with hand-written + `path_params`, so a path or converter rename -- `{thread_id:str}` to + `{tid:str}`, say -- would leave the whole unit suite green while the real + handler raised `KeyError` (neither `TypeError` nor `ValueError`, so it + escapes the 422 block as a bare 500). + """ + + def test_offload_and_cancel_paths_are_registered(self) -> None: + from starlette.testclient import TestClient + + from deepagents_code import offload_api + from deepagents_code.offload_middleware import unchanged_offload_result + + calls: list[tuple[str, str]] = [] + + async def fake_execute( # noqa: RUF029 # replaces an async callee + thread_id: str, + *, + operation_id: str, + context: dict[str, object], # noqa: ARG001 + hook_responses: dict[str, object], # noqa: ARG001 + ) -> dict[str, object]: + calls.append((thread_id, operation_id)) + return { + "status": "complete", + "result": unchanged_offload_result("noop", messages=1, tokens=5), + } + + with ( + patch.object(offload_api, "_execute_offload", new=fake_execute), + TestClient(offload_api.app) as client, + ): + response = client.post( + "/dcode/threads/thread-42/offload", + json={ + "operation_id": "op-1", + "context": {}, + "hook_responses": {}, + }, + ) + + assert response.status_code == 200, response.text + # The handler read the id out of the real path params, so the route's + # converter name and the key it indexes agree. + assert calls == [("thread-42", "op-1")] + + def test_get_on_the_offload_path_is_not_allowed(self) -> None: + """Only POST is registered; the capability probe was removed.""" + from starlette.testclient import TestClient + + from deepagents_code import offload_api + + with TestClient(offload_api.app) as client: + response = client.get("/dcode/threads/thread-42/offload") + + assert response.status_code == 405 + + def test_cancel_path_is_registered(self) -> None: + from starlette.testclient import TestClient + + from deepagents_code import offload_api + + with TestClient(offload_api.app) as client: + response = client.post( + "/dcode/threads/thread-42/offload/op-1/cancel", + ) + + # No such operation is active, but the route resolved and its handler + # answered rather than 404-ing on an unmatched path. + assert response.status_code == 200, response.text + + +class TestThreadLock: + """Concurrent offloads of one thread are serialized in-process. + + The whole design is read-check-execute-recheck-write against a checkpoint. + The per-thread lock is what makes the recheck meaningful for two requests in + the same process: without it both can pass the status and checkpoint gates + before either writes. + """ + + def test_each_thread_gets_its_own_lock(self) -> None: + from deepagents_code import offload_api + + first = offload_api._thread_lock("thread-1") + + assert offload_api._thread_lock("thread-1") is first + assert offload_api._thread_lock("thread-2") is not first + + async def test_execute_waits_for_the_threads_lock(self) -> None: + """An offload must not touch thread state while the lock is held. + + Holding the lock externally proves the `async with` is on the path: + remove it, or key it on `operation_id` instead of `thread_id`, and the + operation reads state immediately. + """ + from deepagents_code import offload_api + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "busy"}), + get_state=AsyncMock(), + update_state=AsyncMock(), + ) + + async def run() -> None: + with contextlib.suppress(offload_api._OffloadConflictError): + await offload_api._execute_offload( + "thread-1", + operation_id="op-2", + context={}, + hook_responses={}, + ) + + with patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ): + async with offload_api._thread_lock("thread-1"): + blocked = asyncio.create_task(run()) + for _ in range(5): + await asyncio.sleep(0) + threads.get.assert_not_awaited() + + await asyncio.wait_for(blocked, timeout=5) + + threads.get.assert_awaited_once() + + async def test_a_different_thread_is_not_blocked(self) -> None: + """The lock is per thread, so unrelated threads must not serialize.""" + from deepagents_code import offload_api + + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "busy"}), + get_state=AsyncMock(), + update_state=AsyncMock(), + ) + + with patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ): + async with offload_api._thread_lock("thread-1"): + with pytest.raises(offload_api._OffloadConflictError): + await asyncio.wait_for( + offload_api._execute_offload( + "thread-2", + operation_id="op-1", + context={}, + hook_responses={}, + ), + timeout=5, + ) + + threads.get.assert_awaited_once() + + +class TestOffloadRoute: + """The HTTP layer maps operation outcomes onto distinct status codes.""" + + @staticmethod + def _request(payload: object) -> SimpleNamespace: + """Build a minimal Starlette-like request for the route handler.""" + return SimpleNamespace( + path_params={"thread_id": "thread-1"}, + json=AsyncMock(return_value=payload), + ) + + @staticmethod + def _cancel_request(operation_id: str = "op-1") -> SimpleNamespace: + """Build a minimal request for the cancellation route.""" + return SimpleNamespace( + path_params={"thread_id": "thread-1", "operation_id": operation_id} + ) + + async def test_malformed_request_is_422(self) -> None: + import json + + from deepagents_code import offload_api + + response = await offload_api.offload(self._request({"operation_id": ""})) # ty: ignore[invalid-argument-type] + + assert response.status_code == 422 + assert "operation_id" in json.loads(bytes(response.body))["detail"] + + async def test_cancel_stops_and_joins_an_active_operation(self) -> None: + """The cancel response is sent only after the operation task exits.""" + import json + + from deepagents_code import offload_api + + started = asyncio.Event() + stopped = asyncio.Event() + + async def execute(*_args: object, **_kwargs: object) -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + stopped.set() + + with patch.object( + offload_api, "_execute_offload", new=AsyncMock(side_effect=execute) + ): + operation = asyncio.create_task( + offload_api.offload( + self._request({"operation_id": "op-1", "context": {}}) # ty: ignore[invalid-argument-type] + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + response = await offload_api.cancel_offload(self._cancel_request()) # ty: ignore[invalid-argument-type] + + assert response.status_code == 200 + assert json.loads(bytes(response.body)) == {"status": "cancelled"} + assert stopped.is_set() + with pytest.raises(asyncio.CancelledError): + await operation + + async def test_cancel_before_request_prevents_operation_start(self) -> None: + """A reordered cancel closes the disconnect-before-register race.""" + import json + + from deepagents_code import offload_api + + cancel = await offload_api.cancel_offload(self._cancel_request()) # ty: ignore[invalid-argument-type] + execute = AsyncMock() + with patch.object(offload_api, "_execute_offload", new=execute): + response = await offload_api.offload( + self._request({"operation_id": "op-1", "context": {}}) # ty: ignore[invalid-argument-type] + ) + + assert json.loads(bytes(cancel.body)) == {"status": "cancelled"} + assert response.status_code == 409 + assert "cancelled" in json.loads(bytes(response.body))["detail"] + execute.assert_not_awaited() + + async def test_conflict_is_409(self) -> None: + import json + + from deepagents_code import offload_api + + with patch.object( + offload_api, + "_execute_offload", + new=AsyncMock( + side_effect=offload_api._OffloadConflictError("thread is busy") + ), + ): + response = await offload_api.offload( + self._request({"operation_id": "op-1", "context": {}}) # ty: ignore[invalid-argument-type] + ) + + assert response.status_code == 409 + assert json.loads(bytes(response.body))["detail"] == "thread is busy" + + async def test_indeterminate_write_is_500_with_its_own_detail(self) -> None: + import json + + from deepagents_code import offload_api + + with patch.object( + offload_api, + "_execute_offload", + new=AsyncMock( + side_effect=offload_api._OffloadIndeterminateError("cannot confirm") + ), + ): + response = await offload_api.offload( + self._request({"operation_id": "op-1", "context": {}}) # ty: ignore[invalid-argument-type] + ) + + assert response.status_code == 500 + assert json.loads(bytes(response.body))["detail"] == "cannot confirm" + + async def test_unbuildable_runtime_is_503_and_does_not_exit(self) -> None: + """The startup barrier must not kill the process from a request handler. + + `get_server_runtime` answers a construction failure with `sys.exit(1)`, + which is correct for the `langgraph.json` graph factory and fatal here: + `SystemExit` is a `BaseException`, so without containment it escapes the + route entirely and takes the server down mid-request. + """ + import json + + from deepagents_code import offload_api + + with patch.object( + offload_api, + "_execute_offload", + new=AsyncMock( + side_effect=offload_api._OffloadUnavailableError("runtime failed") + ), + ): + response = await offload_api.offload( + self._request({"operation_id": "op-1", "context": {}}) # ty: ignore[invalid-argument-type] + ) + + assert response.status_code == 503 + assert json.loads(bytes(response.body))["detail"] == "runtime failed" + + async def test_a_startup_exit_becomes_unavailable_not_a_process_exit( + self, + ) -> None: + """`SystemExit` from the runtime resolves to a typed error, not an exit.""" + from deepagents_code import offload_api + + before = _thread_state() + threads = SimpleNamespace( + get=AsyncMock(return_value={"status": "idle"}), + get_state=AsyncMock(return_value=before), + update_state=AsyncMock(), + ) + with ( + patch.object( + offload_api, + "get_client", + return_value=SimpleNamespace(threads=threads), + ), + patch.object( + offload_api, + "get_server_runtime", + new=AsyncMock(side_effect=SystemExit(1)), + ), + pytest.raises(offload_api._OffloadUnavailableError, match="unavailable"), + ): + await offload_api._execute_offload( + "thread-1", + operation_id="operation-1", + context={}, + hook_responses={}, + ) + + threads.update_state.assert_not_awaited() + + async def test_internal_type_error_is_500_not_422(self) -> None: + """A server-side shape fault must not be reported as a client error.""" + import json + + from deepagents_code import offload_api + + with patch.object( + offload_api, + "_execute_offload", + new=AsyncMock(side_effect=TypeError("LangGraph returned non-object state")), + ): + response = await offload_api.offload( + self._request({"operation_id": "op-1", "context": {}}) # ty: ignore[invalid-argument-type] + ) + + assert response.status_code == 500 + assert "server log" in json.loads(bytes(response.body))["detail"] diff --git a/libs/code/tests/unit_tests/test_remote_client.py b/libs/code/tests/unit_tests/test_remote_client.py index 84f2ff1b9e..b0394e3c12 100644 --- a/libs/code/tests/unit_tests/test_remote_client.py +++ b/libs/code/tests/unit_tests/test_remote_client.py @@ -1,5 +1,8 @@ """Tests for RemoteAgent, _convert_message_data, and helpers.""" +import asyncio +import itertools +import logging import uuid from collections.abc import Sequence from types import SimpleNamespace @@ -24,6 +27,18 @@ _TEST_THREAD_ID = "01966f3a-0000-7000-8000-000000000001" +_COMPACTED_RESULT = { + "status": "compacted", + "messages_offloaded": 2, + "messages_kept": 3, + "tokens_before": 100, + "tokens_after": 40, + "archive_path": "/conversation_history/thread.md", + "archive_ephemeral": False, + "error": None, +} +"""A well-formed `compacted` result, for tests that perturb one field.""" + # --------------------------------------------------------------------------- # _prepare_config @@ -1086,3 +1101,393 @@ 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" + + +def _offload_graph(http: SimpleNamespace) -> SimpleNamespace: + """Build a graph stub that also satisfies `aensure_thread`. + + `aoffload` registers the thread before its first POST, so a stub that only + carries `client.http` no longer suffices. `threads.create` is recorded so + tests can assert the registration happened, and happened first. + """ + threads = SimpleNamespace(create=AsyncMock(return_value=None)) + client = SimpleNamespace(http=http, threads=threads) + return SimpleNamespace( + client=client, + _validate_client=lambda: client, + ) + + +class TestServerOffload: + """The remote client transports operation data without graph state.""" + + async def test_cancellation_waits_for_server_acknowledgement(self) -> None: + """Esc must not release the caller while server offload is still live.""" + request_started = asyncio.Event() + cancel_started = asyncio.Event() + acknowledge_cancel = asyncio.Event() + + async def post(path: str, **_kwargs: object) -> dict[str, object]: + if path.endswith("/cancel"): + cancel_started.set() + await acknowledge_cancel.wait() + return {"status": "cancelled"} + request_started.set() + await asyncio.Event().wait() + return {} + + http = SimpleNamespace(post=AsyncMock(side_effect=post)) + graph = _offload_graph(http) + agent = RemoteAgent("http://localhost:1234") + + with patch.object(agent, "_get_graph", return_value=graph): + task = asyncio.create_task( + agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={}, + fulfill_hook=AsyncMock(), + ) + ) + await asyncio.wait_for(request_started.wait(), timeout=1) + task.cancel() + await asyncio.wait_for(cancel_started.wait(), timeout=1) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + acknowledge_cancel.set() + with pytest.raises(asyncio.CancelledError): + await task + + assert http.post.await_count == 2 + request_call, cancel_call = http.post.await_args_list + operation_id = request_call.kwargs["json"]["operation_id"] + assert cancel_call.args[0] == ( + f"/dcode/threads/thread/offload/{operation_id}/cancel" + ) + + async def test_registers_the_thread_before_the_first_request(self) -> None: + """The operation must not be requested against an unregistered thread. + + Checkpoint persistence and HTTP thread registration are separate on the + dev server, so a resumed thread has state on disk and no live row, and + every request below would 404. Ordering is the whole point -- registering + after the POST would not help -- so assert the call sequence rather than + just that both calls happened. + """ + calls: list[str] = [] + + async def record_post( # noqa: RUF029 -- must satisfy the async post signature + *_args: object, **_kwargs: object + ) -> dict[str, object]: + calls.append("post") + return {"status": "complete", "result": dict(_COMPACTED_RESULT)} + + async def record_create( # noqa: RUF029 -- must satisfy the async create signature + *_args: object, **_kwargs: object + ) -> None: + calls.append("create") + + http = SimpleNamespace(post=AsyncMock(side_effect=record_post)) + graph = _offload_graph(http) + graph.client.threads.create.side_effect = record_create + + agent = RemoteAgent("http://localhost:1234") + with patch.object(agent, "_get_graph", return_value=graph): + await agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={"model": "test:model"}, + fulfill_hook=AsyncMock(), + ) + + assert calls == ["create", "post"] + create_kwargs = graph.client.threads.create.await_args.kwargs + assert create_kwargs["thread_id"] == "thread" + assert create_kwargs["if_exists"] == "do_nothing" + + async def test_fulfills_hook_and_returns_typed_result(self) -> None: + agent = RemoteAgent("http://localhost:1234") + result = { + "status": "compacted", + "messages_offloaded": 2, + "messages_kept": 3, + "tokens_before": 100, + "tokens_after": 40, + "archive_path": "/conversation_history/thread.md", + "archive_ephemeral": False, + "error": None, + } + http = SimpleNamespace( + post=AsyncMock( + side_effect=[ + { + "status": "interrupt", + "request": { + "type": "hook_invocation", + "request": {"invocation_id": "hook-1"}, + }, + }, + {"status": "complete", "result": result}, + ] + ) + ) + graph = _offload_graph(http) + fulfill = AsyncMock(return_value={"decision": "allow"}) + + with patch.object(agent, "_get_graph", return_value=graph): + actual = await agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={"model": "test:model"}, + fulfill_hook=fulfill, + ) + + assert actual == result + assert http.post.await_count == 2 + first = http.post.await_args_list[0].kwargs["json"] + second = http.post.await_args_list[1].kwargs["json"] + assert first["context"] == {"model": "test:model"} + assert "messages" not in first + assert first["operation_id"] == second["operation_id"] + assert second["hook_responses"] == {"hook-1": {"decision": "allow"}} + fulfill.assert_awaited_once() + + async def test_missing_route_names_the_cause(self) -> None: + """A server without the route must not surface a bare "404 Not Found". + + A custom `graph_ref` server never registers dcode's HTTP app. An + unregistered thread cannot reach here as a 404 -- the server answers + 409 for that -- so a 404 means the route is absent, and the message + should say so and name a fix. + """ + import httpx + from langgraph_sdk.errors import NotFoundError + + agent = RemoteAgent("http://localhost:1234") + request = httpx.Request("POST", "http://localhost/dcode/threads/t/offload") + http = SimpleNamespace( + post=AsyncMock( + side_effect=NotFoundError( + "404 Not Found", + response=httpx.Response(404, request=request), + body=None, + ) + ) + ) + graph = _offload_graph(http) + + with ( + patch.object(agent, "_get_graph", return_value=graph), + pytest.raises(RuntimeError, match="does not provide dcode's /offload"), + ): + await agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={}, + fulfill_hook=AsyncMock(), + ) + + async def test_hook_interrupt_payload_round_trips_from_the_server(self) -> None: + """A real server-built interrupt payload must survive the client's parse. + + Uses `build_hook_interrupt_payload` output rather than a hand-written + dict, and feeds the client's reply back through the server-side lookup + key, so a payload-field rename or a UUID/str key mismatch fails here + instead of breaking `/offload` only for users with hooks configured. + """ + from datetime import UTC, datetime, timedelta + from pathlib import Path + from uuid import uuid4 + + from deepagents_code.hooks.interrupt import build_hook_interrupt_payload + from deepagents_code.hooks.models.domain import ( + ApprovalMode, + HookContext, + HookEvent, + HookInvocation, + PreCompactEvent, + ) + from deepagents_code.hooks.models.transport import HookInvocationRequest + + invocation_id = uuid4() + request = HookInvocationRequest( + protocol_version=1, + invocation_id=invocation_id, + snapshot_id="snapshot-1", + run_id="run-1", + invocation=HookInvocation( + context=HookContext( + thread_id="thread", + cwd=Path("/tmp"), + approval_mode=ApprovalMode.MANUAL, + ), + event=PreCompactEvent(event=HookEvent.PRE_COMPACT, trigger="manual"), + ), + deadline=datetime.now(UTC) + timedelta(seconds=60), + ) + payload = build_hook_interrupt_payload(request) + + agent = RemoteAgent("http://localhost:1234") + result = { + "status": "compacted", + "messages_offloaded": 1, + "messages_kept": 1, + "tokens_before": 10, + "tokens_after": 5, + "archive_path": "/conversation_history/thread.md", + "archive_ephemeral": False, + "error": None, + } + http = SimpleNamespace( + post=AsyncMock( + side_effect=[ + {"status": "interrupt", "request": payload}, + {"status": "complete", "result": result}, + ] + ) + ) + graph = _offload_graph(http) + fulfill = AsyncMock(return_value={"decision": "allow"}) + + with patch.object(agent, "_get_graph", return_value=graph): + actual = await agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={}, + fulfill_hook=fulfill, + ) + + assert actual == result + # The key the client accumulates must be exactly the key the server's + # `_invoke_hook` looks up: `str(request.invocation_id)`. + responses = http.post.await_args_list[1].kwargs["json"]["hook_responses"] + assert responses == {str(invocation_id): {"decision": "allow"}} + + @pytest.mark.parametrize( + ("result", "match"), + [ + ({"status": "compacted"}, "messages_offloaded"), + ({**_COMPACTED_RESULT, "tokens_before": "100"}, "tokens_before"), + ({**_COMPACTED_RESULT, "tokens_after": True}, "tokens_after"), + ({}, "no status"), + ("not a dict", "without a typed result"), + ], + ) + async def test_malformed_complete_result_is_refused( + self, result: object, match: str + ) -> None: + """A drifted payload must fail naming the field, not `KeyError` later. + + The renderer indexes these fields unguarded, and the server has already + committed the compaction by this point, so a `KeyError` here would be + reported as "Offload failed" for work that actually succeeded. + """ + agent = RemoteAgent("http://localhost:1234") + http = SimpleNamespace( + post=AsyncMock(return_value={"status": "complete", "result": result}) + ) + graph = _offload_graph(http) + + with ( + patch.object(agent, "_get_graph", return_value=graph), + pytest.raises(RuntimeError, match=match), + ): + await agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={}, + fulfill_hook=AsyncMock(), + ) + + async def test_non_compacted_result_needs_no_statistics(self) -> None: + """`empty`/`noop`/`denied` results carry no stats the renderer reads.""" + agent = RemoteAgent("http://localhost:1234") + result = {"status": "denied", "error": "Blocked by a compaction hook"} + http = SimpleNamespace( + post=AsyncMock(return_value={"status": "complete", "result": result}) + ) + graph = _offload_graph(http) + + with patch.object(agent, "_get_graph", return_value=graph): + actual = await agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={}, + fulfill_hook=AsyncMock(), + ) + + assert actual == result + + async def test_round_limit_logs_the_ids_it_saw( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Exhaustion must be diagnosable and must not assert a cause.""" + from deepagents_code.client.remote_client import _OFFLOAD_MAX_RESUME_ROUNDS + + agent = RemoteAgent("http://localhost:1234") + counter = itertools.count() + + async def _always_interrupt( # noqa: RUF029 # must be awaitable + *_args: object, **_kwargs: object + ) -> dict: + return { + "status": "interrupt", + "request": { + "type": "hook_invocation", + "request": {"invocation_id": f"hook-{next(counter)}"}, + }, + } + + http = SimpleNamespace(post=_always_interrupt) + graph = _offload_graph(http) + + with ( + patch.object(agent, "_get_graph", return_value=graph), + caplog.at_level(logging.WARNING), + pytest.raises(RuntimeError, match="hook rounds"), + ): + await agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={}, + fulfill_hook=AsyncMock(return_value={}), + ) + + assert f"exceeded {_OFFLOAD_MAX_RESUME_ROUNDS} hook rounds" in caplog.text + assert "hook-0" in caplog.text + + async def test_round_limit_does_not_fulfill_an_extra_hook( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """The final round reads a result; it must not answer another hook. + + The loop runs `_OFFLOAD_MAX_RESUME_ROUNDS + 1` times because the extra + iteration exists to POST the last fulfillment and read the reply. Drop + the guarding `break` and it fulfills one hook too many while still + reporting the lower number, so assert the count, not just the message. + """ + from deepagents_code.client.remote_client import _OFFLOAD_MAX_RESUME_ROUNDS + + agent = RemoteAgent("http://localhost:1234") + counter = itertools.count() + + async def _always_interrupt( # noqa: RUF029 # must be awaitable + *_args: object, **_kwargs: object + ) -> dict[str, object]: + return { + "status": "interrupt", + "request": { + "type": "hook_invocation", + "request": {"invocation_id": f"hook-{next(counter)}"}, + }, + } + + http = SimpleNamespace(post=_always_interrupt) + graph = _offload_graph(http) + fulfill = AsyncMock(return_value={}) + + with ( + patch.object(agent, "_get_graph", return_value=graph), + caplog.at_level(logging.WARNING), + pytest.raises(RuntimeError, match="hook rounds"), + ): + await agent.aoffload( + config={"configurable": {"thread_id": "thread"}}, + context={}, + fulfill_hook=fulfill, + ) + + assert fulfill.await_count == _OFFLOAD_MAX_RESUME_ROUNDS diff --git a/libs/code/tests/unit_tests/test_server_graph.py b/libs/code/tests/unit_tests/test_server_graph.py index 2fe1d14b2c..a8865b262a 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 from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -29,6 +30,15 @@ def _module_with_attrs(name: str, **attrs: object) -> ModuleType: return module +def _backend_with_offload(default: object) -> SimpleNamespace: + """Build a minimal backend carrying the server operation resource.""" + from deepagents_code.offload_middleware import OffloadOperation + + backend = SimpleNamespace(default=default) + backend._dcode_offload_operation = OffloadOperation(MagicMock(), MagicMock()) + return backend + + class TestServerGraph: """Tests for server-mode graph bootstrap.""" @@ -38,13 +48,43 @@ 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(), 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_concurrent_resolution_builds_one_runtime(self) -> None: + """Concurrent requests share the single graph runtime.""" + import asyncio + + module = _import_fresh_server_graph() + graph_obj = object() + calls = 0 + + async def build() -> object: + nonlocal calls + calls += 1 + await asyncio.sleep(0) + return module.ServerRuntime(graph_obj, object(), object()) + + factory = module._build_graph_factory(build) + results = await asyncio.gather(factory(), factory(), factory()) + + assert calls == 1 + assert results == [graph_obj, graph_obj, graph_obj] + + def test_server_runtime_slots_are_named(self) -> None: + """Both opaque runtime slots are named to prevent transposition.""" + module = _import_fresh_server_graph() + + assert module.ServerRuntime._fields == ("agent", "backend", "offload") + 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 +160,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, @@ -154,7 +194,7 @@ async def test_auto_discovery_loads_mcp_without_explicit_config(self) -> None: def create_cli_agent_side_effect(**_: object) -> tuple[object, object]: create_cli_agent_thread_ids.append(threading.get_ident()) - return graph_obj, SimpleNamespace(default=repository_backend) + return graph_obj, _backend_with_offload(repository_backend) def create_model_side_effect(*_: object, **__: object) -> object: create_model_thread_ids.append(threading.get_ident()) @@ -225,7 +265,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"], ) @@ -380,7 +420,7 @@ def create_cli_agent_side_effect(**_: object) -> tuple[object, object]: observed["interpreter_ptc"] = settings.interpreter_ptc observed["acknowledge"] = settings.interpreter_ptc_acknowledge_unsafe observed["enable_interpreter"] = settings.enable_interpreter - return graph_obj, SimpleNamespace(default=object()) + return graph_obj, _backend_with_offload(object()) settings_obj = SimpleNamespace( has_tavily=False, diff --git a/libs/code/tests/unit_tests/test_server_manager.py b/libs/code/tests/unit_tests/test_server_manager.py index 974e9609aa..4c896806e8 100644 --- a/libs/code/tests/unit_tests/test_server_manager.py +++ b/libs/code/tests/unit_tests/test_server_manager.py @@ -550,6 +550,38 @@ 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_builtin_server_registers_only_the_agent_graph( + self, tmp_path: Path + ) -> None: + """Operations use an authenticated route, not addressable siblings.""" + 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["http"] == { + "app": "deepagents_code.offload_api:app", + "enable_custom_route_auth": True, + } + + 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"} + assert "http" not in config + class TestWritePyproject: """Tests for the generated runtime pyproject."""