From 484ae704057f0803a34ea72eaafc0cc1f1f141e4 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 12:18:17 +0000 Subject: [PATCH 01/11] docs: specify Oz lifecycle hooks --- specs/APP-4344/PRODUCT.md | 538 ++++++++++++++++++++++++++ specs/APP-4344/TECH.md | 776 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 1314 insertions(+) create mode 100644 specs/APP-4344/PRODUCT.md create mode 100644 specs/APP-4344/TECH.md diff --git a/specs/APP-4344/PRODUCT.md b/specs/APP-4344/PRODUCT.md new file mode 100644 index 00000000000..ac9ecf1f68f --- /dev/null +++ b/specs/APP-4344/PRODUCT.md @@ -0,0 +1,538 @@ +# Oz Lifecycle Hooks + +Linear: [APP-4344](https://linear.app/warpdotdev/issue/APP-4344/add-claude-codecodex-style-lifecycle-hooks-to-the-oz-warp-agent) + +Origin: [Slack request](https://warpdev.slack.com/archives/C0BDQDW8V5E/p1788608487182989?thread_ts=1788608487.182989&cid=C0BDQDW8V5E) + +References: +- [Claude Code hooks](https://code.claude.com/docs/en/hooks) +- [Codex hooks](https://learn.chatgpt.com/docs/hooks) + +## Summary + +Add deterministic lifecycle hooks to Warp's first-party Oz harness. Users can run trusted commands when an Oz session starts or ends, when a prompt is submitted, around every tool call, before compaction, and when a turn stops. The first version supports command handlers only. `PreToolUse` can deny one tool call. No hook can grant permission, change tool input, stop the agent, or modify model context. + +## Problem + +Oz users cannot attach deterministic automation to the first-party harness lifecycle. They must ask the model to run checks or external automation and rely on the model to comply. + +Claude Code and Codex support lifecycle hooks for policy checks, audit logging, state capture, and notifications. Warp exposes those third-party harnesses, but their native hook systems do not apply to Oz. Oz also splits execution across the Warp client, the multi-agent server, and cloud workers. A client-only hook implementation would miss server-owned compaction and tool boundaries. + +The feature must provide one clear contract for local and cloud Oz runs. It must not weaken Warp permissions or expose secrets to hook processes. + +## Goals + +- Support these events for first-party Oz runs: + - `SessionStart` + - `SessionEnd` + - `UserPromptSubmit` + - `Stop` + - `PreToolUse` + - `PostToolUse` + - `PreCompact` +- Use event names and a JSON payload subset that is familiar to Claude Code and Codex hook authors. +- Execute local hooks on the local host. +- Execute cloud hooks inside the cloud worker sandbox. +- Cover client-executed and server-executed Oz tools. +- Let `PreToolUse` deny a tool call before it has side effects. +- Preserve Warp permission prompts, denials, and sandbox boundaries. +- Require explicit trust for project hook definitions. +- Redact and bound all hook payloads and outputs. +- Make configuration, execution, denial, timeout, and failure outcomes observable. + +## Non-goals + +- HTTP, MCP tool, prompt, agent, callback, or asynchronous hook handlers. +- Hook-driven tool-input or tool-output mutation. +- Hook-driven permission grants or persistent allow rules. +- Blocking `UserPromptSubmit`, `Stop`, `SessionEnd`, or `PreCompact`. +- Injecting hook output into model context, except for the reason from a `PreToolUse` denial. +- Claude Code or Codex feature parity beyond the named events and compatible payload subset. +- Replacing Warp execution profiles, permissions, enterprise policy, or sandboxing. +- Changing native Claude Code, Codex, Gemini, or OpenCode hook behavior. +- Automatically copying a user's local hook configuration into a cloud environment. +- A visual hook editor. Configuration is file-based in v1. +- Visual or computer-use validation. + +## Product behavior + +### 1. Configuration discovery + +1. Oz reads user hooks from `~/.warp/hooks.json` on the host that executes Oz. +2. Oz reads project hooks from `/.warp/hooks.json`. +3. Oz determines `` from the session's initial working directory. +4. Oz does not search parent directories above that Git root. +5. Oz does not load a project file when the initial working directory is not inside a Git repository. +6. The cloud runtime reads files from the sandbox filesystem. It does not read `~/.warp/hooks.json` from the user's laptop. +7. Environment setup may provision a cloud user hook file before Oz starts. +8. A checked-out repository may provide a cloud project hook file. +9. Oz snapshots configuration before `SessionStart`. File changes apply to the next session. +10. User and project hooks compose. One layer never replaces the other. +11. Oz evaluates user hooks before project hooks. +12. Oz preserves declaration order within each file. + +The v1 configuration schema is: + +```json +{ + "schema_version": "warp.oz_hooks.config.v1", + "hooks": { + "PreToolUse": [ + { + "matcher": "^(run_shell_command|apply_patch)$", + "hooks": [ + { + "type": "command", + "command": "python3 .warp/hooks/check_tool.py", + "command_windows": "py -3 .warp/hooks/check_tool.py", + "timeout": 10, + "on_failure": "deny" + } + ] + } + ] + } +} +``` + +The schema rules are: +- `schema_version` is required and must equal `warp.oz_hooks.config.v1`. +- `hooks` is required. +- Event keys must be one of the seven v1 event names. +- Each event contains ordered matcher groups. +- `matcher` is optional. +- An omitted matcher, an empty matcher, and `*` match every event occurrence. +- Every other matcher is a case-sensitive regular expression. +- `hooks` in a matcher group is a non-empty ordered array. +- `type` must equal `command`. +- `command` is required and must be non-empty. +- `command_windows` is optional. Windows uses it instead of `command` when present. +- `timeout` is an integer number of seconds. +- The default timeout is 10 seconds. +- A non-`SessionEnd` timeout must be from 1 through 120 seconds. +- The `SessionEnd` default is 1 second. Its configured maximum is 3 seconds. +- `on_failure` is optional and defaults to `continue`. +- `on_failure: "deny"` is valid only for `PreToolUse`. +- Unknown fields, invalid regular expressions, unsupported events, and unsupported values invalidate the containing file. +- An invalid file contributes no hooks. Hooks from another valid file still run. +- Oz reports one configuration diagnostic per invalid file. +- Each file is limited to 256 KiB and 64 command handlers. + +### 2. Matcher behavior + +Oz matches each event against one subject: +- `SessionStart`: `source`, with v1 values `startup` or `resume`. +- `SessionEnd`: `reason`, with v1 values `completed`, `failed`, `cancelled`, or `shutdown`. +- `PreToolUse` and `PostToolUse`: the canonical Oz tool name. +- `PreCompact`: `trigger`, with v1 values `auto` or `manual`. +- `UserPromptSubmit` and `Stop`: no match subject. Oz ignores `matcher` and runs every declared handler for the event. + +Oz uses the same canonical tool name in matching, payloads, diagnostics, and tests. It does not silently map Oz tool names to Claude Code aliases such as `Bash`, `Read`, or `Write`. + +### 3. Project hook trust + +1. User hooks are trusted because the user controls the host-level file. +2. Project hooks are disabled until the user explicitly trusts the exact project hook definition. +3. The trust record includes the canonical Git root, config path, and SHA-256 hash of the validated file bytes. +4. Oz displays the source path, event, matcher, command, timeout, failure mode, and hash before accepting trust. +5. A new file, changed byte, changed command, changed matcher, changed timeout, or changed failure mode creates a new hash. +6. A new hash requires a new trust decision. +7. Untrusted project hooks are skipped. The session continues with a visible diagnostic. +8. Headless runs never auto-trust project hooks. +9. A cloud project hook runs only when the sandbox receives a matching explicit trust record as part of the run or environment configuration. +10. The runtime never accepts a project trust decision from the project repository itself. +11. Revoking trust prevents the definition from running in future sessions. + +### 4. Lifecycle timing + +Oz emits events at these boundaries: + +1. `SessionStart` + - Fires once after environment setup, configuration validation, and trust evaluation complete. + - Fires before Oz processes the first prompt. + - Uses `source: "startup"` for a new conversation. + - Uses `source: "resume"` when an existing conversation resumes. +2. `UserPromptSubmit` + - Fires once for every user prompt. + - Fires after prompt attachments are resolved and before the prompt is sent to MAA. + - Does not fire for internal retries or synthetic model messages. +3. `PreToolUse` + - Fires once after a complete tool name and input exist. + - Fires after Warp computes its native permission classification. + - Fires before a permission prompt and before any tool side effect. + - Fires for every Oz tool that passes native Warp denial, including MCP, file, shell, document, computer-use, orchestration, and server-executed tools. +4. `PostToolUse` + - Fires once after an executed tool reaches a terminal success, failure, timeout, or cancellation result. + - Fires before the result is supplied to the next model inference. + - Does not fire when `PreToolUse` denied the tool because no tool executed. +5. `PreCompact` + - Fires immediately before MAA begins manual or automatic context compaction. + - MAA waits for the observational hook outcome or timeout before compaction begins. +6. `Stop` + - Fires once when an Oz turn has produced its final assistant output. + - Fires before the turn changes to an idle, blocked, failed, or completed state. + - Does not force another model turn. +7. `SessionEnd` + - Fires once during graceful session teardown after the final `Stop`, when applicable. + - Uses the terminal `reason` value. + - Is best effort for process crashes, force kills, host loss, and worker loss. + - Never delays teardown for more than 3 seconds. + +### 5. Execution order + +1. Oz maintains one FIFO hook event queue per conversation. +2. Events in one conversation never overtake each other. +3. Different conversations may execute hooks concurrently. +4. For one event, Oz runs matching command handlers sequentially. +5. Oz runs user handlers before project handlers. +6. Oz preserves file and declaration order. +7. A `PreToolUse` denial stops that event's handler chain immediately. +8. Oz does not start later matching handlers after a denial. +9. Session cancellation terminates the active hook process group or Windows Job Object, removes pending hook events, and ignores late results. +10. Cancelling a tool execution emits `PostToolUse` with a cancelled result when the conversation runtime remains active. + +This differs intentionally from Codex, which launches matching commands concurrently. Oz chooses deterministic order and denial short-circuiting so policy side effects and audit records are reproducible. + +### 6. Command execution + +1. A command hook receives one UTF-8 JSON object on stdin. +2. A command runs with the active session working directory. +3. macOS and Linux run `command` through `SHELL`, with `/bin/sh` as the fallback. +4. Windows runs `command_windows` when present and otherwise runs `command` through `COMSPEC`, with `cmd.exe` as the fallback. +5. A cloud command runs inside the worker task sandbox, not on the user's laptop and not in the worker control plane. +6. The hook process receives a rebuilt environment containing only: + - `HOME`, or `USERPROFILE` on Windows + - `PATH` + - `SHELL` on Unix + - `COMSPEC` and `SystemRoot` on Windows + - `TMPDIR`, `TMP`, or `TEMP` when present + - locale variables required for UTF-8 operation + - `WARP_HOOK_EVENT_NAME` + - `WARP_RUN_ID` + - `WARP_CONVERSATION_ID` +7. The hook process does not inherit managed secret values, API keys, cloud credentials, Git credentials, MCP credentials, or the complete Oz task environment. +8. The hook remains inside the same operating-system user, filesystem, network, container, and sandbox boundaries as the Oz runtime. +9. A project hook receives no privilege that the project does not already have inside that runtime. + +### 7. Input payload + +Every command receives this common envelope: + +```json +{ + "schema_version": "warp.oz_hook.v1", + "hook_event_name": "PreToolUse", + "session_id": "opaque-session-id", + "run_id": "opaque-run-id", + "conversation_id": "opaque-conversation-id", + "cwd": "/workspace/repository", + "hook_source": "user", + "model": "model-id", + "permission_mode": "supervised" +} +``` + +Event-specific fields are: +- `SessionStart`: `source`. +- `SessionEnd`: `reason`. +- `UserPromptSubmit`: `prompt`. +- `PreToolUse`: `tool_name`, `tool_use_id`, and `tool_input`. +- `PostToolUse`: `tool_name`, `tool_use_id`, `tool_input`, and `tool_response`. +- `PreCompact`: `trigger`. +- `Stop`: `turn_status`. + +Compatibility rules: +- Common and event-specific field names use snake case, matching the overlapping Claude Code and Codex command-hook shape. +- Event names preserve Claude Code and Codex capitalization. +- Warp-specific fields are additive. +- `hook_source` identifies the config layer and does not replace the Claude/Codex `SessionStart.source` field. +- Consumers must ignore unknown fields. +- Warp may add optional fields within `warp.oz_hook.v1`. +- A breaking change requires a new `schema_version`. + +### 8. Redaction and size limits + +Oz redacts data before JSON serialization, protocol transport, hook execution, telemetry, or logs. + +The payload must never include: +- the full process environment +- resolved secret values +- managed secret payloads +- API keys or authorization headers +- raw attachment bytes +- file contents +- an absolute transcript path +- a complete conversation transcript +- unbounded tool input or tool output + +The payload preserves useful structure: +- tool name and tool-use ID +- file and directory paths +- argument object keys +- scalar type information +- shell command text with known secrets and credentials replaced +- permission and risk categories +- result status, exit code, duration, byte counts, item counts, and omitted counts +- bounded error and output previews after redaction + +The size contract is: +- Serialized stdin: 256 KiB maximum. +- Redacted prompt: 64 KiB maximum. +- Redacted tool input: 128 KiB maximum. +- Redacted tool response: 64 KiB maximum. +- Captured stdout: 64 KiB maximum. +- Captured stderr: 64 KiB maximum. +- Denial reason delivered to the user and model: 4 KiB maximum. + +Oz truncates only at valid UTF-8 boundaries. It adds explicit truncation metadata. It never writes an unredacted value before truncation. + +### 9. `PreToolUse` decisions + +A `PreToolUse` handler can return no decision or deny the current tool. + +Continue: +- Exit code 0 with empty stdout means no decision. +- Exit code 0 with `{}` means no decision. + +Deny with structured JSON: + +```json +{ + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": "Repository policy blocks this operation." + } +} +``` + +Deny with an exit code: +- Exit code 2 denies the tool. +- Non-empty stderr becomes the denial reason. + +Rules: +- `permissionDecision: "allow"` is invalid. +- `permissionDecision: "ask"` is invalid. +- `updatedInput`, `updatedMCPToolOutput`, `additionalContext`, `continue`, and `decision` control fields are invalid in v1. +- A denial blocks only the current tool call. +- Oz gives the bounded reason to the model so it can choose another action. +- Oz also surfaces the reason to the user. +- A denial never ends the session. +- A denial never creates an allow rule. + +Outputs from the other six events are observational: +- Exit code 0 is success. +- Stdout and stderr are bounded diagnostics only. +- Oz does not parse them as model context or control instructions. +- Exit code 2 has no blocking meaning outside `PreToolUse`. + +### 10. Permission composition + +Warp permissions remain authoritative. + +For each tool call: +1. Warp computes the native permission result. +2. If Warp denies the tool, Oz rejects it. A hook cannot override that result. +3. If Warp allows or would prompt, Oz runs `PreToolUse`. +4. If a hook denies, Oz rejects the tool without showing a permission prompt. +5. If hooks return no decision, Warp applies the original allow or prompt result. +6. The tool executes only after both systems permit it. + +Hooks can reduce authority. Hooks cannot increase authority. + +### 11. Failure and timeout behavior + +A hook failure includes: +- process spawn failure +- timeout +- non-zero exit other than a valid `PreToolUse` exit 2 +- invalid or oversized `PreToolUse` JSON +- unsupported output fields +- mismatched `hookEventName` +- stdout or stderr above its limit +- protocol correlation failure + +Behavior: +- Observational event failures always continue the Oz lifecycle. +- A `PreToolUse` failure with `on_failure: "continue"` records the failure and continues to the original Warp permission decision. +- A `PreToolUse` failure with `on_failure: "deny"` denies the tool with a generic bounded reason. +- `on_failure` applies independently to each handler. +- A failed handler does not stop later handlers unless its failure mode denies. +- A timeout kills the command's process group before the lifecycle continues. +- `SessionEnd` ignores `on_failure`, does not retry, and never exceeds its teardown cap. +- Oz does not retry command hooks automatically. + +### 12. Observability + +For each hook invocation, Oz provides a local diagnostic record containing: +- event name +- config source +- config path +- definition hash +- matcher +- start and finish timestamps +- duration +- result: `succeeded`, `continued`, `denied`, `failed`, `timed_out`, or `cancelled` +- exit code when available +- whether output was truncated +- failure category + +Remote telemetry omits the config path, matcher, and timestamps. It uses only the metadata allowlist in the technical specification. + +Oz does not record raw payloads, raw stdout, raw stderr, raw prompts, raw tool inputs, raw tool responses, or secret values in telemetry. + +The user can distinguish: +- no configured hook +- unmatched hook +- untrusted project hook +- invalid configuration +- successful hook +- denied tool +- failed-open hook +- failed-closed hook +- timed-out hook + +### 13. Local and cloud parity + +1. The seven event names and JSON schemas are identical for local and cloud Oz. +2. The execution host differs: + - Local Oz uses the local Warp host. + - Cloud Oz uses the worker task sandbox. +3. `cwd` is a path in the execution host. +4. A local path is never sent to a cloud hook as its working directory. +5. A cloud hook cannot execute on the worker daemon host outside the task sandbox. +6. Worker Direct, Docker, Kubernetes, and command-dispatched backends preserve the same contract when they support first-party Oz. +7. A backend that cannot execute the contract must reject hook-enabled runs rather than silently run with partial lifecycle coverage. +8. An MAA deployment that does not acknowledge `warp.oz_hook.v1` must reject or be rejected by a hook-enabled runtime before the runtime applies any tool action. + +## Decisions + +### Warp-native configuration + +Options: +- Read Claude Code and Codex configuration directly. +- Define Warp configuration with familiar event names. + +Decision: +- Use `~/.warp/hooks.json` and `.warp/hooks.json`. + +Why: +- Oz needs Warp-specific trust, redaction, permissions, cloud, and protocol semantics. +- Reusing third-party files would imply compatibility that v1 does not provide. + +### Full tool coverage + +Options: +- Hook only client-executed actions. +- Add a blocking MAA protocol round trip for server-owned tool and compaction boundaries. + +Decision: +- Add the protocol round trip. + +Why: +- Partial coverage would make policy hooks unreliable. +- A hook author must not need to know where a tool happens to execute. +- The added protocol work is preferable to a false enforcement guarantee. + +### Deny-only control + +Options: +- Support allow, ask, deny, and input mutation. +- Support deny only. + +Decision: +- Support deny only in `PreToolUse`. + +Why: +- Deny composes safely with Warp permissions. +- Allow and mutation could bypass user intent or invalidate the audited tool request. + +### Sequential execution + +Options: +- Match Codex and launch all handlers concurrently. +- Run handlers sequentially in deterministic order. + +Decision: +- Run handlers sequentially and stop after a denial. + +Why: +- Deterministic side effects and diagnostics are easier to reason about. +- Policy hooks can avoid unnecessary work after a denial. +- The trade-off is higher cumulative latency. + +### Fail-open default with explicit fail-closed policy + +Options: +- Always fail open. +- Always fail closed. +- Default to fail open and allow `PreToolUse` handlers to opt into fail closed. + +Decision: +- Use the third option. + +Why: +- Observability hooks must not break agent work. +- Security policy hooks need an explicit availability-over-progress choice. + +### Host-local execution + +Options: +- Execute every hook on the user's local machine. +- Execute hooks where the Oz runtime executes. + +Decision: +- Execute hooks where Oz executes. + +Why: +- Cloud paths, files, tools, and sandbox boundaries exist only in the worker task environment. +- Sending cloud tool data to a laptop would add latency and a new data boundary. + +## Assumptions + +- The exact trust-review presentation can use Warp's existing confirmation patterns. A separate visual editor is not required. +- Cloud environments that need user hooks will provision `~/.warp/hooks.json` and matching trust material before the Oz session starts. +- `manual` is reserved for a future explicit compaction action even if the first implementation only emits `auto`. +- Hook event and handler limits are sufficient for v1 policy and observability use cases. +- The 120-second non-`SessionEnd` timeout maximum is sufficient for synchronous v1 handlers. + +## Out of scope + +- Enterprise-managed hook layers and mandatory organization policy. +- Hook credential injection. +- Plugin-packaged hooks. +- Per-hook enable and disable controls. +- Live configuration reload. +- Hook retries. +- Cross-session hook state managed by Warp. +- A guarantee that `SessionEnd` runs after a process crash or infrastructure loss. +- Native alias matching for Claude Code tool names. + +## Validation criteria + +1. Config parser tests run with `cargo test -p warp oz_hooks_config` and cover: + - valid user and project files + - invalid schema versions + - unknown fields and events + - invalid regular expressions + - timeout bounds + - `on_failure` restrictions + - file and handler limits +2. Merge-order tests run with `cargo test -p warp oz_hooks_ordering` and prove user-before-project and declaration-order execution. +3. Trust tests run with `cargo test -p warp oz_hooks_trust` and prove new, changed, revoked, and untrusted project definitions do not run without the exact hash. +4. Payload golden tests run with `cargo test -p warp oz_hooks_payload` and cover all seven events. +5. Redaction tests run with `cargo test -p warp oz_hooks_redaction` and prove secrets, environments, attachments, file contents, transcripts, and oversized values never reach serialized payloads or logs. +6. Runtime tests run with `cargo test -p warp oz_hooks_runtime` and cover sequential execution, deny short-circuiting, process-group cancellation, timeout, spawn failure, non-zero exit, malformed JSON, unsupported fields, oversized output, fail-open, and fail-closed behavior. +7. Permission tests run with `cargo test -p warp oz_hooks_permissions` and prove hooks cannot upgrade a Warp deny or bypass a Warp prompt. +8. Local Oz integration tests emit all seven events, prove exact ordering, and prove a denied tool produces no side effect and no `PostToolUse`. +9. Multi-agent server tests run with `go test ./logic/ai/multi_agent/...` and prove server-owned tools and `PreCompact` pause for a correlated hook result. +10. Proto generation runs with `./script/generate -a multi_agent -v v1`. A following `git diff --exit-code` in `warp-proto-apis` proves generated bindings are current. +11. Cloud worker tests run with `go test ./internal/worker/...` and prove hook commands execute in the task workspace for Direct and containerized backends without inheriting worker credentials. +12. A cloud Oz integration test emits all seven events from the worker sandbox and proves its payload contract matches the local golden payloads. +13. Regression tests start Oz, Claude Code, Codex, Gemini, and OpenCode harnesses and prove only `HarnessKind::Oz` activates this hook runtime. +14. `./script/presubmit` passes in `warp`. +15. No visual or computer-use validation is required. + +## Open questions + +None. The requester approved this v1 direction. Implementation should begin only after the spec PR is approved. diff --git a/specs/APP-4344/TECH.md b/specs/APP-4344/TECH.md new file mode 100644 index 00000000000..6a47c48da1c --- /dev/null +++ b/specs/APP-4344/TECH.md @@ -0,0 +1,776 @@ +# Oz Lifecycle Hooks — Technical Specification + +Linear: [APP-4344](https://linear.app/warpdotdev/issue/APP-4344/add-claude-codecodex-style-lifecycle-hooks-to-the-oz-warp-agent) + +Product spec: `specs/APP-4344/PRODUCT.md` + +External references: +- [Claude Code hooks](https://code.claude.com/docs/en/hooks) +- [Codex hooks](https://learn.chatgpt.com/docs/hooks) + +## Summary + +Implement a first-party Oz hook runtime in the Warp client and embedded cloud Oz process. Add protocol gates so the same runtime can execute hooks for server-owned tool and compaction boundaries. Keep Warp permissions authoritative. Apply redaction before data crosses a process or network boundary. Do not modify third-party harness setup or native hooks. + +## Relevant code + +### Warp + +- `app/src/ai/agent_sdk/driver/harness/mod.rs (191-273)` — `HarnessKind::Oz` is separate from `ThirdPartyHarness`. +- `app/src/ai/agent_sdk/driver.rs (2070-2319)` — Oz setup, MCP startup, environment preparation, and skill loading. +- `app/src/ai/agent_sdk/driver.rs (3416-3615)` — `AgentDriver::execute_run` and conversation lifecycle subscriptions. +- `app/src/ai/blocklist/action_model/execute.rs (593-760)` — `BlocklistAIActionExecutor::try_to_execute_action` computes native permission behavior and dispatches client actions. +- `app/src/ai/blocklist/permissions.rs (1-132)` — typed command, read, and write permission results. +- `app/src/ai/blocklist/permissions.rs (352-486)` — execution-profile command permissions and deny lists. +- `app/src/ai/blocklist/permissions.rs (1217-1248)` — system-protected executable configuration paths. +- `app/src/ai/mcp/mod.rs (482-599)` — Warp user and project config-path conventions. +- `app/src/ai/mcp/file_based_manager.rs (246-444)` — content hashing, source scope, and project auto-start restrictions for executable configuration. +- `crates/warp_cli/src/agent.rs (211-279)` — first-party and third-party harness enum. +- `crates/warp_cli/src/lib.rs (111-299)` — shared Oz/Warp CLI argument tree. + +### Multi-agent server + +- `logic/ai/multi_agent/utils/output/tool_call_processor.go (86-286)` — native tool-call parsing and client-action production. +- `logic/ai/multi_agent/compression/summarize/summarize.go (48-78)` — summarization lifecycle callback interface. +- `logic/ai/multi_agent/compression/summarize/summarize.go (199-219)` — entry to server-owned context-window summarization. + +### Protocol + +- `apis/multi_agent/v1/request.proto (21-124)` — request input and user-input variants. +- `apis/multi_agent/v1/request.proto (460-701)` — client capability settings. +- `apis/multi_agent/v1/response.proto (14-59)` — streamed response envelope. +- `apis/multi_agent/v1/response.proto (352-454)` — server-to-client actions. +- `apis/multi_agent/v1/task.proto (24-53)` — task state and opaque server data. + +### Cloud worker + +- `internal/worker/backend.go (125-181)` — backend-neutral task parameters and task execution interface. +- `internal/worker/direct.go (58-72)` — minimal host environment inherited by direct tasks. +- `internal/worker/direct.go (152-284)` — direct workspace setup and embedded Oz execution. +- `internal/worker/docker.go (118-199)` — container task command, environment, and `/workspace` working directory. +- `internal/worker/dispatch_payload.go (9-55)` — versioned command-backend task payload. + +## Current state + +### Harness boundaries + +`HarnessKind` routes Oz through Warp's MAA-backed runtime. Claude Code, Codex, and Gemini use `ThirdPartyHarness` implementations and their own CLI configuration. The new hook runtime must be constructed only for `HarnessKind::Oz`. + +### Local client actions + +MAA emits `ClientAction` values. Warp converts them to `AIAgentAction` values. `BlocklistAIActionExecutor::try_to_execute_action` computes whether Warp can auto-execute the action, whether it needs user confirmation, and which action executor runs it. + +This is the final common client-side dispatch boundary. It is synchronous at entry and returns an async execution variant for actions that need one. Adding a blocking hook requires a staged asynchronous preflight before the existing confirmation and execution branch. + +### Server-owned boundaries + +The server parses native model tool calls in `nativeToolCallProcessor.ProduceActions`. Most calls become client actions, but the tool abstraction also permits server-owned processing. Context compaction is explicitly server-owned in `SummarizeMessagesForContextWindow`. + +A client-only hook runtime cannot guarantee `PreToolUse`, `PostToolUse`, or `PreCompact` at these boundaries. The protocol must let MAA request a hook invocation and await a correlated result before it continues. + +### Cloud execution + +The worker launches the same Warp/Oz binary inside the task execution environment: +- Direct uses the task workspace as `cmd.Dir`. +- Docker uses `/workspace` as the container working directory. +- Kubernetes and command-dispatched backends also launch the task runtime outside the worker control plane. + +The embedded Oz process, not the worker daemon, must own hook discovery and execution. The worker only needs to preserve required task metadata, trust material, cancellation, and sandbox placement. + +## Technical design + +### 1. Add a shared Oz hook module + +Add an Oz-only module under `app/src/ai/agent_sdk/hooks/` with these responsibilities: +- `config`: discover, parse, validate, hash, and merge hook files. +- `trust`: evaluate exact project-definition trust. +- `matcher`: compile matchers and select handlers. +- `payload`: define the versioned event envelope. +- `redaction`: convert internal prompt, action, result, and compaction data to safe payloads. +- `runtime`: queue events, spawn commands, enforce limits, parse `PreToolUse` output, and aggregate outcomes. +- `telemetry`: emit metadata-only execution events. + +Expose a narrow runtime interface: + +```rust +pub(crate) trait OzHookRuntime { + async fn observe(&self, event: OzHookEvent) -> OzHookObservation; + async fn pre_tool_use(&self, event: OzPreToolUseEvent) -> OzPreToolUseDecision; + fn cancel(&self, scope: OzHookCancellationScope); +} +``` + +`observe` never returns control effects. `pre_tool_use` returns only `Continue` or `Deny { reason, source }`. + +Do not add hook methods to `ThirdPartyHarness`. This runtime is a first-party Oz service, not a generalized wrapper around native third-party hook systems. + +### 2. Configuration model + +Deserialize with strict unknown-field rejection. + +The in-memory model should distinguish: +- `HookConfigFile` +- `HookEventName` +- `MatcherGroup` +- `CommandHandler` +- `HookConfigSource::User` +- `HookConfigSource::Project` +- `FailureMode::Continue` +- `FailureMode::Deny` + +Validation happens before the file contributes handlers: +1. Enforce the 256 KiB byte limit. +2. Parse JSON. +3. Validate `schema_version`. +4. Reject unknown fields and events. +5. Validate handler counts. +6. Compile regular expressions. +7. Validate timeout bounds. +8. Reject `on_failure: "deny"` outside `PreToolUse`. +9. Compute SHA-256 over the exact validated file bytes. +10. Apply trust to the project file. +11. Merge user then project handlers without deduplication. + +Do not canonicalize JSON before hashing. An exact byte change must require a new project trust decision. This matches the product contract and avoids ambiguity about semantically equivalent but differently audited files. + +Snapshot the merged config once per conversation. Store the source path and definition hash on every configured handler for diagnostics. + +### 3. Trust material + +Add a `HookTrustStore` abstraction. A trust key contains: +- canonical Git root +- canonical config path +- SHA-256 file hash + +The store must not accept trust data from `.warp/hooks.json` or another project file. + +Local runs use private user state. Cloud runs receive signed or server-authenticated trust records in task metadata. The embedded runtime verifies that the record matches the canonical repository identity, config path, and hash it discovers inside the sandbox. + +When a cloud run has no matching record: +- Skip the project file. +- Emit an `untrusted_project_hooks` setup diagnostic. +- Continue the run with valid user hooks. + +The run-launch or environment-management surface that records cloud trust must show the exact validated handler definitions described in the product spec. The trust transport must contain hashes and identity only. It must not contain command output, hook payloads, or secrets. + +Add `.warp/hooks.json` and the host trust store to the same system-protected write classification used for MCP configuration in `app/src/ai/blocklist/permissions.rs`. An Oz action must not auto-write either path. A user-confirmed write still creates a new project hash that remains untrusted. + +### 4. Runtime ownership and lifecycle + +Create one `OzHookRuntimeHandle` per Oz conversation. + +For local Oz: +- Construct it when the conversation has an initial working directory and run identifiers. +- Load configuration before the first prompt. +- Store the handle with conversation-scoped state used by prompt and action execution. + +For cloud Oz: +- Construct it in `AgentDriver` after terminal bootstrap and environment preparation. +- Construct it after repositories and setup commands are complete so project files exist. +- Construct it before the initial prompt enters `execute_run`. +- Execute commands through the embedded Warp process in the task sandbox. + +Fire `SessionStart` after configuration and trust evaluation. Register graceful teardown so `SessionEnd` receives the final reason and a hard 3-second cap. + +`SessionEnd` cannot be guaranteed after SIGKILL, container loss, worker loss, or host loss. Do not report a synthetic successful `SessionEnd` in those cases. + +### 5. Event queue and ordering + +Each runtime owns one FIFO queue keyed to the conversation. + +Processing rules: +- Dequeue one lifecycle event at a time. +- Resolve matching handlers from the immutable config snapshot. +- Run handlers sequentially. +- Preserve source, group, and handler declaration order. +- Stop a `PreToolUse` chain after the first explicit deny or fail-closed failure. +- Continue after fail-open failures. +- Let different conversation runtimes process independently. + +Every event has an opaque invocation ID. Tool events also carry the stable tool-use ID. Protocol results must match both identifiers before the runtime applies them. + +Hook or session cancellation: +- Cancel the event future. +- Kill the active process group. +- Drop queued events scoped to the cancelled operation. +- Mark the invocation cancelled. +- Reject late local and protocol results. + +Tool execution cancellation is a tool result, not automatic hook-runtime cancellation. Emit `PostToolUse` with terminal status `cancelled` when the conversation runtime remains active. + +### 6. Command runner + +Use a dedicated subprocess runner. Do not reuse the normal agent shell action executor because hook commands must not create recursive `PreToolUse` events. + +Runner behavior: +- Select `command_windows` on Windows when present. +- Otherwise use `command`. +- Start the command through the active session shell. +- Set the command working directory to the event `cwd`. +- Create a process group on Unix or a Job Object on Windows. +- Build a new environment from the explicit allowlist in the product spec. +- Write one serialized payload to stdin and close stdin. +- Capture stdout and stderr independently. +- Enforce 64 KiB per stream while reading. +- Kill the process group or Job Object on timeout, cancellation, or output overflow. +- Decode output as UTF-8. +- Record duration and exit status. +- Never log the command's stdin or raw output. + +The cloud Direct backend already starts the task from a minimal host environment in `internal/worker/direct.go (58-72, 256-264)`. The hook runner must still rebuild its own environment because task-level environment variables include resolved secrets. + +### 7. Payload and redaction + +Define typed Rust payload structs with a common envelope and event-specific flattened fields. Define equivalent typed Go payload templates for server-owned events. Serialize only after redaction. + +Common fields: +- `schema_version` +- `hook_event_name` +- `session_id` +- `run_id` +- `conversation_id` +- `cwd` +- `hook_source` +- `model` +- `permission_mode` + +Event fields: +- `SessionStart`: `source` +- `SessionEnd`: `reason` +- `UserPromptSubmit`: `prompt` +- `PreToolUse`: `tool_name`, `tool_use_id`, `tool_input` +- `PostToolUse`: `tool_name`, `tool_use_id`, `tool_input`, `tool_response` +- `PreCompact`: `trigger` +- `Stop`: `turn_status` + +Add a Rust redaction adapter for every `AIAgentActionTypeDiscriminants` value. Add a Go redaction adapter for every server-owned tool category. Do not generically serialize an internal action or tool-call object as the hook payload. Generic serialization can expose fields that have not received a redaction review. + +The adapters should retain: +- stable tool name +- paths +- argument keys +- non-sensitive enum values +- command text after secret replacement +- risk category +- status and numeric metadata +- bounded safe previews + +Replace omitted values with explicit objects such as: + +```json +{ + "content": { + "redacted": true, + "reason": "file_content", + "byte_count": 18432 + } +} +``` + +Redaction order: +1. Convert the internal event to an allowlisted intermediate representation. +2. Replace known secret values and credential patterns. +3. Remove prohibited fields. +4. Truncate bounded leaf values. +5. Add omitted and truncation metadata. +6. Enforce the event-specific size budget. +7. Serialize. +8. Enforce the final 256 KiB limit. + +Use explicit allowlisted intermediate types for local execution and protocol transport. Test the Rust and Go types against the same canonical JSON golden fixtures. + +### 8. `PreToolUse` output parser + +Accept: +- Exit 0 with empty stdout. +- Exit 0 with `{}`. +- Exit 0 with the exact compatible deny subset. +- Exit 2 with non-empty stderr. + +The structured subset is: + +```json +{ + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": "bounded reason" + } +} +``` + +Use strict parsing: +- Reject unknown top-level fields. +- Reject unknown `hookSpecificOutput` fields. +- Require `hookEventName` to equal `PreToolUse`. +- Accept only `permissionDecision: "deny"`. +- Require a non-empty reason for structured denial. +- Reject mutation and context fields. +- Treat multiple JSON values or non-whitespace trailing data as malformed. + +Exit precedence: +- Timeout, cancellation, spawn failure, and output overflow are execution failures. +- Exit 2 with non-empty stderr is denial. +- Exit 2 with empty stderr is failure. +- Exit 0 parses stdout. +- Every other exit code is failure. + +Apply the handler's `on_failure` only to failures. Explicit denial always denies. + +### 9. Warp permission composition + +Refactor the client action path into stages: +1. Compute the existing Warp permission classification without side effects. +2. Reject an existing Warp deny. +3. Build the redacted hook event. +4. Await `PreToolUse`. +5. Reject a hook deny. +6. Show the existing Warp confirmation when the native classification requires it. +7. Execute the original action object without modification. +8. Build and await `PostToolUse`. +9. Send the original tool result to MAA. + +`BlocklistAIActionExecutor::try_to_execute_action` currently combines permission decisions and dispatch. Introduce an async preflight state rather than blocking the model thread. Preserve the existing `NotExecutedReason::NeedsConfirmation` behavior after hook continuation. + +The hook runtime must never return an action object. This type boundary prevents input mutation by construction. + +When Warp denies before the hook stage: +- Do not execute the tool. +- Do not emit `PreToolUse`. +- Do not emit `PostToolUse`. +- Return the existing Warp denial result. + +When a hook denies: +- Produce a tool result that clearly attributes the denial to a trusted Oz hook. +- Include the bounded reason. +- Let the model continue and choose another tool. +- Do not change the native permission profile. + +### 10. Protocol contract + +Extend `warp-proto-apis` with first-party Oz hook messages. + +Add `apis/multi_agent/v1/oz_hooks.proto`. Import it from `request.proto` and `response.proto`. Define the shared event enum with an unspecified zero value and one value for each v1 event. + +Add a client capability: +- Add `bool supports_oz_lifecycle_hooks = 34` to `Request.Settings`. + +Add `repeated string supported_oz_hook_payload_schema_versions = 4` to `ResponseEvent.StreamInit`. A hook-enabled client must receive `warp.oz_hook.v1` in this field before it applies any action from the stream. + +Add `OzHookContext oz_hook_context = 7` to `Request`: + +```proto +message OzHookContext { + repeated OzHookEvent enabled_events = 1; + repeated string supported_payload_schema_versions = 2; +} +``` + +Add `OzHookResult oz_hook_result = 9` to `Request.Input.UserInputs.UserInput.input`. + +Do not send hook commands, matchers, failure modes, or trust records to MAA. The execution host owns configuration and matching. When an event name is enabled, MAA emits a gate for every server-owned occurrence of that event. The execution host returns continue without spawning a command when no local matcher selects a handler. + +Add `RunOzHook run_oz_hook = 15` to `ClientAction.action`: + +```proto +message RunOzHook { + string invocation_id = 1; + string tool_use_id = 2; + OzHookEvent event = 3; + string schema_version = 4; + google.protobuf.Struct redacted_payload = 5 [(sensitive) = true]; +} +``` + +Define the client-to-server input: + +```proto +message OzHookResult { + string invocation_id = 1; + string tool_use_id = 2; + oneof outcome { + Continue continue = 3; + Deny deny = 4; + Failed failed = 5; + Cancelled cancelled = 6; + } +} +``` + +`Deny` contains only a bounded reason and source identity. `Failed` contains a category and an explicit resolved action of continue or deny. It does not contain raw stdout or stderr. + +Use explicit enum values for event and outcome names in protobuf. Keep the command-facing PascalCase event name in JSON. + +Protocol rules: +- MAA assigns `invocation_id`. +- The client must echo it unchanged. +- MAA rejects missing, duplicate-with-different-content, stale, or mismatched results. +- Replaying an identical result is idempotent. +- A pending gate is scoped to the conversation, request, event, and tool-use ID. +- Pending state must survive the request boundary in server-owned task state. +- MAA must not execute or release the gated operation before a valid result. +- Cancellation clears pending gates. +- Old clients that do not advertise the capability never receive hook actions. +- A new server that receives an enabled hook context but cannot select a mutually supported payload schema finishes the stream before inference with an explicit incompatibility error. +- A hook-enabled client cancels a stream whose `StreamInit` omits its requested schema version. +- A hook-enabled runtime must reject a server that cannot provide required server-owned event coverage. It must not silently downgrade to client-only coverage. + +Apply redaction before constructing `RunOzHook`. The protocol payload is a source-neutral template because MAA does not know which user or project handlers will match. The execution host adds `hook_source` for each selected handler before it serializes command stdin. It reapplies the final payload limit after adding that field. A field marked sensitive prevents accidental logging, but it is not a substitute for redaction. + +### 11. Server-owned tool gate + +Wrap server-owned tool execution in this state machine: +1. Parse and validate the complete tool call. +2. Determine the canonical tool name and stable tool-use ID. +3. Apply existing server policy. +4. Check whether the request hook context enables `PreToolUse`. +5. If it is not enabled, execute normally. +6. If it is enabled, persist a pending gate and emit `RunOzHook(PreToolUse)`. +7. End the current response at a resumable boundary. +8. On the next request, validate `OzHookResult`. +9. On continue, execute the exact stored tool input. +10. On deny, do not execute. Persist a synthetic denied tool result for the model. +11. After execution reaches a terminal result, repeat the gate when the request context enables `PostToolUse`. +12. Resume inference only after `PostToolUse` continues or fails according to the resolved client outcome. + +Integrate the gate before a server tool implementation can produce side effects. Do not add it only after `nativeToolCallProcessor.ProduceActions`; that is too late for tool implementations that execute during action production. + +Store a hash of the original canonical tool input with the pending gate. Validate the hash before execution after resume. This proves the executed input is the input that the hook observed. + +Client-executed actions do not need the server round trip for pre/post execution. They use the local runtime stages in section 9. The server protocol is reserved for boundaries the execution host cannot otherwise intercept. + +### 12. `PreCompact` gate + +Compaction is server-owned. Add a resumable hook boundary immediately before `SummarizeMessagesForContextWindow`. + +Flow: +1. MAA decides compaction is required. +2. MAA builds a redacted `PreCompact` payload with `trigger`. +3. MAA persists a pending compaction gate. +4. MAA emits `RunOzHook(PreCompact)`. +5. The client or embedded cloud Oz runtime executes matching observational handlers. +6. The runtime returns continue or failed-open. +7. MAA validates the invocation and begins compaction. + +`PreCompact` cannot deny. A deny-shaped or malformed result is a failure and compaction continues after diagnostics. MAA must never wait indefinitely for this event. + +Keep the existing `SummarizationEventHandler.Start` callback for summarization output. The hook gate is an earlier control-plane boundary, not a replacement for current telemetry callbacks. + +### 13. Prompt, stop, and session events + +`UserPromptSubmit`: +- Run after attachment and skill resolution has produced the prompt representation. +- Redact attachments and secrets. +- Await observation before sending the prompt to MAA. +- Continue on every failure. + +`Stop`: +- Subscribe to the existing conversation status and final-output lifecycle. +- Emit once per user turn. +- Use a turn-generation token to prevent duplicate events from retries or repeated status notifications. +- Await observation before publishing the final terminal turn status. + +`SessionStart`: +- Emit once per runtime after configuration and trust resolution. +- Do not inject stdout into the model. + +`SessionEnd`: +- Emit from graceful shutdown and cancellation paths. +- Use a 1-second default and 3-second total teardown cap. +- Never extend worker or app shutdown beyond that cap. + +### 14. Cloud worker integration + +Do not add a second hook executor to `oz-agent-worker`. + +Required worker changes are limited to: +- carry hook capability and cloud trust material into `TaskParams` +- preserve the metadata through Direct, Docker, Kubernetes, and command dispatch +- ensure the embedded Oz process can read the task sandbox's home and project hook files +- keep hook processes inside the task cancellation context +- reject a hook-enabled task when a backend cannot preserve the embedded runtime contract +- add backend tests that prove worker-control-plane credentials are not inherited + +Direct: +- The embedded runtime runs under the task process and task workspace. +- It inherits only task environment into Oz. +- The Oz hook runner clears that environment again before hook spawn. + +Docker and Kubernetes: +- Hook commands run in the task container or pod. +- They use the task working directory. +- They never use Docker daemon or Kubernetes controller credentials unless those credentials are already intentionally present inside the task sandbox. + +Command backend: +- Bump `DispatchPayloadVersion` if trust/capability metadata changes the stable dispatch JSON. +- Include non-secret hook capability and trust identifiers. +- Do not include hook payloads or command outputs in dispatch metadata. + +### 15. Observability + +Define structured hook telemetry with allowlisted fields: +- event +- source +- definition hash +- matched +- duration bucket +- result +- exit code category +- failure category +- timeout +- truncation flags +- execution mode +- worker backend when available + +Never attach: +- command string +- payload JSON +- prompt +- tool input +- tool response +- stdout +- stderr +- denial reason +- environment + +User-facing diagnostics can include the configured command during trust review because that is required for informed approval. Routine execution logs should identify the source and hash instead of repeating the command. + +Add correlation fields to client and server tracing: +- conversation ID +- run ID +- hook invocation ID +- tool-use ID + +### 16. Feature gating and rollout + +Gate the behavior independently in Warp and MAA. + +Rollout order: +1. Land protocol definitions and generated bindings. +2. Deploy MAA support with the server flag off. +3. Land client and embedded cloud runtime support with the client flag off. +4. Land worker metadata propagation. +5. Enable internal local Oz runs. +6. Enable internal cloud Oz runs. +7. Verify denial, timeout, and redaction telemetry. +8. Expand availability. + +The server must handle capability negotiation throughout rollout. Third-party harness requests must never enter the Oz hook state machine. + +## End-to-end flows + +### Local client tool + +1. Local Oz loads host-local configuration. +2. The user submits a prompt. +3. Oz runs `UserPromptSubmit`. +4. MAA returns a client action. +5. Warp computes native permission classification. +6. Warp runs local `PreToolUse`. +7. Warp applies any existing permission prompt. +8. Warp executes the original action. +9. Warp runs local `PostToolUse`. +10. Warp returns the result to MAA. +11. The turn ends and Oz runs `Stop`. + +### Cloud client tool + +1. The worker starts embedded Oz in the task sandbox. +2. Embedded Oz loads sandbox-local configuration and trust. +3. The same client-tool flow runs inside the sandbox. +4. Hook commands run in the task workspace. +5. The worker daemon does not execute the commands. + +### Server-owned tool + +1. MAA receives a complete model tool call. +2. MAA persists a `PreToolUse` gate. +3. The execution host receives `RunOzHook`. +4. The host runs the configured command chain. +5. The host returns a correlated continue or deny. +6. MAA executes only on continue. +7. MAA persists and emits a `PostToolUse` gate after execution. +8. MAA resumes inference after the observational result. + +### Compaction + +1. MAA decides to compact. +2. MAA emits a resumable `PreCompact` gate. +3. The execution host runs matching hooks. +4. MAA receives a correlated continue or failure. +5. MAA starts summarization. + +## Decisions and trade-offs + +### Reuse event names, not full third-party semantics + +The command-facing event and field names follow the shared Claude Code and Codex shape. Oz rejects unsupported control fields instead of pretending they work. + +Trade-off: +- Existing scripts that only observe the compatible subset are portable. +- Scripts that grant, mutate, inject context, or rely on third-party tool aliases need an Oz adapter. + +### Add protocol gates + +The implementation adds resumable MAA state and extra request latency for server-owned events. + +Trade-off: +- More protocol and state-machine complexity. +- Complete coverage and reliable deny semantics. + +### Sequential handlers + +The runtime does not match Codex concurrency. + +Trade-off: +- Higher worst-case latency when many hooks match. +- Stable ordering, simple fail-closed semantics, and reproducible side effects. + +### Redacted structural payloads + +Tool payloads omit file contents, attachments, full transcripts, and raw output. + +Trade-off: +- Some third-party hooks cannot inspect every byte. +- The hook boundary does not become a general data-exfiltration channel. + +### Fail-open default + +Operational failures continue by default. + +Trade-off: +- Observability automation cannot break normal work. +- Policy authors must explicitly select `on_failure: "deny"` when availability of the policy is mandatory. + +## Assumptions + +- The execution host can persist or receive project trust material before a session starts. +- Existing secret-redaction utilities can supply known secret values to the hook redactor without exposing them in logs. +- MAA can persist pending hook gates in request/task state without changing the external conversation model. +- Every server-owned tool has a stable tool-use ID before side effects. +- The first release may emit only `trigger: "auto"` for `PreCompact`. + +## Out of scope + +- Full Claude Code or Codex configuration-file compatibility. +- Managed organization hooks. +- Hook-distributed credentials. +- Hook output persistence. +- Server-side execution of user command hooks. +- A remote hook service. +- Live trust prompts inside unattended cloud tasks. +- Tool-input mutation and allow decisions. +- Third-party harness regression fixes. + +## Validation criteria + +### Warp unit tests + +Run: + +```bash +cargo test -p warp oz_hooks_config +cargo test -p warp oz_hooks_trust +cargo test -p warp oz_hooks_ordering +cargo test -p warp oz_hooks_payload +cargo test -p warp oz_hooks_redaction +cargo test -p warp oz_hooks_runtime +cargo test -p warp oz_hooks_permissions +``` + +Required coverage: +- strict config parsing and limits +- user/project merge order +- exact-byte hash trust +- all matcher subjects +- all seven payload goldens +- every redaction and size limit +- sequential execution +- explicit deny +- exit-2 deny +- fail-open and fail-closed +- timeout and process-group kill +- cancellation and late-result rejection +- no mutation or allow output +- native permission composition + +### Warp integration tests + +Add local Oz integration tests that: +- emit the seven events in lifecycle order +- verify one event per documented boundary +- deny a shell, file, MCP, and orchestration tool before side effects +- verify denied tools produce no `PostToolUse` +- verify failed tools produce `PostToolUse` with terminal status +- prove third-party harness startup does not construct the Oz hook runtime + +Run the focused integration target, then run: + +```bash +./script/presubmit +``` + +### Protocol tests + +In `warp-proto-apis`, run: + +```bash +./script/generate -a multi_agent -v v1 +git diff --exit-code +``` + +Add compatibility tests for: +- old clients without the capability +- unknown future enum values +- duplicate identical results +- duplicate conflicting results +- stale and mismatched invocation IDs +- cancellation + +### MAA tests + +In `warp-server`, run: + +```bash +go test ./logic/ai/multi_agent/... +``` + +Add focused tests that: +- pause a server tool before side effects +- preserve and verify the original input hash +- synthesize a denied tool result +- wait for `PostToolUse` before the next inference +- pause before compaction +- continue compaction after observational hook failure +- resume idempotently after request replay +- clear pending gates on cancellation +- never gate third-party harnesses + +### Worker tests + +In `oz-agent-worker`, run: + +```bash +go test ./internal/worker/... +``` + +Add backend tests that: +- run a hook in the Direct task workspace +- run a hook in Docker and Kubernetes task workspaces +- preserve hook metadata through command dispatch +- reject an incompatible backend +- cancel hook subprocesses with the task +- prove worker API keys and control-plane credentials are absent from the hook environment + +### End-to-end acceptance + +Run one local Oz session and one cloud Oz session with a fixture hook set. Both must: +- emit all seven events +- produce payloads matching the same versioned golden schema +- preserve user-before-project order +- deny the same tool without side effects +- continue after the same observational failure +- deny after the same fail-closed `PreToolUse` failure +- record metadata-only diagnostics + +No visual or computer-use validation is required. From 0aa290876dde53b68a8741a0fb2bf044b0af2152 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 12:53:08 +0000 Subject: [PATCH 02/11] Add first-party Oz hook runtime core --- app/src/ai/agent_sdk/hooks/config.rs | 359 ++++++++++ app/src/ai/agent_sdk/hooks/config_tests.rs | 242 +++++++ app/src/ai/agent_sdk/hooks/mod.rs | 94 +++ app/src/ai/agent_sdk/hooks/mod_tests.rs | 17 + app/src/ai/agent_sdk/hooks/payload.rs | 207 ++++++ app/src/ai/agent_sdk/hooks/payload_tests.rs | 121 ++++ app/src/ai/agent_sdk/hooks/permissions.rs | 38 ++ .../ai/agent_sdk/hooks/permissions_tests.rs | 25 + app/src/ai/agent_sdk/hooks/redaction.rs | 204 ++++++ app/src/ai/agent_sdk/hooks/redaction_tests.rs | 40 ++ app/src/ai/agent_sdk/hooks/runtime.rs | 638 ++++++++++++++++++ app/src/ai/agent_sdk/hooks/runtime_tests.rs | 228 +++++++ app/src/ai/agent_sdk/hooks/trust.rs | 44 ++ app/src/ai/agent_sdk/mod.rs | 13 + crates/warp_features/src/lib.rs | 3 + 15 files changed, 2273 insertions(+) create mode 100644 app/src/ai/agent_sdk/hooks/config.rs create mode 100644 app/src/ai/agent_sdk/hooks/config_tests.rs create mode 100644 app/src/ai/agent_sdk/hooks/mod.rs create mode 100644 app/src/ai/agent_sdk/hooks/mod_tests.rs create mode 100644 app/src/ai/agent_sdk/hooks/payload.rs create mode 100644 app/src/ai/agent_sdk/hooks/payload_tests.rs create mode 100644 app/src/ai/agent_sdk/hooks/permissions.rs create mode 100644 app/src/ai/agent_sdk/hooks/permissions_tests.rs create mode 100644 app/src/ai/agent_sdk/hooks/redaction.rs create mode 100644 app/src/ai/agent_sdk/hooks/redaction_tests.rs create mode 100644 app/src/ai/agent_sdk/hooks/runtime.rs create mode 100644 app/src/ai/agent_sdk/hooks/runtime_tests.rs create mode 100644 app/src/ai/agent_sdk/hooks/trust.rs diff --git a/app/src/ai/agent_sdk/hooks/config.rs b/app/src/ai/agent_sdk/hooks/config.rs new file mode 100644 index 00000000000..6c66b315013 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/config.rs @@ -0,0 +1,359 @@ +use std::collections::BTreeMap; +use std::fs; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use regex::Regex; +use serde::Deserialize; +use sha2::{Digest as _, Sha256}; + +use super::trust::{HookTrustKey, HookTrustStore}; +use super::{ + CONFIG_SCHEMA_VERSION, FailureMode, HookConfigSource, HookEventName, MAX_CONFIG_BYTES, + MAX_HANDLERS_PER_FILE, +}; + +const USER_CONFIG_RELATIVE_PATH: &str = ".warp/hooks.json"; +const PROJECT_CONFIG_RELATIVE_PATH: &str = ".warp/hooks.json"; + +#[derive(Clone, Debug)] +pub(crate) struct ConfiguredHook { + pub(crate) event: HookEventName, + pub(crate) matcher_text: Option, + matcher: Option, + pub(crate) command: String, + pub(crate) command_windows: Option, + pub(crate) timeout: Duration, + pub(crate) on_failure: FailureMode, + pub(crate) source: HookConfigSource, + pub(crate) config_path: PathBuf, + pub(crate) definition_hash: String, +} + +impl ConfiguredHook { + pub(crate) fn matches(&self, subject: Option<&str>) -> bool { + if self.event.ignores_matcher() { + return true; + } + self.matcher + .as_ref() + .is_none_or(|matcher| subject.is_some_and(|subject| matcher.is_match(subject))) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) enum HookConfigDiagnosticKind { + Invalid, + UntrustedProject, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct HookConfigDiagnostic { + pub(crate) path: PathBuf, + pub(crate) kind: HookConfigDiagnosticKind, + pub(crate) definition_hash: Option, + pub(crate) message: String, +} + +#[derive(Clone, Debug, Default)] +pub(crate) struct HookConfigSnapshot { + handlers: Arc>, + pub(crate) diagnostics: Arc>, +} + +impl HookConfigSnapshot { + pub(crate) fn matching_handlers( + &self, + event: HookEventName, + subject: Option<&str>, + ) -> impl Iterator { + self.handlers + .iter() + .filter(move |hook| hook.event == event && hook.matches(subject)) + } + + pub(crate) fn enabled_events(&self) -> impl Iterator + '_ { + HookEventName::ALL + .into_iter() + .filter(|event| self.handlers.iter().any(|handler| handler.event == *event)) + } + + #[cfg(test)] + pub(crate) fn handlers(&self) -> &[ConfiguredHook] { + &self.handlers + } +} + +pub(crate) fn discover_hook_config( + initial_cwd: &Path, + trust_store: &dyn HookTrustStore, +) -> HookConfigSnapshot { + let user_path = dirs::home_dir().map(|home| home.join(USER_CONFIG_RELATIVE_PATH)); + let project = git2::Repository::discover(initial_cwd) + .ok() + .and_then(|repository| repository.workdir().map(Path::to_path_buf)) + .and_then(|git_root| { + let canonical_git_root = fs::canonicalize(git_root).ok()?; + Some(ProjectConfig { + path: canonical_git_root.join(PROJECT_CONFIG_RELATIVE_PATH), + git_root: canonical_git_root, + }) + }); + load_hook_config(user_path.as_deref(), project.as_ref(), trust_store) +} + +#[derive(Clone, Debug)] +pub(crate) struct ProjectConfig { + pub(crate) path: PathBuf, + pub(crate) git_root: PathBuf, +} + +pub(crate) fn load_hook_config( + user_path: Option<&Path>, + project: Option<&ProjectConfig>, + trust_store: &dyn HookTrustStore, +) -> HookConfigSnapshot { + let mut handlers = Vec::new(); + let mut diagnostics = Vec::new(); + + if let Some(path) = user_path { + load_file( + path, + HookConfigSource::User, + None, + trust_store, + &mut handlers, + ) + .err() + .into_iter() + .for_each(|diagnostic| diagnostics.push(diagnostic)); + } + if let Some(project) = project { + load_file( + &project.path, + HookConfigSource::Project, + Some(&project.git_root), + trust_store, + &mut handlers, + ) + .err() + .into_iter() + .for_each(|diagnostic| diagnostics.push(diagnostic)); + } + + HookConfigSnapshot { + handlers: Arc::new(handlers), + diagnostics: Arc::new(diagnostics), + } +} + +fn load_file( + path: &Path, + source: HookConfigSource, + git_root: Option<&Path>, + trust_store: &dyn HookTrustStore, + handlers: &mut Vec, +) -> Result<(), HookConfigDiagnostic> { + let bytes = match fs::read(path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(error) => { + return Err(invalid_diagnostic( + path, + format!("failed to read file: {error}"), + )); + } + }; + if bytes.len() > MAX_CONFIG_BYTES { + return Err(invalid_diagnostic( + path, + format!("file exceeds {MAX_CONFIG_BYTES} bytes"), + )); + } + + let raw: RawHookConfig = serde_json::from_slice(&bytes) + .map_err(|error| invalid_diagnostic(path, format!("invalid JSON: {error}")))?; + let parsed = validate_config(raw, source, path, &bytes)?; + if source == HookConfigSource::Project { + let canonical_git_root = git_root + .and_then(|root| fs::canonicalize(root).ok()) + .ok_or_else(|| invalid_diagnostic(path, "project Git root is unavailable".into()))?; + let canonical_path = fs::canonicalize(path).map_err(|error| { + invalid_diagnostic(path, format!("failed to canonicalize path: {error}")) + })?; + let key = HookTrustKey { + git_root: canonical_git_root, + config_path: canonical_path, + definition_hash: parsed.definition_hash.clone(), + }; + if !trust_store.is_trusted(&key) { + return Err(HookConfigDiagnostic { + path: path.to_path_buf(), + kind: HookConfigDiagnosticKind::UntrustedProject, + definition_hash: Some(parsed.definition_hash), + message: "project hook definition is not trusted".into(), + }); + } + } + handlers.extend(parsed.handlers); + Ok(()) +} + +fn invalid_diagnostic(path: &Path, message: String) -> HookConfigDiagnostic { + HookConfigDiagnostic { + path: path.to_path_buf(), + kind: HookConfigDiagnosticKind::Invalid, + definition_hash: None, + message, + } +} + +struct ValidatedConfig { + handlers: Vec, + definition_hash: String, +} + +fn validate_config( + raw: RawHookConfig, + source: HookConfigSource, + path: &Path, + bytes: &[u8], +) -> Result { + if raw.schema_version != CONFIG_SCHEMA_VERSION { + return Err(invalid_diagnostic( + path, + format!("unsupported schema_version {:?}", raw.schema_version), + )); + } + let handler_count = raw + .hooks + .values() + .flat_map(|groups| groups.iter()) + .map(|group| group.hooks.len()) + .sum::(); + if handler_count > MAX_HANDLERS_PER_FILE { + return Err(invalid_diagnostic( + path, + format!("file exceeds {MAX_HANDLERS_PER_FILE} command handlers"), + )); + } + + let definition_hash = hex::encode(Sha256::digest(bytes)); + let mut handlers = Vec::with_capacity(handler_count); + for (event, groups) in raw.hooks { + for group in groups { + if group.hooks.is_empty() { + return Err(invalid_diagnostic( + path, + format!("{event} matcher group has no handlers"), + )); + } + let matcher_text = group + .matcher + .filter(|matcher| !matcher.is_empty() && matcher != "*"); + let matcher = matcher_text + .as_deref() + .map(Regex::new) + .transpose() + .map_err(|error| { + invalid_diagnostic(path, format!("invalid {event} matcher: {error}")) + })?; + for raw_handler in group.hooks { + let RawCommandHandler { + handler_type, + command, + command_windows, + timeout, + on_failure, + } = raw_handler; + match handler_type { + CommandHandlerType::Command => {} + } + if command.is_empty() { + return Err(invalid_diagnostic( + path, + format!("{event} command must be non-empty"), + )); + } + if on_failure == FailureMode::Deny && event != HookEventName::PreToolUse { + return Err(invalid_diagnostic( + path, + format!("on_failure deny is unsupported for {event}"), + )); + } + let timeout_seconds = timeout.unwrap_or_else(|| { + if event == HookEventName::SessionEnd { + 1 + } else { + 10 + } + }); + let maximum = if event == HookEventName::SessionEnd { + 3 + } else { + 120 + }; + if timeout_seconds == 0 || timeout_seconds > maximum { + return Err(invalid_diagnostic( + path, + format!("{event} timeout must be between 1 and {maximum} seconds"), + )); + } + handlers.push(ConfiguredHook { + event, + matcher_text: matcher_text.clone(), + matcher: matcher.clone(), + command, + command_windows, + timeout: Duration::from_secs(timeout_seconds), + on_failure, + source, + config_path: path.to_path_buf(), + definition_hash: definition_hash.clone(), + }); + } + } + } + + Ok(ValidatedConfig { + handlers, + definition_hash, + }) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RawHookConfig { + schema_version: String, + hooks: BTreeMap>, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RawMatcherGroup { + matcher: Option, + hooks: Vec, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RawCommandHandler { + #[serde(rename = "type")] + handler_type: CommandHandlerType, + command: String, + command_windows: Option, + timeout: Option, + #[serde(default)] + on_failure: FailureMode, +} + +#[derive(Deserialize)] +#[serde(rename_all = "lowercase")] +enum CommandHandlerType { + Command, +} + +#[cfg(test)] +#[path = "config_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/hooks/config_tests.rs b/app/src/ai/agent_sdk/hooks/config_tests.rs new file mode 100644 index 00000000000..bd3be12f076 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/config_tests.rs @@ -0,0 +1,242 @@ +use std::fs; + +use serde_json::json; + +use super::*; +use crate::ai::agent_sdk::hooks::trust::{DenyProjectHookTrust, ExactHookTrustStore}; + +fn write_config(directory: &Path, value: serde_json::Value) -> PathBuf { + let path = directory.join("hooks.json"); + fs::write(&path, serde_json::to_vec(&value).unwrap()).unwrap(); + path +} + +fn valid_config() -> serde_json::Value { + json!({ + "schema_version": CONFIG_SCHEMA_VERSION, + "hooks": { + "PreToolUse": [{ + "matcher": "^(run_shell_command|apply_patch)$", + "hooks": [{ + "type": "command", + "command": "check", + "timeout": 10, + "on_failure": "deny" + }] + }] + } + }) +} + +#[test] +fn oz_hooks_config_parses_strict_valid_file() { + let temp = tempfile::tempdir().unwrap(); + let path = write_config(temp.path(), valid_config()); + + let snapshot = load_hook_config(Some(&path), None, &DenyProjectHookTrust); + + assert_eq!(snapshot.handlers().len(), 1); + assert!(snapshot.diagnostics.is_empty()); + assert!(snapshot.handlers()[0].matches(Some("apply_patch"))); + assert!(!snapshot.handlers()[0].matches(Some("read_files"))); +} + +#[test] +fn oz_hooks_config_rejects_unknown_fields_events_and_schema_versions() { + for value in [ + json!({"schema_version": "future", "hooks": {}}), + json!({"schema_version": CONFIG_SCHEMA_VERSION, "hooks": {}, "unknown": true}), + json!({"schema_version": CONFIG_SCHEMA_VERSION, "hooks": {"Future": []}}), + ] { + let temp = tempfile::tempdir().unwrap(); + let path = write_config(temp.path(), value); + + let snapshot = load_hook_config(Some(&path), None, &DenyProjectHookTrust); + + assert!(snapshot.handlers().is_empty()); + assert_eq!(snapshot.diagnostics.len(), 1); + assert_eq!( + snapshot.diagnostics[0].kind, + HookConfigDiagnosticKind::Invalid + ); + } +} + +#[test] +fn oz_hooks_config_rejects_invalid_regex_timeout_and_failure_mode() { + let cases = [ + json!({ + "schema_version": CONFIG_SCHEMA_VERSION, + "hooks": {"PreToolUse": [{"matcher": "(", "hooks": [{"type": "command", "command": "x"}]}]} + }), + json!({ + "schema_version": CONFIG_SCHEMA_VERSION, + "hooks": {"PreToolUse": [{"hooks": [{"type": "command", "command": "x", "timeout": 121}]}]} + }), + json!({ + "schema_version": CONFIG_SCHEMA_VERSION, + "hooks": {"SessionEnd": [{"hooks": [{"type": "command", "command": "x", "timeout": 4}]}]} + }), + json!({ + "schema_version": CONFIG_SCHEMA_VERSION, + "hooks": {"Stop": [{"hooks": [{"type": "command", "command": "x", "on_failure": "deny"}]}]} + }), + ]; + for value in cases { + let temp = tempfile::tempdir().unwrap(); + let path = write_config(temp.path(), value); + + let snapshot = load_hook_config(Some(&path), None, &DenyProjectHookTrust); + + assert!(snapshot.handlers().is_empty()); + assert_eq!(snapshot.diagnostics.len(), 1); + } +} + +#[test] +fn oz_hooks_config_enforces_file_and_handler_limits() { + let temp = tempfile::tempdir().unwrap(); + let oversized = temp.path().join("oversized.json"); + fs::write(&oversized, vec![b' '; MAX_CONFIG_BYTES + 1]).unwrap(); + let handlers = (0..=MAX_HANDLERS_PER_FILE) + .map(|_| json!({"type": "command", "command": "x"})) + .collect::>(); + let too_many = write_config( + temp.path(), + json!({ + "schema_version": CONFIG_SCHEMA_VERSION, + "hooks": {"Stop": [{"hooks": handlers}]} + }), + ); + + let oversized_snapshot = load_hook_config(Some(&oversized), None, &DenyProjectHookTrust); + let handler_snapshot = load_hook_config(Some(&too_many), None, &DenyProjectHookTrust); + + assert_eq!(oversized_snapshot.diagnostics.len(), 1); + assert_eq!(handler_snapshot.diagnostics.len(), 1); +} + +#[test] +fn oz_hooks_ordering_preserves_user_then_project_declarations() { + let user_dir = tempfile::tempdir().unwrap(); + let project_dir = tempfile::tempdir().unwrap(); + let user_path = write_config( + user_dir.path(), + json!({ + "schema_version": CONFIG_SCHEMA_VERSION, + "hooks": {"Stop": [{"hooks": [ + {"type": "command", "command": "user-one"}, + {"type": "command", "command": "user-two"} + ]}]} + }), + ); + let project_path = write_config( + project_dir.path(), + json!({ + "schema_version": CONFIG_SCHEMA_VERSION, + "hooks": {"Stop": [{"hooks": [ + {"type": "command", "command": "project"} + ]}]} + }), + ); + let bytes = fs::read(&project_path).unwrap(); + let trust = ExactHookTrustStore::default(); + trust.trust(HookTrustKey { + git_root: fs::canonicalize(project_dir.path()).unwrap(), + config_path: fs::canonicalize(&project_path).unwrap(), + definition_hash: hex::encode(Sha256::digest(bytes)), + }); + + let snapshot = load_hook_config( + Some(&user_path), + Some(&ProjectConfig { + path: project_path, + git_root: project_dir.path().to_path_buf(), + }), + &trust, + ); + + assert_eq!( + snapshot + .handlers() + .iter() + .map(|handler| handler.command.as_str()) + .collect::>(), + ["user-one", "user-two", "project"] + ); +} + +#[test] +fn oz_hooks_trust_requires_exact_bytes_and_supports_revocation() { + let project_dir = tempfile::tempdir().unwrap(); + let path = write_config(project_dir.path(), valid_config()); + let project = ProjectConfig { + path: path.clone(), + git_root: project_dir.path().to_path_buf(), + }; + let trust = ExactHookTrustStore::default(); + let key = HookTrustKey { + git_root: fs::canonicalize(project_dir.path()).unwrap(), + config_path: fs::canonicalize(&path).unwrap(), + definition_hash: hex::encode(Sha256::digest(fs::read(&path).unwrap())), + }; + + assert!( + load_hook_config(None, Some(&project), &trust) + .handlers() + .is_empty() + ); + trust.trust(key.clone()); + assert_eq!( + load_hook_config(None, Some(&project), &trust) + .handlers() + .len(), + 1 + ); + fs::write(&path, [fs::read(&path).unwrap(), b"\n".to_vec()].concat()).unwrap(); + assert!( + load_hook_config(None, Some(&project), &trust) + .handlers() + .is_empty() + ); + fs::write(&path, serde_json::to_vec(&valid_config()).unwrap()).unwrap(); + trust.revoke(&key); + assert!( + load_hook_config(None, Some(&project), &trust) + .handlers() + .is_empty() + ); +} + +#[test] +fn oz_hooks_config_matcher_subject_rules_cover_all_events() { + let temp = tempfile::tempdir().unwrap(); + let hooks = HookEventName::ALL + .into_iter() + .map(|event| { + ( + event.as_str().to_owned(), + json!([{"matcher": "^wanted$", "hooks": [{"type": "command", "command": "x"}]}]), + ) + }) + .collect::>(); + let path = write_config( + temp.path(), + json!({"schema_version": CONFIG_SCHEMA_VERSION, "hooks": hooks}), + ); + let snapshot = load_hook_config(Some(&path), None, &DenyProjectHookTrust); + + for event in HookEventName::ALL { + let matched = snapshot.matching_handlers(event, Some("other")).count(); + if event.ignores_matcher() { + assert_eq!(matched, 1, "{event}"); + } else { + assert_eq!(matched, 0, "{event}"); + assert_eq!( + snapshot.matching_handlers(event, Some("wanted")).count(), + 1, + "{event}" + ); + } + } +} diff --git a/app/src/ai/agent_sdk/hooks/mod.rs b/app/src/ai/agent_sdk/hooks/mod.rs new file mode 100644 index 00000000000..b1beea7cc06 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/mod.rs @@ -0,0 +1,94 @@ +use std::fmt; + +use serde::{Deserialize, Serialize}; + +pub(crate) mod config; +pub(crate) mod payload; +pub(crate) mod permissions; +pub(crate) mod redaction; +pub(crate) mod runtime; +pub(crate) mod trust; + +pub(crate) const CONFIG_SCHEMA_VERSION: &str = "warp.oz_hooks.config.v1"; +pub(crate) const PAYLOAD_SCHEMA_VERSION: &str = "warp.oz_hook.v1"; +pub(crate) const MAX_CONFIG_BYTES: usize = 256 * 1024; +pub(crate) const MAX_HANDLERS_PER_FILE: usize = 64; +pub(crate) const MAX_PAYLOAD_BYTES: usize = 256 * 1024; +pub(crate) const MAX_PROMPT_BYTES: usize = 64 * 1024; +pub(crate) const MAX_TOOL_INPUT_BYTES: usize = 128 * 1024; +pub(crate) const MAX_TOOL_RESPONSE_BYTES: usize = 64 * 1024; +pub(crate) const MAX_OUTPUT_BYTES: usize = 64 * 1024; +pub(crate) const MAX_DENIAL_REASON_BYTES: usize = 4 * 1024; + +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +pub(crate) enum HookEventName { + SessionStart, + SessionEnd, + UserPromptSubmit, + Stop, + PreToolUse, + PostToolUse, + PreCompact, +} + +impl HookEventName { + pub(crate) const ALL: [Self; 7] = [ + Self::SessionStart, + Self::SessionEnd, + Self::UserPromptSubmit, + Self::Stop, + Self::PreToolUse, + Self::PostToolUse, + Self::PreCompact, + ]; + + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::SessionStart => "SessionStart", + Self::SessionEnd => "SessionEnd", + Self::UserPromptSubmit => "UserPromptSubmit", + Self::Stop => "Stop", + Self::PreToolUse => "PreToolUse", + Self::PostToolUse => "PostToolUse", + Self::PreCompact => "PreCompact", + } + } + + pub(crate) const fn ignores_matcher(self) -> bool { + matches!(self, Self::UserPromptSubmit | Self::Stop) + } +} + +impl fmt::Display for HookEventName { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum HookConfigSource { + User, + Project, +} + +impl HookConfigSource { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::User => "user", + Self::Project => "project", + } + } +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum FailureMode { + #[default] + Continue, + Deny, +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/hooks/mod_tests.rs b/app/src/ai/agent_sdk/hooks/mod_tests.rs new file mode 100644 index 00000000000..7d3cbc53109 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/mod_tests.rs @@ -0,0 +1,17 @@ +use super::HookEventName; + +#[test] +fn oz_hooks_config_event_names_are_stable() { + assert_eq!( + HookEventName::ALL.map(HookEventName::as_str), + [ + "SessionStart", + "SessionEnd", + "UserPromptSubmit", + "Stop", + "PreToolUse", + "PostToolUse", + "PreCompact", + ] + ); +} diff --git a/app/src/ai/agent_sdk/hooks/payload.rs b/app/src/ai/agent_sdk/hooks/payload.rs new file mode 100644 index 00000000000..393f4742ae2 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/payload.rs @@ -0,0 +1,207 @@ +use serde::Serialize; + +use super::redaction::{RedactedText, RedactedValue, TruncationMetadata}; +use super::{HookConfigSource, HookEventName, MAX_PAYLOAD_BYTES, PAYLOAD_SCHEMA_VERSION}; + +#[derive(Clone, Debug)] +pub(crate) struct HookPayloadContext { + pub(crate) session_id: String, + pub(crate) run_id: String, + pub(crate) conversation_id: String, + pub(crate) cwd: String, + pub(crate) model: String, + pub(crate) permission_mode: String, +} + +#[derive(Clone, Debug)] +pub(crate) struct HookPayloadTemplate { + pub(crate) context: HookPayloadContext, + pub(crate) event: HookEventFields, +} + +impl HookPayloadTemplate { + pub(crate) fn event_name(&self) -> HookEventName { + self.event.event_name() + } + + pub(crate) fn matcher_subject(&self) -> Option<&str> { + self.event.matcher_subject() + } + + pub(crate) fn serialize_for_source( + &self, + source: HookConfigSource, + ) -> Result, PayloadError> { + let payload = HookPayload { + schema_version: PAYLOAD_SCHEMA_VERSION, + hook_event_name: self.event.event_name(), + session_id: &self.context.session_id, + run_id: &self.context.run_id, + conversation_id: &self.context.conversation_id, + cwd: &self.context.cwd, + hook_source: source, + model: &self.context.model, + permission_mode: &self.context.permission_mode, + event: &self.event, + }; + let bytes = serde_json::to_vec(&payload).map_err(PayloadError::Serialize)?; + if bytes.len() > MAX_PAYLOAD_BYTES { + return Err(PayloadError::Oversized(bytes.len())); + } + Ok(bytes) + } +} + +#[derive(Clone, Debug, Serialize)] +#[serde(untagged)] +pub(crate) enum HookEventFields { + SessionStart { + source: SessionStartSource, + }, + SessionEnd { + reason: SessionEndReason, + }, + UserPromptSubmit { + prompt: String, + #[serde(skip_serializing_if = "Option::is_none")] + prompt_truncation: Option, + }, + Stop { + turn_status: TurnStatus, + }, + PreToolUse { + tool_name: String, + tool_use_id: String, + tool_input: RedactedValue, + }, + PostToolUse { + tool_name: String, + tool_use_id: String, + tool_input: RedactedValue, + tool_response: RedactedValue, + }, + PreCompact { + trigger: CompactTrigger, + }, +} + +impl HookEventFields { + pub(crate) fn user_prompt(prompt: RedactedText) -> Self { + Self::UserPromptSubmit { + prompt: prompt.value, + prompt_truncation: prompt.truncation, + } + } + + pub(crate) const fn event_name(&self) -> HookEventName { + match self { + Self::SessionStart { .. } => HookEventName::SessionStart, + Self::SessionEnd { .. } => HookEventName::SessionEnd, + Self::UserPromptSubmit { .. } => HookEventName::UserPromptSubmit, + Self::Stop { .. } => HookEventName::Stop, + Self::PreToolUse { .. } => HookEventName::PreToolUse, + Self::PostToolUse { .. } => HookEventName::PostToolUse, + Self::PreCompact { .. } => HookEventName::PreCompact, + } + } + + pub(crate) fn matcher_subject(&self) -> Option<&str> { + match self { + Self::SessionStart { source } => Some(source.as_str()), + Self::SessionEnd { reason } => Some(reason.as_str()), + Self::PreToolUse { tool_name, .. } | Self::PostToolUse { tool_name, .. } => { + Some(tool_name) + } + Self::PreCompact { trigger } => Some(trigger.as_str()), + Self::UserPromptSubmit { .. } | Self::Stop { .. } => None, + } + } +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum SessionStartSource { + Startup, + Resume, +} + +impl SessionStartSource { + const fn as_str(self) -> &'static str { + match self { + Self::Startup => "startup", + Self::Resume => "resume", + } + } +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum SessionEndReason { + Completed, + Failed, + Cancelled, + Shutdown, +} + +impl SessionEndReason { + const fn as_str(self) -> &'static str { + match self { + Self::Completed => "completed", + Self::Failed => "failed", + Self::Cancelled => "cancelled", + Self::Shutdown => "shutdown", + } + } +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum CompactTrigger { + Auto, + Manual, +} + +impl CompactTrigger { + const fn as_str(self) -> &'static str { + match self { + Self::Auto => "auto", + Self::Manual => "manual", + } + } +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum TurnStatus { + Idle, + Blocked, + Failed, + Completed, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum PayloadError { + #[error("failed to serialize hook payload: {0}")] + Serialize(serde_json::Error), + #[error("serialized hook payload is {0} bytes")] + Oversized(usize), +} + +#[derive(Serialize)] +struct HookPayload<'a> { + schema_version: &'static str, + hook_event_name: HookEventName, + session_id: &'a str, + run_id: &'a str, + conversation_id: &'a str, + cwd: &'a str, + hook_source: HookConfigSource, + model: &'a str, + permission_mode: &'a str, + #[serde(flatten)] + event: &'a HookEventFields, +} + +#[cfg(test)] +#[path = "payload_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/hooks/payload_tests.rs b/app/src/ai/agent_sdk/hooks/payload_tests.rs new file mode 100644 index 00000000000..3e08b46a405 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/payload_tests.rs @@ -0,0 +1,121 @@ +use serde_json::json; + +use super::*; +use crate::ai::agent_sdk::hooks::redaction::RedactedValue; + +fn context() -> HookPayloadContext { + HookPayloadContext { + session_id: "session".into(), + run_id: "run".into(), + conversation_id: "conversation".into(), + cwd: "/workspace/repo".into(), + model: "model".into(), + permission_mode: "supervised".into(), + } +} + +fn payload(event: HookEventFields) -> serde_json::Value { + serde_json::from_slice( + &HookPayloadTemplate { + context: context(), + event, + } + .serialize_for_source(HookConfigSource::User) + .unwrap(), + ) + .unwrap() +} + +#[test] +fn oz_hooks_payload_golden_covers_all_seven_events() { + let events = [ + ( + HookEventName::SessionStart, + HookEventFields::SessionStart { + source: SessionStartSource::Startup, + }, + ), + ( + HookEventName::SessionEnd, + HookEventFields::SessionEnd { + reason: SessionEndReason::Completed, + }, + ), + ( + HookEventName::UserPromptSubmit, + HookEventFields::user_prompt(RedactedText { + value: "prompt".into(), + truncation: None, + }), + ), + ( + HookEventName::Stop, + HookEventFields::Stop { + turn_status: TurnStatus::Completed, + }, + ), + ( + HookEventName::PreToolUse, + HookEventFields::PreToolUse { + tool_name: "run_shell_command".into(), + tool_use_id: "tool".into(), + tool_input: RedactedValue::object([("command", "pwd".into())]), + }, + ), + ( + HookEventName::PostToolUse, + HookEventFields::PostToolUse { + tool_name: "run_shell_command".into(), + tool_use_id: "tool".into(), + tool_input: RedactedValue::object([("command", "pwd".into())]), + tool_response: RedactedValue::object([("status", "succeeded".into())]), + }, + ), + ( + HookEventName::PreCompact, + HookEventFields::PreCompact { + trigger: CompactTrigger::Auto, + }, + ), + ]; + + for (event_name, fields) in events { + let value = payload(fields); + assert_eq!(value["schema_version"], PAYLOAD_SCHEMA_VERSION); + assert_eq!(value["hook_event_name"], event_name.as_str()); + assert_eq!(value["session_id"], "session"); + assert_eq!(value["run_id"], "run"); + assert_eq!(value["conversation_id"], "conversation"); + assert_eq!(value["cwd"], "/workspace/repo"); + assert_eq!(value["hook_source"], "user"); + assert_eq!(value["model"], "model"); + assert_eq!(value["permission_mode"], "supervised"); + } +} + +#[test] +fn oz_hooks_payload_matches_command_facing_shape() { + let value = payload(HookEventFields::PreToolUse { + tool_name: "apply_patch".into(), + tool_use_id: "tool-1".into(), + tool_input: RedactedValue::object([("path", "src/lib.rs".into())]), + }); + + assert_eq!( + value, + json!({ + "schema_version": "warp.oz_hook.v1", + "hook_event_name": "PreToolUse", + "session_id": "session", + "run_id": "run", + "conversation_id": "conversation", + "cwd": "/workspace/repo", + "hook_source": "user", + "model": "model", + "permission_mode": "supervised", + "tool_name": "apply_patch", + "tool_use_id": "tool-1", + "tool_input": {"path": "src/lib.rs"} + }) + ); +} diff --git a/app/src/ai/agent_sdk/hooks/permissions.rs b/app/src/ai/agent_sdk/hooks/permissions.rs new file mode 100644 index 00000000000..d16e469f116 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/permissions.rs @@ -0,0 +1,38 @@ +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum NativePermission { + Deny, + Allow, + Prompt, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum HookPermission { + Continue, + Deny, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ComposedPermission { + DeniedByWarp, + DeniedByHook, + Allow, + Prompt, +} + +pub(crate) fn compose_permission( + native: NativePermission, + hook: HookPermission, +) -> ComposedPermission { + match native { + NativePermission::Deny => ComposedPermission::DeniedByWarp, + NativePermission::Allow | NativePermission::Prompt if hook == HookPermission::Deny => { + ComposedPermission::DeniedByHook + } + NativePermission::Allow => ComposedPermission::Allow, + NativePermission::Prompt => ComposedPermission::Prompt, + } +} + +#[cfg(test)] +#[path = "permissions_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/hooks/permissions_tests.rs b/app/src/ai/agent_sdk/hooks/permissions_tests.rs new file mode 100644 index 00000000000..96f878ef1d1 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/permissions_tests.rs @@ -0,0 +1,25 @@ +use super::*; + +#[test] +fn oz_hooks_permissions_never_upgrade_a_native_denial() { + assert_eq!( + compose_permission(NativePermission::Deny, HookPermission::Continue), + ComposedPermission::DeniedByWarp + ); + assert_eq!( + compose_permission(NativePermission::Deny, HookPermission::Deny), + ComposedPermission::DeniedByWarp + ); +} + +#[test] +fn oz_hooks_permissions_preserve_prompt_after_hook_continuation() { + assert_eq!( + compose_permission(NativePermission::Prompt, HookPermission::Continue), + ComposedPermission::Prompt + ); + assert_eq!( + compose_permission(NativePermission::Prompt, HookPermission::Deny), + ComposedPermission::DeniedByHook + ); +} diff --git a/app/src/ai/agent_sdk/hooks/redaction.rs b/app/src/ai/agent_sdk/hooks/redaction.rs new file mode 100644 index 00000000000..1b5aad99987 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/redaction.rs @@ -0,0 +1,204 @@ +use std::collections::BTreeMap; + +use regex::Regex; +use serde::Serialize; +use serde_json::Value; + +#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +#[serde(untagged)] +pub(crate) enum RedactedValue { + Null, + Bool(bool), + Number(serde_json::Number), + String(String), + Array(Vec), + Object(BTreeMap), +} + +impl RedactedValue { + pub(crate) fn object( + fields: impl IntoIterator, RedactedValue)>, + ) -> Self { + Self::Object( + fields + .into_iter() + .map(|(key, value)| (key.into(), value)) + .collect(), + ) + } + + pub(crate) fn redacted(reason: &str, byte_count: usize) -> Self { + Self::object([ + ("redacted", Self::Bool(true)), + ("reason", Self::String(reason.into())), + ( + "byte_count", + Self::Number(serde_json::Number::from(byte_count)), + ), + ]) + } + + pub(crate) fn serialized_len(&self) -> usize { + serde_json::to_vec(self).map_or(usize::MAX, |bytes| bytes.len()) + } +} + +impl From<&str> for RedactedValue { + fn from(value: &str) -> Self { + Self::String(value.into()) + } +} + +impl From for RedactedValue { + fn from(value: String) -> Self { + Self::String(value) + } +} + +impl From for RedactedValue { + fn from(value: bool) -> Self { + Self::Bool(value) + } +} + +impl From for RedactedValue { + fn from(value: u64) -> Self { + Self::Number(value.into()) + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +pub(crate) struct TruncationMetadata { + pub(crate) truncated: bool, + pub(crate) original_bytes: usize, + pub(crate) included_bytes: usize, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct RedactedText { + pub(crate) value: String, + pub(crate) truncation: Option, +} + +#[derive(Clone)] +pub(crate) struct HookRedactor { + known_secrets: Vec, + credential_patterns: Vec, +} + +impl HookRedactor { + pub(crate) fn new(known_secrets: impl IntoIterator) -> Self { + Self { + known_secrets: known_secrets + .into_iter() + .filter(|secret| !secret.is_empty()) + .collect(), + credential_patterns: vec![ + Regex::new(r"(?i)(authorization\s*:\s*(?:bearer|basic)\s+)\S+").unwrap(), + Regex::new( + r#"(?i)((?:api[_-]?key|access[_-]?token|secret)\s*[=:]\s*["']?)[^\s"',;]+"#, + ) + .unwrap(), + Regex::new(r"\b(?:gh[opusr]_[A-Za-z0-9_]{20,}|sk-[A-Za-z0-9_-]{20,})\b").unwrap(), + ], + } + } + + pub(crate) fn redact_text(&self, input: &str, maximum_bytes: usize) -> RedactedText { + let mut redacted = input.to_owned(); + for secret in &self.known_secrets { + redacted = redacted.replace(secret, "[REDACTED]"); + } + for pattern in &self.credential_patterns { + redacted = pattern + .replace_all(&redacted, |captures: ®ex::Captures<'_>| { + if captures.len() > 1 { + format!("{}[REDACTED]", captures.get(1).unwrap().as_str()) + } else { + "[REDACTED]".into() + } + }) + .into_owned(); + } + crate::ai::agent::redaction::redact_secrets(&mut redacted); + let original_bytes = redacted.len(); + if original_bytes <= maximum_bytes { + return RedactedText { + value: redacted, + truncation: None, + }; + } + let included_bytes = floor_utf8_boundary(&redacted, maximum_bytes); + redacted.truncate(included_bytes); + RedactedText { + value: redacted, + truncation: Some(TruncationMetadata { + truncated: true, + original_bytes, + included_bytes, + }), + } + } + + pub(crate) fn redact_json_preview( + &self, + value: &RedactedValue, + maximum_bytes: usize, + ) -> RedactedValue { + if value.serialized_len() <= maximum_bytes { + return value.clone(); + } + RedactedValue::object([ + ("truncated", RedactedValue::Bool(true)), + ( + "original_bytes", + RedactedValue::Number(value.serialized_len().into()), + ), + ( + "preview", + RedactedValue::redacted("size_limit", value.serialized_len()), + ), + ]) + } +} + +pub(crate) fn truncate_utf8(input: &str, maximum_bytes: usize) -> String { + let boundary = floor_utf8_boundary(input, maximum_bytes); + input[..boundary].to_owned() +} + +fn floor_utf8_boundary(input: &str, maximum_bytes: usize) -> usize { + let mut boundary = maximum_bytes.min(input.len()); + while !input.is_char_boundary(boundary) { + boundary -= 1; + } + boundary +} + +pub(crate) fn contains_prohibited_payload_key(value: &Value) -> bool { + const PROHIBITED_KEYS: [&str; 9] = [ + "environment", + "env", + "authorization", + "api_key", + "secret", + "attachment_bytes", + "file_content", + "transcript", + "transcript_path", + ]; + match value { + Value::Object(object) => object.iter().any(|(key, value)| { + PROHIBITED_KEYS + .iter() + .any(|prohibited| key.eq_ignore_ascii_case(prohibited)) + || contains_prohibited_payload_key(value) + }), + Value::Array(values) => values.iter().any(contains_prohibited_payload_key), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => false, + } +} + +#[cfg(test)] +#[path = "redaction_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/hooks/redaction_tests.rs b/app/src/ai/agent_sdk/hooks/redaction_tests.rs new file mode 100644 index 00000000000..2a2aa49673a --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/redaction_tests.rs @@ -0,0 +1,40 @@ +use super::*; + +#[test] +fn oz_hooks_redaction_removes_known_secrets_and_credentials_before_truncation() { + let redactor = HookRedactor::new(["known-secret".to_string()]); + let input = "known-secret Authorization: Bearer top-secret api_key=another-secret"; + + let output = redactor.redact_text(input, 128); + + assert!(!output.value.contains("known-secret")); + assert!(!output.value.contains("top-secret")); + assert!(!output.value.contains("another-secret")); +} + +#[test] +fn oz_hooks_redaction_truncates_only_at_utf8_boundaries_with_metadata() { + let output = HookRedactor::new([]).redact_text("abc😀def", 5); + + assert_eq!(output.value, "abc"); + assert_eq!( + output.truncation, + Some(TruncationMetadata { + truncated: true, + original_bytes: 10, + included_bytes: 3, + }) + ); +} + +#[test] +fn oz_hooks_redaction_represents_omitted_file_content_structurally() { + let value = + RedactedValue::object([("content", RedactedValue::redacted("file_content", 18_432))]); + let serialized = serde_json::to_value(value).unwrap(); + + assert_eq!(serialized["content"]["redacted"], true); + assert_eq!(serialized["content"]["reason"], "file_content"); + assert_eq!(serialized["content"]["byte_count"], 18_432); + assert!(!contains_prohibited_payload_key(&serialized)); +} diff --git a/app/src/ai/agent_sdk/hooks/runtime.rs b/app/src/ai/agent_sdk/hooks/runtime.rs new file mode 100644 index 00000000000..cd79621c048 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/runtime.rs @@ -0,0 +1,638 @@ +use std::collections::HashMap; +use std::ffi::OsString; +use std::path::Path; +use std::process::Stdio; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant, SystemTime}; + +use async_trait::async_trait; +use serde::Deserialize; +use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWriteExt as _}; +use tokio::process::Command; +use tokio::sync::{Mutex as AsyncMutex, mpsc}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +use super::config::{ConfiguredHook, HookConfigSnapshot}; +use super::payload::HookPayloadTemplate; +use super::redaction::truncate_utf8; +use super::{ + FailureMode, HookConfigSource, HookEventName, MAX_DENIAL_REASON_BYTES, MAX_OUTPUT_BYTES, +}; + +const SESSION_END_TOTAL_TIMEOUT: Duration = Duration::from_secs(3); + +#[derive(Clone, Debug)] +pub(crate) struct OzHookEvent { + pub(crate) invocation_id: String, + pub(crate) tool_use_id: Option, + pub(crate) payload: HookPayloadTemplate, +} + +#[derive(Clone, Debug)] +pub(crate) struct OzPreToolUseEvent(OzHookEvent); + +impl OzPreToolUseEvent { + pub(crate) fn new(event: OzHookEvent) -> Result { + if event.payload.event_name() != HookEventName::PreToolUse { + return Err(HookRuntimeError::WrongEvent); + } + Ok(Self(event)) + } +} + +#[derive(Clone, Debug)] +pub(crate) enum OzHookCancellationScope { + Session, + Invocation(String), + Tool(String), +} + +#[derive(Clone, Debug, Default)] +pub(crate) struct OzHookObservation { + pub(crate) diagnostics: Vec, +} + +#[derive(Clone, Debug)] +pub(crate) enum OzPreToolUseDecision { + Continue { + diagnostics: Vec, + }, + Deny { + reason: String, + source: HookConfigSource, + diagnostics: Vec, + }, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum HookInvocationResult { + Succeeded, + Continued, + Denied, + Failed, + TimedOut, + Cancelled, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum HookFailureCategory { + Spawn, + Stdin, + Timeout, + Cancelled, + OutputOverflow, + InvalidUtf8, + NonZeroExit, + InvalidDecision, + Payload, +} + +#[derive(Clone, Debug)] +pub(crate) struct HookInvocationDiagnostic { + pub(crate) event: HookEventName, + pub(crate) source: HookConfigSource, + pub(crate) config_path: std::path::PathBuf, + pub(crate) definition_hash: String, + pub(crate) matcher: Option, + pub(crate) started_at: SystemTime, + pub(crate) finished_at: SystemTime, + pub(crate) duration: Duration, + pub(crate) result: HookInvocationResult, + pub(crate) exit_code: Option, + pub(crate) output_truncated: bool, + pub(crate) failure_category: Option, +} + +#[async_trait] +pub(crate) trait OzHookRuntime: Send + Sync { + async fn observe(&self, event: OzHookEvent) -> OzHookObservation; + async fn pre_tool_use(&self, event: OzPreToolUseEvent) -> OzPreToolUseDecision; + fn cancel(&self, scope: OzHookCancellationScope); +} + +pub(crate) struct OzHookRuntimeService { + config: HookConfigSnapshot, + queue: AsyncMutex<()>, + cancellation: Mutex, +} + +struct RuntimeCancellation { + session: CancellationToken, + invocations: HashMap, +} + +struct PendingInvocation { + tool_use_id: Option, + token: CancellationToken, +} + +impl OzHookRuntimeService { + pub(crate) fn new(config: HookConfigSnapshot) -> Self { + Self { + config, + queue: AsyncMutex::new(()), + cancellation: Mutex::new(RuntimeCancellation { + session: CancellationToken::new(), + invocations: HashMap::new(), + }), + } + } + + async fn run_event(&self, event: OzHookEvent, pre_tool: bool) -> EventOutcome { + let token = { + let mut cancellation = self.cancellation.lock().unwrap(); + let token = cancellation.session.child_token(); + cancellation.invocations.insert( + event.invocation_id.clone(), + PendingInvocation { + tool_use_id: event.tool_use_id.clone(), + token: token.clone(), + }, + ); + token + }; + let _queue = tokio::select! { + guard = self.queue.lock() => guard, + () = token.cancelled() => { + self.remove_pending(&event.invocation_id); + return EventOutcome::default(); + } + }; + + let event_name = event.payload.event_name(); + let session_end_deadline = (event_name == HookEventName::SessionEnd) + .then(|| Instant::now() + SESSION_END_TOTAL_TIMEOUT); + let mut outcome = EventOutcome::default(); + for handler in self + .config + .matching_handlers(event_name, event.payload.matcher_subject()) + { + if token.is_cancelled() { + break; + } + let Some(timeout) = effective_timeout(handler.timeout, session_end_deadline) else { + break; + }; + let started_at = SystemTime::now(); + let started = Instant::now(); + let result = match event.payload.serialize_for_source(handler.source) { + Ok(payload) => { + run_command(handler, &event.payload, &payload, timeout, token.clone()).await + } + Err(_) => Err(CommandFailure { + category: HookFailureCategory::Payload, + exit_code: None, + }), + }; + let mut diagnostic = HookInvocationDiagnostic { + event: event_name, + source: handler.source, + config_path: handler.config_path.clone(), + definition_hash: handler.definition_hash.clone(), + matcher: handler.matcher_text.clone(), + started_at, + finished_at: SystemTime::now(), + duration: started.elapsed(), + result: HookInvocationResult::Succeeded, + exit_code: None, + output_truncated: false, + failure_category: None, + }; + match result { + Ok(CommandOutcome::Continue { exit_code }) => { + diagnostic.exit_code = exit_code; + } + Ok(CommandOutcome::Deny { reason, exit_code }) if pre_tool => { + diagnostic.result = HookInvocationResult::Denied; + diagnostic.exit_code = exit_code; + outcome.diagnostics.push(diagnostic); + outcome.denial = Some((reason, handler.source)); + break; + } + Ok(CommandOutcome::Deny { exit_code, .. }) => { + diagnostic.result = HookInvocationResult::Failed; + diagnostic.exit_code = exit_code; + diagnostic.failure_category = Some(HookFailureCategory::InvalidDecision); + } + Err(failure) => { + diagnostic.exit_code = failure.exit_code; + diagnostic.failure_category = Some(failure.category); + diagnostic.result = match failure.category { + HookFailureCategory::Timeout => HookInvocationResult::TimedOut, + HookFailureCategory::Cancelled => HookInvocationResult::Cancelled, + HookFailureCategory::Spawn + | HookFailureCategory::Stdin + | HookFailureCategory::OutputOverflow + | HookFailureCategory::InvalidUtf8 + | HookFailureCategory::NonZeroExit + | HookFailureCategory::InvalidDecision + | HookFailureCategory::Payload => { + if pre_tool && handler.on_failure == FailureMode::Deny { + HookInvocationResult::Denied + } else { + HookInvocationResult::Continued + } + } + }; + if pre_tool + && handler.on_failure == FailureMode::Deny + && failure.category != HookFailureCategory::Cancelled + { + outcome.diagnostics.push(diagnostic); + outcome.denial = Some(( + "An Oz hook failed closed and denied this tool.".into(), + handler.source, + )); + break; + } + } + } + outcome.diagnostics.push(diagnostic); + } + self.remove_pending(&event.invocation_id); + outcome + } + + fn remove_pending(&self, invocation_id: &str) { + self.cancellation + .lock() + .unwrap() + .invocations + .remove(invocation_id); + } +} + +#[async_trait] +impl OzHookRuntime for OzHookRuntimeService { + async fn observe(&self, event: OzHookEvent) -> OzHookObservation { + OzHookObservation { + diagnostics: self.run_event(event, false).await.diagnostics, + } + } + + async fn pre_tool_use(&self, event: OzPreToolUseEvent) -> OzPreToolUseDecision { + let outcome = self.run_event(event.0, true).await; + match outcome.denial { + Some((reason, source)) => OzPreToolUseDecision::Deny { + reason, + source, + diagnostics: outcome.diagnostics, + }, + None => OzPreToolUseDecision::Continue { + diagnostics: outcome.diagnostics, + }, + } + } + + fn cancel(&self, scope: OzHookCancellationScope) { + let cancellation = self.cancellation.lock().unwrap(); + match scope { + OzHookCancellationScope::Session => cancellation.session.cancel(), + OzHookCancellationScope::Invocation(invocation_id) => { + if let Some(invocation) = cancellation.invocations.get(&invocation_id) { + invocation.token.cancel(); + } + } + OzHookCancellationScope::Tool(tool_use_id) => { + for invocation in cancellation.invocations.values() { + if invocation.tool_use_id.as_deref() == Some(&tool_use_id) { + invocation.token.cancel(); + } + } + } + } + } +} + +#[derive(Default)] +struct EventOutcome { + diagnostics: Vec, + denial: Option<(String, HookConfigSource)>, +} + +fn effective_timeout( + configured: Duration, + session_end_deadline: Option, +) -> Option { + let Some(deadline) = session_end_deadline else { + return Some(configured); + }; + deadline + .checked_duration_since(Instant::now()) + .map(|remaining| remaining.min(configured)) + .filter(|remaining| !remaining.is_zero()) +} + +enum CommandOutcome { + Continue { + exit_code: Option, + }, + Deny { + reason: String, + exit_code: Option, + }, +} + +struct CommandFailure { + category: HookFailureCategory, + exit_code: Option, +} + +async fn run_command( + handler: &ConfiguredHook, + payload: &HookPayloadTemplate, + stdin_payload: &[u8], + timeout: Duration, + cancellation: CancellationToken, +) -> Result { + let mut command = hook_command(handler); + command + .current_dir(Path::new(&payload.context.cwd)) + .env_clear() + .envs(hook_environment(payload)) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + configure_process_group(&mut command); + let mut child = command.spawn().map_err(|_| CommandFailure { + category: HookFailureCategory::Spawn, + exit_code: None, + })?; + let process_id = child.id(); + let mut stdin = child.stdin.take().unwrap(); + let stdin_payload = stdin_payload.to_vec(); + let stdin_task = tokio::spawn(async move { + stdin.write_all(&stdin_payload).await?; + stdin.shutdown().await + }); + let (overflow_tx, mut overflow_rx) = mpsc::unbounded_channel(); + let stdout_task = spawn_bounded_reader(child.stdout.take().unwrap(), overflow_tx.clone()); + let stderr_task = spawn_bounded_reader(child.stderr.take().unwrap(), overflow_tx); + + enum Completion { + Exited(std::io::Result), + Timeout, + Cancelled, + Overflow, + } + let completion = tokio::select! { + status = child.wait() => Completion::Exited(status), + () = tokio::time::sleep(timeout) => Completion::Timeout, + () = cancellation.cancelled() => Completion::Cancelled, + Some(()) = overflow_rx.recv() => Completion::Overflow, + }; + let status = match completion { + Completion::Exited(status) => status.map_err(|_| CommandFailure { + category: HookFailureCategory::NonZeroExit, + exit_code: None, + })?, + Completion::Timeout => { + kill_process_tree(process_id, &mut child).await; + return Err(CommandFailure { + category: HookFailureCategory::Timeout, + exit_code: None, + }); + } + Completion::Cancelled => { + kill_process_tree(process_id, &mut child).await; + return Err(CommandFailure { + category: HookFailureCategory::Cancelled, + exit_code: None, + }); + } + Completion::Overflow => { + kill_process_tree(process_id, &mut child).await; + return Err(CommandFailure { + category: HookFailureCategory::OutputOverflow, + exit_code: None, + }); + } + }; + if stdin_task.await.is_err() { + return Err(CommandFailure { + category: HookFailureCategory::Stdin, + exit_code: status.code(), + }); + } + let stdout = join_output(stdout_task, status.code()).await?; + let stderr = join_output(stderr_task, status.code()).await?; + parse_command_result(payload.event_name(), status.code(), stdout, stderr) +} + +fn hook_command(handler: &ConfiguredHook) -> Command { + #[cfg(windows)] + { + let shell = std::env::var_os("COMSPEC").unwrap_or_else(|| OsString::from("cmd.exe")); + let selected = handler + .command_windows + .as_deref() + .unwrap_or(&handler.command); + let mut command = Command::new(shell); + command.arg("/C").arg(selected); + command + } + #[cfg(not(windows))] + { + let shell = std::env::var_os("SHELL").unwrap_or_else(|| OsString::from("/bin/sh")); + let mut command = Command::new(shell); + command.arg("-c").arg(&handler.command); + command + } +} + +fn hook_environment(payload: &HookPayloadTemplate) -> HashMap { + #[cfg(windows)] + const ALLOWED: &[&str] = &[ + "USERPROFILE", + "PATH", + "COMSPEC", + "SystemRoot", + "TMP", + "TEMP", + "LANG", + "LC_ALL", + "LC_CTYPE", + ]; + #[cfg(not(windows))] + const ALLOWED: &[&str] = &[ + "HOME", "PATH", "SHELL", "TMPDIR", "TMP", "TEMP", "LANG", "LC_ALL", "LC_CTYPE", + ]; + + let mut environment = ALLOWED + .iter() + .filter_map(|key| std::env::var_os(key).map(|value| (OsString::from(key), value))) + .collect::>(); + environment.insert( + "WARP_HOOK_EVENT_NAME".into(), + payload.event_name().as_str().into(), + ); + environment.insert("WARP_RUN_ID".into(), payload.context.run_id.clone().into()); + environment.insert( + "WARP_CONVERSATION_ID".into(), + payload.context.conversation_id.clone().into(), + ); + environment +} + +fn spawn_bounded_reader( + mut reader: impl AsyncRead + Unpin + Send + 'static, + overflow: mpsc::UnboundedSender<()>, +) -> JoinHandle, ()>> { + tokio::spawn(async move { + let mut output = Vec::new(); + let mut buffer = [0_u8; 8192]; + loop { + let read = reader.read(&mut buffer).await.map_err(|_| ())?; + if read == 0 { + return Ok(output); + } + if output.len() + read > MAX_OUTPUT_BYTES { + let _ = overflow.send(()); + return Err(()); + } + output.extend_from_slice(&buffer[..read]); + } + }) +} + +async fn join_output( + task: JoinHandle, ()>>, + exit_code: Option, +) -> Result { + let bytes = task + .await + .map_err(|_| CommandFailure { + category: HookFailureCategory::OutputOverflow, + exit_code, + })? + .map_err(|_| CommandFailure { + category: HookFailureCategory::OutputOverflow, + exit_code, + })?; + String::from_utf8(bytes).map_err(|_| CommandFailure { + category: HookFailureCategory::InvalidUtf8, + exit_code, + }) +} + +fn parse_command_result( + event: HookEventName, + exit_code: Option, + stdout: String, + stderr: String, +) -> Result { + match (event, exit_code) { + (HookEventName::PreToolUse, Some(2)) if !stderr.is_empty() => Ok(CommandOutcome::Deny { + reason: truncate_utf8(&stderr, MAX_DENIAL_REASON_BYTES), + exit_code, + }), + (HookEventName::PreToolUse, Some(2)) => Err(CommandFailure { + category: HookFailureCategory::InvalidDecision, + exit_code, + }), + (HookEventName::PreToolUse, Some(0)) => parse_pre_tool_stdout(&stdout, exit_code), + (_, Some(0)) => Ok(CommandOutcome::Continue { exit_code }), + _ => Err(CommandFailure { + category: HookFailureCategory::NonZeroExit, + exit_code, + }), + } +} + +fn parse_pre_tool_stdout( + stdout: &str, + exit_code: Option, +) -> Result { + if stdout.trim().is_empty() { + return Ok(CommandOutcome::Continue { exit_code }); + } + let output: PreToolOutput = serde_json::from_str(stdout).map_err(|_| CommandFailure { + category: HookFailureCategory::InvalidDecision, + exit_code, + })?; + let Some(specific) = output.hook_specific_output else { + return Ok(CommandOutcome::Continue { exit_code }); + }; + if specific.hook_event_name != HookEventName::PreToolUse + || specific.permission_decision != PermissionDecision::Deny + || specific.permission_decision_reason.is_empty() + { + return Err(CommandFailure { + category: HookFailureCategory::InvalidDecision, + exit_code, + }); + } + Ok(CommandOutcome::Deny { + reason: truncate_utf8( + &specific.permission_decision_reason, + MAX_DENIAL_REASON_BYTES, + ), + exit_code, + }) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct PreToolOutput { + #[serde(rename = "hookSpecificOutput")] + hook_specific_output: Option, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct PreToolSpecificOutput { + #[serde(rename = "hookEventName")] + hook_event_name: HookEventName, + #[serde(rename = "permissionDecision")] + permission_decision: PermissionDecision, + #[serde(rename = "permissionDecisionReason")] + permission_decision_reason: String, +} + +#[derive(Deserialize, Eq, PartialEq)] +#[serde(rename_all = "lowercase")] +enum PermissionDecision { + Deny, +} + +#[cfg(unix)] +fn configure_process_group(command: &mut Command) { + use std::os::unix::process::CommandExt as _; + command.as_std_mut().process_group(0); +} + +#[cfg(windows)] +fn configure_process_group(command: &mut Command) { + use std::os::windows::process::CommandExt as _; + command.creation_flags(windows::Win32::System::Threading::CREATE_NEW_PROCESS_GROUP.0); +} + +#[cfg(unix)] +async fn kill_process_tree(process_id: Option, child: &mut tokio::process::Child) { + if let Some(process_id) = process_id { + let _ = nix::sys::signal::killpg( + nix::unistd::Pid::from_raw(process_id as i32), + nix::sys::signal::Signal::SIGKILL, + ); + } + let _ = child.start_kill(); + let _ = child.wait().await; +} + +#[cfg(windows)] +async fn kill_process_tree(_process_id: Option, child: &mut tokio::process::Child) { + let _ = child.start_kill(); + let _ = child.wait().await; +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum HookRuntimeError { + #[error("expected a PreToolUse event")] + WrongEvent, +} + +#[cfg(test)] +#[path = "runtime_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/hooks/runtime_tests.rs b/app/src/ai/agent_sdk/hooks/runtime_tests.rs new file mode 100644 index 00000000000..ad873cac553 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/runtime_tests.rs @@ -0,0 +1,228 @@ +use std::fs; + +use serde_json::json; + +use super::*; +use crate::ai::agent_sdk::hooks::config::load_hook_config; +use crate::ai::agent_sdk::hooks::payload::{ + HookEventFields, HookPayloadContext, SessionStartSource, +}; +use crate::ai::agent_sdk::hooks::trust::DenyProjectHookTrust; + +fn runtime_with_hooks(hooks: serde_json::Value) -> OzHookRuntimeService { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("hooks.json"); + fs::write( + &path, + serde_json::to_vec(&json!({ + "schema_version": "warp.oz_hooks.config.v1", + "hooks": hooks + })) + .unwrap(), + ) + .unwrap(); + let snapshot = load_hook_config(Some(&path), None, &DenyProjectHookTrust); + OzHookRuntimeService::new(snapshot) +} + +fn event(invocation_id: &str, fields: HookEventFields) -> OzHookEvent { + OzHookEvent { + invocation_id: invocation_id.into(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: HookPayloadContext { + session_id: "session".into(), + run_id: "run".into(), + conversation_id: "conversation".into(), + cwd: std::env::current_dir() + .unwrap() + .to_string_lossy() + .into_owned(), + model: "model".into(), + permission_mode: "supervised".into(), + }, + event: fields, + }, + } +} + +fn pre_tool_event(invocation_id: &str) -> OzPreToolUseEvent { + OzPreToolUseEvent::new(event( + invocation_id, + HookEventFields::PreToolUse { + tool_name: "run_shell_command".into(), + tool_use_id: "tool".into(), + tool_input: super::super::redaction::RedactedValue::object([] as [(&str, _); 0]), + }, + )) + .unwrap() +} + +#[tokio::test] +async fn oz_hooks_runtime_runs_matching_handlers_sequentially() { + let output = tempfile::NamedTempFile::new().unwrap(); + let path = output.path().to_string_lossy(); + let runtime = runtime_with_hooks(json!({ + "SessionStart": [{"hooks": [ + {"type": "command", "command": format!("printf first >> '{path}'")}, + {"type": "command", "command": format!("printf second >> '{path}'")} + ]}] + })); + + let observation = runtime + .observe(event( + "one", + HookEventFields::SessionStart { + source: SessionStartSource::Startup, + }, + )) + .await; + + assert_eq!(fs::read_to_string(output.path()).unwrap(), "firstsecond"); + assert_eq!(observation.diagnostics.len(), 2); +} + +#[tokio::test] +async fn oz_hooks_runtime_structured_deny_short_circuits_later_handlers() { + let marker = tempfile::NamedTempFile::new().unwrap(); + let marker_path = marker.path().to_string_lossy(); + let deny = r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"deny","permissionDecisionReason":"policy"}}"#; + let runtime = runtime_with_hooks(json!({ + "PreToolUse": [{"hooks": [ + {"type": "command", "command": format!("printf '%s' '{deny}'")}, + {"type": "command", "command": format!("printf ran >> '{marker_path}'")} + ]}] + })); + + let decision = runtime.pre_tool_use(pre_tool_event("deny")).await; + + assert!(matches!( + decision, + OzPreToolUseDecision::Deny { ref reason, .. } if reason == "policy" + )); + assert_eq!(fs::read_to_string(marker.path()).unwrap(), ""); +} + +#[tokio::test] +async fn oz_hooks_runtime_exit_two_denies_only_pre_tool_use() { + let runtime = runtime_with_hooks(json!({ + "PreToolUse": [{"hooks": [ + {"type": "command", "command": "printf policy >&2; exit 2"} + ]}], + "Stop": [{"hooks": [ + {"type": "command", "command": "printf ignored >&2; exit 2"} + ]}] + })); + + assert!(matches!( + runtime.pre_tool_use(pre_tool_event("pre")).await, + OzPreToolUseDecision::Deny { ref reason, .. } if reason == "policy" + )); + let observation = runtime + .observe(event( + "stop", + HookEventFields::Stop { + turn_status: super::super::payload::TurnStatus::Completed, + }, + )) + .await; + assert_eq!( + observation.diagnostics[0].failure_category, + Some(HookFailureCategory::NonZeroExit) + ); +} + +#[tokio::test] +async fn oz_hooks_runtime_invalid_allow_output_fails_open_or_closed() { + let allow = r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow","permissionDecisionReason":"no"}}"#; + for (mode, denied) in [("continue", false), ("deny", true)] { + let runtime = runtime_with_hooks(json!({ + "PreToolUse": [{"hooks": [{ + "type": "command", + "command": format!("printf '%s' '{allow}'"), + "on_failure": mode + }]}] + })); + + let decision = runtime.pre_tool_use(pre_tool_event(mode)).await; + + assert_eq!( + matches!(decision, OzPreToolUseDecision::Deny { .. }), + denied + ); + } +} + +#[tokio::test] +async fn oz_hooks_runtime_timeout_kills_and_resolves_failure_mode() { + let runtime = runtime_with_hooks(json!({ + "PreToolUse": [{"hooks": [{ + "type": "command", + "command": "sleep 30", + "timeout": 1, + "on_failure": "deny" + }]}] + })); + + let started = Instant::now(); + let decision = runtime.pre_tool_use(pre_tool_event("timeout")).await; + + assert!(started.elapsed() < Duration::from_secs(5)); + assert!(matches!(decision, OzPreToolUseDecision::Deny { .. })); +} + +#[tokio::test] +async fn oz_hooks_runtime_rejects_oversized_output() { + let runtime = runtime_with_hooks(json!({ + "PreToolUse": [{"hooks": [{ + "type": "command", + "command": "head -c 70000 /dev/zero", + "on_failure": "continue" + }]}] + })); + + let decision = runtime.pre_tool_use(pre_tool_event("overflow")).await; + + let OzPreToolUseDecision::Continue { diagnostics } = decision else { + panic!("overflow should fail open"); + }; + assert_eq!( + diagnostics[0].failure_category, + Some(HookFailureCategory::OutputOverflow) + ); +} + +#[tokio::test] +async fn oz_hooks_runtime_cancellation_removes_pending_event() { + let runtime = Arc::new(runtime_with_hooks(json!({ + "SessionStart": [{"hooks": [{ + "type": "command", + "command": "sleep 30" + }]}] + }))); + let running = { + let runtime = Arc::clone(&runtime); + tokio::spawn(async move { + runtime + .observe(event( + "cancel", + HookEventFields::SessionStart { + source: SessionStartSource::Startup, + }, + )) + .await + }) + }; + tokio::task::yield_now().await; + + runtime.cancel(OzHookCancellationScope::Invocation("cancel".into())); + let observation = tokio::time::timeout(Duration::from_secs(5), running) + .await + .unwrap() + .unwrap(); + + assert_eq!( + observation.diagnostics[0].result, + HookInvocationResult::Cancelled + ); +} diff --git a/app/src/ai/agent_sdk/hooks/trust.rs b/app/src/ai/agent_sdk/hooks/trust.rs new file mode 100644 index 00000000000..035d0a452fe --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/trust.rs @@ -0,0 +1,44 @@ +use std::collections::HashSet; +use std::path::PathBuf; +use std::sync::RwLock; + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(crate) struct HookTrustKey { + pub(crate) git_root: PathBuf, + pub(crate) config_path: PathBuf, + pub(crate) definition_hash: String, +} + +pub(crate) trait HookTrustStore: Send + Sync { + fn is_trusted(&self, key: &HookTrustKey) -> bool; +} + +#[derive(Default)] +pub(crate) struct DenyProjectHookTrust; + +impl HookTrustStore for DenyProjectHookTrust { + fn is_trusted(&self, _key: &HookTrustKey) -> bool { + false + } +} + +#[derive(Default)] +pub(crate) struct ExactHookTrustStore { + trusted: RwLock>, +} + +impl ExactHookTrustStore { + pub(crate) fn trust(&self, key: HookTrustKey) { + self.trusted.write().unwrap().insert(key); + } + + pub(crate) fn revoke(&self, key: &HookTrustKey) { + self.trusted.write().unwrap().remove(key); + } +} + +impl HookTrustStore for ExactHookTrustStore { + fn is_trusted(&self, key: &HookTrustKey) -> bool { + self.trusted.read().unwrap().contains(key) + } +} diff --git a/app/src/ai/agent_sdk/mod.rs b/app/src/ai/agent_sdk/mod.rs index f6e0cd0729a..e42486da10e 100644 --- a/app/src/ai/agent_sdk/mod.rs +++ b/app/src/ai/agent_sdk/mod.rs @@ -91,6 +91,8 @@ pub(crate) mod environment_snapshot; mod federate; mod harness_support; #[cfg(not(target_family = "wasm"))] +pub(crate) mod hooks; +#[cfg(not(target_family = "wasm"))] mod integration; #[cfg(not(target_family = "wasm"))] mod integration_output; @@ -1064,6 +1066,16 @@ impl AgentDriverRunner { .harness .as_ref() .and_then(|h| h.model_config()); + if args.oz_lifecycle_hooks_context.is_some() { + if args.harness != Harness::Oz { + anyhow::bail!( + "--oz-lifecycle-hooks-context is supported only with the Oz harness" + ); + } + if !FeatureFlag::OzLifecycleHooks.is_enabled() { + anyhow::bail!("this Oz runtime does not support lifecycle hooks"); + } + } let driver_options = driver::AgentDriverOptions { working_dir: working_dir.clone(), task_id, @@ -1093,6 +1105,7 @@ impl AgentDriverRunner { skip_initial_turn: args.skip_initial_turn, strict_mcp_startup: args.strict_mcp_startup, mcp_startup_timeout: args.mcp_startup_timeout.map(|duration| duration.into()), + oz_lifecycle_hooks_context: args.oz_lifecycle_hooks_context.clone(), }; Ok((merged_config, task, driver_options)) diff --git a/crates/warp_features/src/lib.rs b/crates/warp_features/src/lib.rs index 49e10d2ae9a..9e33e74b85a 100644 --- a/crates/warp_features/src/lib.rs +++ b/crates/warp_features/src/lib.rs @@ -995,6 +995,9 @@ pub enum FeatureFlag { /// replace inline computer-use screenshot bytes with references to /// Warp-managed object storage. StoredScreenshots, + + /// Enables first-party Oz lifecycle hooks. + OzLifecycleHooks, } static FLAG_STATES: [AtomicBool; cardinality::()] = From 1801edc7e6ddd39db6dbb0aff8372f0e2ea703ef Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:29:12 +0000 Subject: [PATCH 03/11] Integrate Oz lifecycle hook protocol gates --- Cargo.lock | 34 +- Cargo.toml | 2 +- app/src/ai/agent/api.rs | 3 + app/src/ai/agent/api/convert_to.rs | 3 + app/src/ai/agent/api/impl.rs | 2 + app/src/ai/agent/api/impl_tests.rs | 1 + app/src/ai/agent/conversation.rs | 3 +- app/src/ai/agent/mod.rs | 15 +- app/src/ai/agent/redaction.rs | 3 +- app/src/ai/agent_sdk/driver.rs | 237 ++++++++++- app/src/ai/agent_sdk/driver/output.rs | 6 +- app/src/ai/agent_sdk/driver_tests.rs | 1 + app/src/ai/agent_sdk/hooks/config.rs | 1 + app/src/ai/agent_sdk/hooks/mod.rs | 20 + app/src/ai/agent_sdk/hooks/payload.rs | 10 +- app/src/ai/agent_sdk/hooks/permissions.rs | 1 + app/src/ai/agent_sdk/hooks/protocol.rs | 379 ++++++++++++++++++ app/src/ai/agent_sdk/hooks/protocol_tests.rs | 261 ++++++++++++ app/src/ai/agent_sdk/hooks/redaction.rs | 26 +- app/src/ai/agent_sdk/hooks/runtime.rs | 58 ++- app/src/ai/blocklist/block/view_impl.rs | 3 +- .../ai/blocklist/block/view_impl/common.rs | 3 +- app/src/ai/blocklist/controller.rs | 234 ++++++++++- app/src/ai/blocklist/controller_tests.rs | 1 + app/src/ai/blocklist/history_model_tests.rs | 2 + app/src/ai/blocklist/permissions.rs | 12 +- app/src/ai/blocklist/permissions_tests.rs | 1 + app/src/ai/blocklist/persistence.rs | 3 +- app/src/server/telemetry/events.rs | 2 + .../replay_agent_conversations.rs | 1 + .../view/shared_session/view_impl_tests.rs | 1 + crates/warp_cli/src/agent.rs | 96 +++++ crates/warp_cli/src/lib_tests.rs | 46 +++ specs/APP-4344/TECH.md | 27 ++ 34 files changed, 1423 insertions(+), 75 deletions(-) create mode 100644 app/src/ai/agent_sdk/hooks/protocol.rs create mode 100644 app/src/ai/agent_sdk/hooks/protocol_tests.rs diff --git a/Cargo.lock b/Cargo.lock index ef632433429..7f3799d17e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -651,7 +651,7 @@ dependencies = [ "objc2-foundation 0.3.2", "parking_lot", "percent-encoding", - "windows-sys 0.59.0", + "windows-sys 0.52.0", "wl-clipboard-rs", "x11rb", ] @@ -2150,7 +2150,7 @@ dependencies = [ "bitflags 2.13.0", "cexpr", "clang-sys", - "itertools 0.11.0", + "itertools 0.10.5", "lazy_static", "lazycell", "log", @@ -3188,7 +3188,7 @@ version = "3.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fde0e0ec90c9dfb3b4b1a0891a7dcd0e2bffde2f7efed5fe7c9bb00e5bfb915e" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -4527,7 +4527,7 @@ checksum = "6738d2e996274e499bc7b0d693c858b7720b9cd2543a0643a3087e6cb0a4fa16" dependencies = [ "cfg-if", "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -4832,7 +4832,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -7158,7 +7158,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi 0.5.2", "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -7352,7 +7352,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -10597,8 +10597,8 @@ version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ - "heck 0.4.1", - "itertools 0.11.0", + "heck 0.5.0", + "itertools 0.10.5", "log", "multimap", "petgraph", @@ -10617,7 +10617,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -10863,7 +10863,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -11863,7 +11863,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -11876,7 +11876,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.11.0", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -11944,7 +11944,7 @@ dependencies = [ "security-framework 3.5.1", "security-framework-sys", "webpki-root-certs 1.0.1", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -13656,7 +13656,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.1.2", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -15861,7 +15861,7 @@ dependencies = [ [[package]] name = "warp_multi_agent_api" version = "0.0.0" -source = "git+https://github.com/warpdotdev/warp-proto-apis.git?rev=7b638a3bce15b6159116c25caa1dda62cf004b0e#7b638a3bce15b6159116c25caa1dda62cf004b0e" +source = "git+https://github.com/warpdotdev/warp-proto-apis.git?rev=94d0881ca1cad6323c50b28d33bbb8526d434a54#94d0881ca1cad6323c50b28d33bbb8526d434a54" dependencies = [ "prost", "prost-reflect", @@ -17107,7 +17107,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 0dbb0ffba74..aa3fc3f59de 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -349,7 +349,7 @@ version-compare = "0.1" vte = { git = "https://github.com/warpdotdev/vte.git", rev = "4b399c87b63ba88f45709edaa6383fc519f6c900", default-features = false } walkdir = "2" warp-workflows = { git = "https://github.com/warpdotdev/workflows", rev = "793a98ddda6ef19682aed66364faebd2829f0e01" } -warp_multi_agent_api = { git = "https://github.com/warpdotdev/warp-proto-apis.git", rev = "7b638a3bce15b6159116c25caa1dda62cf004b0e" } +warp_multi_agent_api = { git = "https://github.com/warpdotdev/warp-proto-apis.git", rev = "94d0881ca1cad6323c50b28d33bbb8526d434a54" } wasm-bindgen = "0.2.89" wasm-bindgen-futures = "0.4.42" web-sys = { version = "0.3.69", features = [ diff --git a/app/src/ai/agent/api.rs b/app/src/ai/agent/api.rs index 9b2dbeea5dc..2b5f783d0be 100644 --- a/app/src/ai/agent/api.rs +++ b/app/src/ai/agent/api.rs @@ -137,6 +137,7 @@ pub struct RequestParams { pub warp_drive_context_enabled: bool, pub context_window_limit: Option, pub mcp_context: Option, + pub oz_hook_context: Option, pub planning_enabled: bool, should_redact_secrets: bool, @@ -210,6 +211,7 @@ impl RequestParams { warp_drive_context_enabled: false, context_window_limit: None, mcp_context: None, + oz_hook_context: None, planning_enabled: false, should_redact_secrets: false, member_byo_credentials_allowed: false, @@ -430,6 +432,7 @@ impl RequestParams { supported_tools_override: request_input.supported_tools_override.clone(), parent_agent_id: None, agent_name: None, + oz_hook_context: None, } } } diff --git a/app/src/ai/agent/api/convert_to.rs b/app/src/ai/agent/api/convert_to.rs index b22b2992253..84900a88368 100644 --- a/app/src/ai/agent/api/convert_to.rs +++ b/app/src/ai/agent/api/convert_to.rs @@ -340,6 +340,9 @@ fn convert_input_to_user_input( )) } AIAgentInput::ActionResult { result, .. } => result.try_into(), + AIAgentInput::OzHookResult(result) => { + Ok(api::request::input::user_inputs::user_input::Input::OzHookResult(result)) + } AIAgentInput::MessagesReceivedFromAgents { messages } => Ok( api::request::input::user_inputs::user_input::Input::MessagesReceivedFromAgents( api::request::input::user_inputs::MessagesReceivedFromAgents { diff --git a/app/src/ai/agent/api/impl.rs b/app/src/ai/agent/api/impl.rs index f963ec6bfba..f84aa970e52 100644 --- a/app/src/ai/agent/api/impl.rs +++ b/app/src/ai/agent/api/impl.rs @@ -111,6 +111,7 @@ pub async fn generate_multi_agent_output( supports_stored_screenshots: FeatureFlag::StoredScreenshots.is_enabled(), custom_model_providers: params.custom_model_providers, custom_model_routers: params.custom_model_routers, + supports_oz_lifecycle_hooks: params.oz_hook_context.is_some(), }), metadata: Some(api::request::Metadata { logging: logging_metadata, @@ -140,6 +141,7 @@ pub async fn generate_multi_agent_output( .existing_suggestions .map(|suggestions| suggestions.into()), mcp_context: params.mcp_context.map(Into::into), + oz_hook_context: params.oz_hook_context, }; let response_stream = warp_multi_agent_client::generate_multi_agent_output( diff --git a/app/src/ai/agent/api/impl_tests.rs b/app/src/ai/agent/api/impl_tests.rs index b1f64cb5e05..613ff7815a6 100644 --- a/app/src/ai/agent/api/impl_tests.rs +++ b/app/src/ai/agent/api/impl_tests.rs @@ -48,6 +48,7 @@ fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool supported_tools_override: None, parent_agent_id: None, agent_name: None, + oz_hook_context: None, } } diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index ab1a28f5d4a..e74da7e65b1 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -2121,7 +2121,7 @@ impl AIConversation { for (task_id, inputs) in input_messages.into_iter() { let should_hide = inputs .iter() - .any(|input| input.is_passive_suggestion_trigger()); + .any(|input| input.is_passive_suggestion_trigger() || input.is_oz_hook_result()); let new_exchange = AIAgentExchange { id: AIAgentExchangeId::new(), @@ -2829,6 +2829,7 @@ impl AIConversation { log::debug!("Rollback transaction."); self.rollback_transaction(response_stream_id); } + Action::RunOzHook(_) => {} Action::CreateTask(CreateTask { task: Some(task) }) => { let task_id = TaskId::new(task.id.clone()); // Save an empty task to the transaction diff --git a/app/src/ai/agent/mod.rs b/app/src/ai/agent/mod.rs index 30d78bf77ec..978bd19a4e2 100644 --- a/app/src/ai/agent/mod.rs +++ b/app/src/ai/agent/mod.rs @@ -2989,6 +2989,7 @@ pub enum AIAgentInput { config: OrchestrationConfig, status: OrchestrationConfigStatus, }, + OzHookResult(warp_multi_agent_api::OzHookResult), } /// Data for a single message received by an agent from another agent. @@ -3082,6 +3083,7 @@ impl Display for AIAgentInput { } Self::PassiveSuggestionResult { .. } => write!(f, "PassiveSuggestionResult"), Self::OrchestrationConfigUpdate { .. } => write!(f, "OrchestrationConfigUpdate"), + Self::OzHookResult(_) => write!(f, "OzHookResult"), } } } @@ -3143,7 +3145,8 @@ impl AIAgentInput { | Self::MessagesReceivedFromAgents { .. } | Self::EventsFromAgents { .. } | Self::PassiveSuggestionResult { .. } - | Self::OrchestrationConfigUpdate { .. } => None, + | Self::OrchestrationConfigUpdate { .. } + | Self::OzHookResult(_) => None, } } @@ -3211,6 +3214,10 @@ impl AIAgentInput { matches!(self, AIAgentInput::UserQuery { .. }) } + pub fn is_oz_hook_result(&self) -> bool { + matches!(self, AIAgentInput::OzHookResult(_)) + } + pub fn prompt_suggestion_result(&self) -> Option<&String> { if let Some(AIAgentActionResult { result: AIAgentActionResultType::SuggestPrompt(SuggestPromptResult::Accepted { query }), @@ -3248,7 +3255,8 @@ impl AIAgentInput { Self::SummarizeConversation { context, .. } => Some(context), Self::MessagesReceivedFromAgents { .. } | Self::EventsFromAgents { .. } - | Self::OrchestrationConfigUpdate { .. } => None, + | Self::OrchestrationConfigUpdate { .. } + | Self::OzHookResult(_) => None, } } @@ -3279,7 +3287,8 @@ impl AIAgentInput { | Self::MessagesReceivedFromAgents { .. } | Self::EventsFromAgents { .. } | Self::PassiveSuggestionResult { .. } - | Self::OrchestrationConfigUpdate { .. } => None, + | Self::OrchestrationConfigUpdate { .. } + | Self::OzHookResult(_) => None, } } diff --git a/app/src/ai/agent/redaction.rs b/app/src/ai/agent/redaction.rs index 5b328a379e0..6bc85ad027d 100644 --- a/app/src/ai/agent/redaction.rs +++ b/app/src/ai/agent/redaction.rs @@ -90,7 +90,8 @@ pub(crate) fn redact_inputs(inputs: &mut [AIAgentInput]) { // No user-provided text to redact in inter-agent relay inputs. AIAgentInput::MessagesReceivedFromAgents { .. } | AIAgentInput::EventsFromAgents { .. } - | AIAgentInput::OrchestrationConfigUpdate { .. } => {} + | AIAgentInput::OrchestrationConfigUpdate { .. } + | AIAgentInput::OzHookResult(_) => {} AIAgentInput::ActionResult { result, context } => { redact_context(Arc::make_mut(context)); match &mut result.result { diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index 2c20572b66f..b86784b2033 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -26,7 +26,7 @@ use repo_metadata::{RepoMetadataModel, RepositoryIdentifier}; use session_sharing_protocol::sharer::SessionRetentionReason; use tracing::Instrument as _; use uuid::Uuid; -use warp_cli::agent::{Harness, OutputFormat, RepositoryHeadOverride}; +use warp_cli::agent::{Harness, OutputFormat, OzLifecycleHooksContext, RepositoryHeadOverride}; use warp_cli::mcp::MCPSpec; use warp_cli::share::ShareRequest; use warp_cli::skill::SkillSpec; @@ -55,6 +55,15 @@ use crate::ai::agent_sdk::driver::harness::{ use crate::ai::agent_sdk::environment_snapshot::{ EnvironmentSnapshot, EnvironmentSnapshotReporter, }; +use crate::ai::agent_sdk::hooks::config::discover_hook_config; +use crate::ai::agent_sdk::hooks::payload::{ + HookEventFields, HookPayloadContext, HookPayloadTemplate, SessionEndReason, SessionStartSource, + TurnStatus, +}; +use crate::ai::agent_sdk::hooks::redaction::HookRedactor; +use crate::ai::agent_sdk::hooks::runtime::{OzHookEvent, OzHookRuntime, OzHookRuntimeService}; +use crate::ai::agent_sdk::hooks::trust::{ExactHookTrustStore, HookTrustKey}; +use crate::ai::agent_sdk::hooks::{MAX_PROMPT_BYTES, OzHookSession, PAYLOAD_SCHEMA_VERSION}; use crate::ai::agent_sdk::setup_observability::{SetupClientEventReporter, SetupStep}; use crate::ai::ambient_agents::task::HarnessModelConfig; use crate::ai::ambient_agents::{ @@ -617,6 +626,8 @@ pub struct AgentDriverOptions { pub strict_mcp_startup: bool, /// MCP server startup timeout override. pub mcp_startup_timeout: Option, + /// Server-authenticated lifecycle hook capability and project trust. + pub oz_lifecycle_hooks_context: Option, } /// `AgentDriver` is a model for driving an ambient Warp agent to completion. @@ -722,6 +733,8 @@ pub struct AgentDriver { /// How long to wait for MCP servers to start before degrading (or failing, /// in strict mode). mcp_startup_timeout: Duration, + /// Server-authenticated lifecycle hook capability and project trust. + oz_lifecycle_hooks_context: Option, } #[derive(Clone)] @@ -993,6 +1006,219 @@ impl From for AgentDriverError { } impl AgentDriver { + async fn initialize_oz_hook_runtime( + foreground: &ModelSpawner, + ) -> Result>, AgentDriverError> { + let (context, cwd, task_id, hooks_enabled) = foreground + .spawn(|me, _| { + ( + me.oz_lifecycle_hooks_context.clone(), + me.harness_working_dir.clone(), + me.task_id, + FeatureFlag::OzLifecycleHooks.is_enabled(), + ) + }) + .await?; + if !hooks_enabled { + return Ok(None); + } + + let trust_store = ExactHookTrustStore::default(); + if let Some(context) = context { + for trust in context.project_trust { + let (Ok(git_root), Ok(config_path)) = ( + std::fs::canonicalize(trust.git_root), + std::fs::canonicalize(trust.config_path), + ) else { + continue; + }; + trust_store.trust(HookTrustKey { + git_root, + config_path, + definition_hash: trust.sha256, + }); + } + } + let config = discover_hook_config(&cwd, &trust_store); + for diagnostic in config.diagnostics.iter() { + log::warn!( + "Oz hook configuration diagnostic: kind={:?} path={} hash_present={}", + diagnostic.kind, + diagnostic.path.display(), + diagnostic.definition_hash.is_some() + ); + } + let enabled_events = config + .enabled_events() + .map(|event| event.protocol_value().into()) + .collect(); + let runtime: Arc = Arc::new(OzHookRuntimeService::new(config)); + let session = OzHookSession { + runtime: Arc::clone(&runtime), + protocol_context: warp_multi_agent_api::OzHookContext { + enabled_events, + supported_payload_schema_versions: vec![PAYLOAD_SCHEMA_VERSION.into()], + }, + }; + foreground + .spawn(move |me, ctx| { + me.terminal_driver.update(ctx, |driver, ctx| { + driver.with_terminal_view(ctx, |terminal, ctx| { + terminal.ai_controller().update(ctx, |controller, _| { + controller.set_oz_hook_session(Some(session)); + }); + }); + }); + }) + .await?; + + let run_id = task_id + .map(|id| id.to_string()) + .unwrap_or_else(|| Uuid::new_v4().to_string()); + runtime + .observe(OzHookEvent { + invocation_id: Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: HookPayloadContext { + session_id: run_id.clone(), + run_id, + conversation_id: String::new(), + cwd: cwd.to_string_lossy().into_owned(), + model: String::new(), + permission_mode: "supervised".into(), + }, + event: HookEventFields::SessionStart { + source: SessionStartSource::Startup, + }, + }, + }) + .await; + Ok(Some(runtime)) + } + + async fn finish_oz_hook_runtime( + runtime: Option>, + status: &SDKConversationOutputStatus, + foreground: &ModelSpawner, + ) { + let Some(runtime) = runtime else { + return; + }; + let Ok((cwd, task_id)) = foreground + .spawn(|me, _| (me.harness_working_dir.clone(), me.task_id)) + .await + else { + return; + }; + let reason = match status { + SDKConversationOutputStatus::Success => SessionEndReason::Completed, + SDKConversationOutputStatus::Error { .. } + | SDKConversationOutputStatus::Blocked { .. } => SessionEndReason::Failed, + SDKConversationOutputStatus::Cancelled { .. } => SessionEndReason::Cancelled, + }; + let run_id = task_id + .map(|id| id.to_string()) + .unwrap_or_else(|| Uuid::new_v4().to_string()); + runtime + .observe(OzHookEvent { + invocation_id: Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: HookPayloadContext { + session_id: run_id.clone(), + run_id, + conversation_id: String::new(), + cwd: cwd.to_string_lossy().into_owned(), + model: String::new(), + permission_mode: "supervised".into(), + }, + event: HookEventFields::SessionEnd { reason }, + }, + }) + .await; + runtime.cancel(crate::ai::agent_sdk::hooks::runtime::OzHookCancellationScope::Session); + } + async fn observe_oz_prompt( + runtime: Option<&Arc>, + prompt: &AgentRunPrompt, + foreground: &ModelSpawner, + ) { + let (Some(runtime), AgentRunPrompt::Local(prompt)) = (runtime, prompt) else { + return; + }; + let Ok((cwd, task_id)) = foreground + .spawn(|me, _| (me.harness_working_dir.clone(), me.task_id)) + .await + else { + return; + }; + let run_id = task_id + .map(|id| id.to_string()) + .unwrap_or_else(|| Uuid::new_v4().to_string()); + runtime + .observe(OzHookEvent { + invocation_id: Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: HookPayloadContext { + session_id: run_id.clone(), + run_id, + conversation_id: String::new(), + cwd: cwd.to_string_lossy().into_owned(), + model: String::new(), + permission_mode: "supervised".into(), + }, + event: HookEventFields::user_prompt( + HookRedactor::new([]).redact_text(prompt, MAX_PROMPT_BYTES), + ), + }, + }) + .await; + } + + async fn observe_oz_stop( + runtime: Option<&Arc>, + status: &SDKConversationOutputStatus, + foreground: &ModelSpawner, + ) { + let Some(runtime) = runtime else { + return; + }; + let Ok((cwd, task_id)) = foreground + .spawn(|me, _| (me.harness_working_dir.clone(), me.task_id)) + .await + else { + return; + }; + let turn_status = match status { + SDKConversationOutputStatus::Success => TurnStatus::Completed, + SDKConversationOutputStatus::Error { .. } => TurnStatus::Failed, + SDKConversationOutputStatus::Cancelled { .. } => TurnStatus::Idle, + SDKConversationOutputStatus::Blocked { .. } => TurnStatus::Blocked, + }; + let run_id = task_id + .map(|id| id.to_string()) + .unwrap_or_else(|| Uuid::new_v4().to_string()); + runtime + .observe(OzHookEvent { + invocation_id: Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: HookPayloadContext { + session_id: run_id.clone(), + run_id, + conversation_id: String::new(), + cwd: cwd.to_string_lossy().into_owned(), + model: String::new(), + permission_mode: "supervised".into(), + }, + event: HookEventFields::Stop { turn_status }, + }, + }) + .await; + } + #[tracing::instrument(name = "AgentDriver::new", skip_all, err, fields( tags.cloud_agent = true, task_id = ?options.task_id, @@ -1026,6 +1252,7 @@ impl AgentDriver { skip_initial_turn, strict_mcp_startup, mcp_startup_timeout, + oz_lifecycle_hooks_context, } = options; // Split the unified resume option into the two internal slots that the rest of @@ -1194,6 +1421,7 @@ impl AgentDriver { skip_initial_turn, strict_mcp_startup, mcp_startup_timeout: mcp_startup_timeout.unwrap_or(MCP_SERVER_STARTUP_TIMEOUT), + oz_lifecycle_hooks_context, }) } @@ -1244,6 +1472,7 @@ impl AgentDriver { skip_initial_turn: false, strict_mcp_startup: false, mcp_startup_timeout: MCP_SERVER_STARTUP_TIMEOUT, + oz_lifecycle_hooks_context: None, } } @@ -2383,6 +2612,8 @@ impl AgentDriver { &foreground, ) .await?; + let oz_hook_runtime = Self::initialize_oz_hook_runtime(&foreground).await?; + Self::observe_oz_prompt(oz_hook_runtime.as_ref(), &task.prompt, &foreground).await; let status_rx = foreground .spawn(move |me, ctx| me.execute_run(task.prompt, ctx)) @@ -2401,6 +2632,8 @@ impl AgentDriver { &foreground, ) .await?; + Self::observe_oz_stop(oz_hook_runtime.as_ref(), &conversation_status, &foreground) + .await; log::info!( "Ambient agent Oz lifecycle: event=run_exit_received idle_on_complete_elapsed_or_not_configured=true next=terminal_teardown_after_flush" @@ -2413,6 +2646,8 @@ impl AgentDriver { // to send a message when the streams are finished, flushed, and the websocket is disconnected. For now, we'll just sleep for a second, as this seems // to be enough time for the streams to be finished and the events to be flushed. warpui::r#async::Timer::after(Duration::from_secs(1)).await; + Self::finish_oz_hook_runtime(oz_hook_runtime, &conversation_status, &foreground) + .await; conversation_status.into_result() } diff --git a/app/src/ai/agent_sdk/driver/output.rs b/app/src/ai/agent_sdk/driver/output.rs index 05c4b2abd01..b1411a2a7ab 100644 --- a/app/src/ai/agent_sdk/driver/output.rs +++ b/app/src/ai/agent_sdk/driver/output.rs @@ -38,7 +38,8 @@ pub mod text { | AIAgentInput::MessagesReceivedFromAgents { .. } | AIAgentInput::PassiveSuggestionResult { .. } | AIAgentInput::EventsFromAgents { .. } - | AIAgentInput::OrchestrationConfigUpdate { .. } => { + | AIAgentInput::OrchestrationConfigUpdate { .. } + | AIAgentInput::OzHookResult(_) => { // Do not include the user query, since it's already provided as input to the agent. Ok(()) } @@ -823,7 +824,8 @@ pub mod json { | AIAgentInput::MessagesReceivedFromAgents { .. } | AIAgentInput::EventsFromAgents { .. } | AIAgentInput::PassiveSuggestionResult { .. } - | AIAgentInput::OrchestrationConfigUpdate { .. } => None, + | AIAgentInput::OrchestrationConfigUpdate { .. } + | AIAgentInput::OzHookResult(_) => None, // These input types should not occur in a SDK-run agent. AIAgentInput::ResumeConversation { .. } | AIAgentInput::TriggerPassiveSuggestion { .. } => None, diff --git a/app/src/ai/agent_sdk/driver_tests.rs b/app/src/ai/agent_sdk/driver_tests.rs index 2b5b36e24c7..29059178138 100644 --- a/app/src/ai/agent_sdk/driver_tests.rs +++ b/app/src/ai/agent_sdk/driver_tests.rs @@ -1518,6 +1518,7 @@ fn complete_mock_stream_successfully(app: &mut App, stream: &warpui::ModelHandle request_id: "test-request".to_string(), conversation_id: "test-server-conversation".to_string(), run_id: String::new(), + supported_oz_hook_payload_schema_versions: vec![], })), }, ctx, diff --git a/app/src/ai/agent_sdk/hooks/config.rs b/app/src/ai/agent_sdk/hooks/config.rs index 6c66b315013..4f422933e24 100644 --- a/app/src/ai/agent_sdk/hooks/config.rs +++ b/app/src/ai/agent_sdk/hooks/config.rs @@ -23,6 +23,7 @@ pub(crate) struct ConfiguredHook { pub(crate) matcher_text: Option, matcher: Option, pub(crate) command: String, + #[cfg_attr(not(windows), allow(dead_code))] pub(crate) command_windows: Option, pub(crate) timeout: Duration, pub(crate) on_failure: FailureMode, diff --git a/app/src/ai/agent_sdk/hooks/mod.rs b/app/src/ai/agent_sdk/hooks/mod.rs index b1beea7cc06..fe6972c62a0 100644 --- a/app/src/ai/agent_sdk/hooks/mod.rs +++ b/app/src/ai/agent_sdk/hooks/mod.rs @@ -1,10 +1,12 @@ use std::fmt; +use std::sync::Arc; use serde::{Deserialize, Serialize}; pub(crate) mod config; pub(crate) mod payload; pub(crate) mod permissions; +pub(crate) mod protocol; pub(crate) mod redaction; pub(crate) mod runtime; pub(crate) mod trust; @@ -20,6 +22,12 @@ pub(crate) const MAX_TOOL_RESPONSE_BYTES: usize = 64 * 1024; pub(crate) const MAX_OUTPUT_BYTES: usize = 64 * 1024; pub(crate) const MAX_DENIAL_REASON_BYTES: usize = 4 * 1024; +#[derive(Clone)] +pub(crate) struct OzHookSession { + pub(crate) runtime: Arc, + pub(crate) protocol_context: warp_multi_agent_api::OzHookContext, +} + #[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] pub(crate) enum HookEventName { SessionStart, @@ -57,6 +65,18 @@ impl HookEventName { pub(crate) const fn ignores_matcher(self) -> bool { matches!(self, Self::UserPromptSubmit | Self::Stop) } + + pub(crate) const fn protocol_value(self) -> warp_multi_agent_api::OzHookEvent { + match self { + Self::SessionStart => warp_multi_agent_api::OzHookEvent::SessionStart, + Self::SessionEnd => warp_multi_agent_api::OzHookEvent::SessionEnd, + Self::UserPromptSubmit => warp_multi_agent_api::OzHookEvent::UserPromptSubmit, + Self::Stop => warp_multi_agent_api::OzHookEvent::Stop, + Self::PreToolUse => warp_multi_agent_api::OzHookEvent::PreToolUse, + Self::PostToolUse => warp_multi_agent_api::OzHookEvent::PostToolUse, + Self::PreCompact => warp_multi_agent_api::OzHookEvent::PreCompact, + } + } } impl fmt::Display for HookEventName { diff --git a/app/src/ai/agent_sdk/hooks/payload.rs b/app/src/ai/agent_sdk/hooks/payload.rs index 393f4742ae2..fc12f79bcfa 100644 --- a/app/src/ai/agent_sdk/hooks/payload.rs +++ b/app/src/ai/agent_sdk/hooks/payload.rs @@ -1,4 +1,4 @@ -use serde::Serialize; +use serde::{Deserialize, Serialize}; use super::redaction::{RedactedText, RedactedValue, TruncationMetadata}; use super::{HookConfigSource, HookEventName, MAX_PAYLOAD_BYTES, PAYLOAD_SCHEMA_VERSION}; @@ -118,7 +118,7 @@ impl HookEventFields { } } -#[derive(Clone, Copy, Debug, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Serialize)] #[serde(rename_all = "lowercase")] pub(crate) enum SessionStartSource { Startup, @@ -134,7 +134,7 @@ impl SessionStartSource { } } -#[derive(Clone, Copy, Debug, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Serialize)] #[serde(rename_all = "lowercase")] pub(crate) enum SessionEndReason { Completed, @@ -154,7 +154,7 @@ impl SessionEndReason { } } -#[derive(Clone, Copy, Debug, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Serialize)] #[serde(rename_all = "lowercase")] pub(crate) enum CompactTrigger { Auto, @@ -170,7 +170,7 @@ impl CompactTrigger { } } -#[derive(Clone, Copy, Debug, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Serialize)] #[serde(rename_all = "snake_case")] pub(crate) enum TurnStatus { Idle, diff --git a/app/src/ai/agent_sdk/hooks/permissions.rs b/app/src/ai/agent_sdk/hooks/permissions.rs index d16e469f116..2371574e302 100644 --- a/app/src/ai/agent_sdk/hooks/permissions.rs +++ b/app/src/ai/agent_sdk/hooks/permissions.rs @@ -1,4 +1,5 @@ #[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[allow(dead_code)] pub(crate) enum NativePermission { Deny, Allow, diff --git a/app/src/ai/agent_sdk/hooks/protocol.rs b/app/src/ai/agent_sdk/hooks/protocol.rs new file mode 100644 index 00000000000..f87cc09e094 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/protocol.rs @@ -0,0 +1,379 @@ +use std::collections::BTreeMap; + +use prost_types::value::Kind; +use warp_multi_agent_api::oz_hook_result::{self, ResolvedAction}; +use warp_multi_agent_api::{OzHookEvent as ProtocolEvent, OzHookResult, RunOzHook}; + +use super::payload::{HookEventFields, HookPayloadContext, HookPayloadTemplate}; +use super::redaction::RedactedValue; +use super::runtime::{ + HookFailureCategory, HookInvocationDiagnostic, HookInvocationResult, OzHookEvent, + OzPreToolUseDecision, +}; +use super::{ + HookEventName, MAX_DENIAL_REASON_BYTES, MAX_PAYLOAD_BYTES, MAX_TOOL_INPUT_BYTES, + MAX_TOOL_RESPONSE_BYTES, PAYLOAD_SCHEMA_VERSION, +}; + +pub(crate) fn event_from_protocol(action: &RunOzHook) -> Result { + if action.invocation_id.is_empty() { + return Err(ProtocolHookError::InvalidInvocationId); + } + if action.schema_version != PAYLOAD_SCHEMA_VERSION { + return Err(ProtocolHookError::UnsupportedSchema); + } + let protocol_event = + ProtocolEvent::try_from(action.event).map_err(|_| ProtocolHookError::UnknownEvent)?; + let event_name = match protocol_event { + ProtocolEvent::SessionStart => HookEventName::SessionStart, + ProtocolEvent::SessionEnd => HookEventName::SessionEnd, + ProtocolEvent::UserPromptSubmit => HookEventName::UserPromptSubmit, + ProtocolEvent::Stop => HookEventName::Stop, + ProtocolEvent::PreToolUse => HookEventName::PreToolUse, + ProtocolEvent::PostToolUse => HookEventName::PostToolUse, + ProtocolEvent::PreCompact => HookEventName::PreCompact, + ProtocolEvent::Unspecified => return Err(ProtocolHookError::UnknownEvent), + }; + let payload = action + .redacted_payload + .as_ref() + .ok_or(ProtocolHookError::MissingPayload)?; + let mut fields = payload + .fields + .iter() + .map(|(key, value)| Ok((key.clone(), value_from_protocol(value)?))) + .collect::, ProtocolHookError>>()?; + if serde_json::to_vec(&fields).map_or(true, |bytes| bytes.len() > MAX_PAYLOAD_BYTES) { + return Err(ProtocolHookError::OversizedPayload); + } + validate_envelope_fields(&mut fields, event_name)?; + + let context = HookPayloadContext { + session_id: take_string(&mut fields, "session_id")?, + run_id: take_string(&mut fields, "run_id")?, + conversation_id: take_string(&mut fields, "conversation_id")?, + cwd: take_string(&mut fields, "cwd")?, + model: take_string(&mut fields, "model")?, + permission_mode: take_string(&mut fields, "permission_mode")?, + }; + let event = match event_name { + HookEventName::SessionStart => HookEventFields::SessionStart { + source: take_enum(&mut fields, "source")?, + }, + HookEventName::SessionEnd => HookEventFields::SessionEnd { + reason: take_enum(&mut fields, "reason")?, + }, + HookEventName::UserPromptSubmit => HookEventFields::UserPromptSubmit { + prompt: take_string(&mut fields, "prompt")?, + prompt_truncation: take_optional_enum(&mut fields, "prompt_truncation")?, + }, + HookEventName::Stop => HookEventFields::Stop { + turn_status: take_enum(&mut fields, "turn_status")?, + }, + HookEventName::PreToolUse => { + let tool_input = take_value(&mut fields, "tool_input")?; + validate_value_size(&tool_input, MAX_TOOL_INPUT_BYTES)?; + HookEventFields::PreToolUse { + tool_name: take_string(&mut fields, "tool_name")?, + tool_use_id: take_string(&mut fields, "tool_use_id")?, + tool_input, + } + } + HookEventName::PostToolUse => { + let tool_input = take_value(&mut fields, "tool_input")?; + let tool_response = take_value(&mut fields, "tool_response")?; + validate_value_size(&tool_input, MAX_TOOL_INPUT_BYTES)?; + validate_value_size(&tool_response, MAX_TOOL_RESPONSE_BYTES)?; + HookEventFields::PostToolUse { + tool_name: take_string(&mut fields, "tool_name")?, + tool_use_id: take_string(&mut fields, "tool_use_id")?, + tool_input, + tool_response, + } + } + HookEventName::PreCompact => HookEventFields::PreCompact { + trigger: take_enum(&mut fields, "trigger")?, + }, + }; + if !fields.is_empty() { + return Err(ProtocolHookError::UnknownPayloadField); + } + if matches!( + event_name, + HookEventName::PreToolUse | HookEventName::PostToolUse + ) && action.tool_use_id.is_empty() + { + return Err(ProtocolHookError::InvalidToolUseId); + } + if let HookEventFields::PreToolUse { tool_use_id, .. } + | HookEventFields::PostToolUse { tool_use_id, .. } = &event + && tool_use_id != &action.tool_use_id + { + return Err(ProtocolHookError::MismatchedToolUseId); + } + Ok(OzHookEvent { + invocation_id: action.invocation_id.clone(), + tool_use_id: (!action.tool_use_id.is_empty()).then(|| action.tool_use_id.clone()), + payload: HookPayloadTemplate { context, event }, + }) +} + +fn validate_value_size( + value: &RedactedValue, + maximum_bytes: usize, +) -> Result<(), ProtocolHookError> { + if value.serialized_len() > maximum_bytes { + Err(ProtocolHookError::OversizedPayload) + } else { + Ok(()) + } +} + +pub(crate) fn result_for_observation( + action: &RunOzHook, + diagnostics: &[HookInvocationDiagnostic], +) -> OzHookResult { + OzHookResult { + invocation_id: action.invocation_id.clone(), + tool_use_id: action.tool_use_id.clone(), + outcome: outcome_for_diagnostics(diagnostics, ResolvedAction::Continue), + } +} + +pub(crate) fn result_for_pre_tool( + action: &RunOzHook, + decision: OzPreToolUseDecision, +) -> OzHookResult { + let outcome = match decision { + OzPreToolUseDecision::Continue { diagnostics } => { + outcome_for_diagnostics(&diagnostics, ResolvedAction::Continue) + } + OzPreToolUseDecision::Deny { + reason, + source, + diagnostics, + } => diagnostics + .iter() + .rev() + .find(|diagnostic| diagnostic.failure_category.is_some()) + .map(|diagnostic| Some(failed_outcome(diagnostic, ResolvedAction::Deny))) + .unwrap_or_else(|| { + Some(oz_hook_result::Outcome::Deny(oz_hook_result::Deny { + reason: super::redaction::truncate_utf8(&reason, MAX_DENIAL_REASON_BYTES), + source: source.as_str().into(), + })) + }), + }; + OzHookResult { + invocation_id: action.invocation_id.clone(), + tool_use_id: action.tool_use_id.clone(), + outcome, + } +} + +pub(crate) fn failed_result(action: &RunOzHook, error: ProtocolHookError) -> OzHookResult { + OzHookResult { + invocation_id: action.invocation_id.clone(), + tool_use_id: action.tool_use_id.clone(), + outcome: Some(oz_hook_result::Outcome::Failed(oz_hook_result::Failed { + category: error.category().into(), + resolved_action: ResolvedAction::Continue.into(), + })), + } +} + +fn outcome_for_diagnostics( + diagnostics: &[HookInvocationDiagnostic], + resolved_action: ResolvedAction, +) -> Option { + let Some(diagnostic) = diagnostics + .iter() + .rev() + .find(|diagnostic| diagnostic.failure_category.is_some()) + else { + return Some(oz_hook_result::Outcome::Continue( + oz_hook_result::Continue {}, + )); + }; + if diagnostic.result == HookInvocationResult::Cancelled { + return Some(oz_hook_result::Outcome::Cancelled( + oz_hook_result::Cancelled {}, + )); + } + Some(failed_outcome(diagnostic, resolved_action)) +} + +fn failed_outcome( + diagnostic: &HookInvocationDiagnostic, + resolved_action: ResolvedAction, +) -> oz_hook_result::Outcome { + oz_hook_result::Outcome::Failed(oz_hook_result::Failed { + category: diagnostic + .failure_category + .map(failure_category_name) + .unwrap_or("unknown") + .into(), + resolved_action: resolved_action.into(), + }) +} + +fn failure_category_name(category: HookFailureCategory) -> &'static str { + match category { + HookFailureCategory::Spawn => "spawn", + HookFailureCategory::Stdin => "stdin", + HookFailureCategory::Timeout => "timeout", + HookFailureCategory::Cancelled => "cancelled", + HookFailureCategory::OutputOverflow => "output_overflow", + HookFailureCategory::OutputRead => "output_read", + HookFailureCategory::InvalidUtf8 => "invalid_utf8", + HookFailureCategory::NonZeroExit => "nonzero_exit", + HookFailureCategory::InvalidDecision => "invalid_decision", + HookFailureCategory::Payload => "payload", + } +} + +fn validate_envelope_fields( + fields: &mut BTreeMap, + event: HookEventName, +) -> Result<(), ProtocolHookError> { + if fields.contains_key("hook_source") { + return Err(ProtocolHookError::UnexpectedHookSource); + } + if let Some(schema_version) = fields.remove("schema_version") + && schema_version != RedactedValue::String(PAYLOAD_SCHEMA_VERSION.into()) + { + return Err(ProtocolHookError::UnsupportedSchema); + } + if let Some(hook_event_name) = fields.remove("hook_event_name") + && hook_event_name != RedactedValue::String(event.as_str().into()) + { + return Err(ProtocolHookError::MismatchedEvent); + } + Ok(()) +} + +fn take_string( + fields: &mut BTreeMap, + key: &'static str, +) -> Result { + match fields.remove(key) { + Some(RedactedValue::String(value)) => Ok(value), + Some(_) => Err(ProtocolHookError::InvalidPayloadField(key)), + None => Err(ProtocolHookError::MissingPayloadField(key)), + } +} + +fn take_value( + fields: &mut BTreeMap, + key: &'static str, +) -> Result { + fields + .remove(key) + .ok_or(ProtocolHookError::MissingPayloadField(key)) +} + +fn take_enum( + fields: &mut BTreeMap, + key: &'static str, +) -> Result { + let value = take_value(fields, key)?; + serde_json::from_value( + serde_json::to_value(value).map_err(|_| ProtocolHookError::InvalidPayloadField(key))?, + ) + .map_err(|_| ProtocolHookError::InvalidPayloadField(key)) +} + +fn take_optional_enum( + fields: &mut BTreeMap, + key: &'static str, +) -> Result, ProtocolHookError> { + let Some(value) = fields.remove(key) else { + return Ok(None); + }; + serde_json::from_value( + serde_json::to_value(value).map_err(|_| ProtocolHookError::InvalidPayloadField(key))?, + ) + .map(Some) + .map_err(|_| ProtocolHookError::InvalidPayloadField(key)) +} + +fn value_from_protocol(value: &prost_types::Value) -> Result { + match value.kind.as_ref() { + Some(Kind::NullValue(_)) => Ok(RedactedValue::Null), + Some(Kind::NumberValue(value)) => serde_json::Number::from_f64(*value) + .map(RedactedValue::Number) + .ok_or(ProtocolHookError::InvalidNumber), + Some(Kind::StringValue(value)) => Ok(RedactedValue::String(value.clone())), + Some(Kind::BoolValue(value)) => Ok(RedactedValue::Bool(*value)), + Some(Kind::StructValue(value)) => value + .fields + .iter() + .map(|(key, value)| Ok((key.clone(), value_from_protocol(value)?))) + .collect::, _>>() + .map(RedactedValue::Object), + Some(Kind::ListValue(value)) => value + .values + .iter() + .map(value_from_protocol) + .collect::, _>>() + .map(RedactedValue::Array), + None => Err(ProtocolHookError::InvalidValue), + } +} + +#[derive(Clone, Copy, Debug, thiserror::Error)] +pub(crate) enum ProtocolHookError { + #[error("missing invocation ID")] + InvalidInvocationId, + #[error("invalid tool-use ID")] + InvalidToolUseId, + #[error("tool-use ID does not match payload")] + MismatchedToolUseId, + #[error("unsupported payload schema")] + UnsupportedSchema, + #[error("unknown hook event")] + UnknownEvent, + #[error("payload event does not match action")] + MismatchedEvent, + #[error("missing hook payload")] + MissingPayload, + #[error("hook payload exceeds the size limit")] + OversizedPayload, + #[error("source-neutral payload included hook_source")] + UnexpectedHookSource, + #[error("unknown hook payload field")] + UnknownPayloadField, + #[error("missing hook payload field {0}")] + MissingPayloadField(&'static str), + #[error("invalid hook payload field {0}")] + InvalidPayloadField(&'static str), + #[error("invalid protocol value")] + InvalidValue, + #[error("invalid protocol number")] + InvalidNumber, +} + +impl ProtocolHookError { + fn category(self) -> &'static str { + match self { + Self::InvalidInvocationId => "invalid_invocation_id", + Self::InvalidToolUseId => "invalid_tool_use_id", + Self::MismatchedToolUseId => "mismatched_tool_use_id", + Self::UnsupportedSchema => "unsupported_schema", + Self::UnknownEvent => "unknown_event", + Self::MismatchedEvent => "mismatched_event", + Self::MissingPayload => "missing_payload", + Self::OversizedPayload => "oversized_payload", + Self::UnexpectedHookSource => "unexpected_hook_source", + Self::UnknownPayloadField => "unknown_payload_field", + Self::MissingPayloadField(_) => "missing_payload_field", + Self::InvalidPayloadField(_) => "invalid_payload_field", + Self::InvalidValue => "invalid_value", + Self::InvalidNumber => "invalid_number", + } + } +} + +#[cfg(test)] +#[path = "protocol_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/hooks/protocol_tests.rs b/app/src/ai/agent_sdk/hooks/protocol_tests.rs new file mode 100644 index 00000000000..eac1a3de80b --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/protocol_tests.rs @@ -0,0 +1,261 @@ +use std::collections::BTreeMap; +use std::time::{Duration, SystemTime}; + +use prost_types::value::Kind; +use prost_types::{Struct, Value}; +use warp_multi_agent_api::oz_hook_result::{Outcome, ResolvedAction}; +use warp_multi_agent_api::{OzHookEvent as ProtocolEvent, RunOzHook}; + +use super::*; +use crate::ai::agent_sdk::hooks::runtime::{ + HookFailureCategory, HookInvocationDiagnostic, HookInvocationResult, +}; +use crate::ai::agent_sdk::hooks::{HookConfigSource, HookEventName}; + +fn string(value: &str) -> Value { + Value { + kind: Some(Kind::StringValue(value.into())), + } +} + +fn object(fields: impl IntoIterator) -> Value { + Value { + kind: Some(Kind::StructValue(Struct { + fields: fields + .into_iter() + .map(|(key, value)| (key.into(), value)) + .collect(), + })), + } +} + +fn common_fields() -> BTreeMap { + [ + ("session_id", string("session")), + ("run_id", string("run")), + ("conversation_id", string("conversation")), + ("cwd", string("/workspace")), + ("model", string("model")), + ("permission_mode", string("supervised")), + ] + .into_iter() + .map(|(key, value)| (key.into(), value)) + .collect() +} + +fn action( + event: ProtocolEvent, + event_fields: impl IntoIterator, +) -> RunOzHook { + let mut fields = common_fields(); + fields.extend( + event_fields + .into_iter() + .map(|(key, value)| (key.into(), value)), + ); + let tool_use_id = matches!( + event, + ProtocolEvent::PreToolUse | ProtocolEvent::PostToolUse + ) + .then(|| "tool-use".into()) + .unwrap_or_default(); + RunOzHook { + invocation_id: "invocation".into(), + tool_use_id, + event: event.into(), + schema_version: PAYLOAD_SCHEMA_VERSION.into(), + redacted_payload: Some(Struct { fields }), + } +} + +#[test] +fn oz_hooks_protocol_parses_all_seven_protocol_events() { + let cases = [ + ( + action(ProtocolEvent::SessionStart, [("source", string("startup"))]), + HookEventName::SessionStart, + ), + ( + action(ProtocolEvent::SessionEnd, [("reason", string("completed"))]), + HookEventName::SessionEnd, + ), + ( + action( + ProtocolEvent::UserPromptSubmit, + [("prompt", string("hello"))], + ), + HookEventName::UserPromptSubmit, + ), + ( + action(ProtocolEvent::Stop, [("turn_status", string("completed"))]), + HookEventName::Stop, + ), + ( + action( + ProtocolEvent::PreToolUse, + [ + ("tool_name", string("run_shell_command")), + ("tool_use_id", string("tool-use")), + ("tool_input", object([("command", string("pwd"))])), + ], + ), + HookEventName::PreToolUse, + ), + ( + action( + ProtocolEvent::PostToolUse, + [ + ("tool_name", string("run_shell_command")), + ("tool_use_id", string("tool-use")), + ("tool_input", object([("command", string("pwd"))])), + ("tool_response", object([("exit_code", string("0"))])), + ], + ), + HookEventName::PostToolUse, + ), + ( + action(ProtocolEvent::PreCompact, [("trigger", string("auto"))]), + HookEventName::PreCompact, + ), + ]; + + for (action, expected) in cases { + let event = event_from_protocol(&action).expect("event should parse"); + assert_eq!(event.payload.event_name(), expected); + } +} + +#[test] +fn oz_hooks_protocol_rejects_unknown_mismatched_and_source_specific_fields() { + let mut unknown = action(ProtocolEvent::Stop, [("turn_status", string("idle"))]); + unknown + .redacted_payload + .as_mut() + .unwrap() + .fields + .insert("unknown".into(), string("value")); + assert!(matches!( + event_from_protocol(&unknown), + Err(ProtocolHookError::UnknownPayloadField) + )); + + let mut mismatched = action(ProtocolEvent::Stop, [("turn_status", string("idle"))]); + mismatched + .redacted_payload + .as_mut() + .unwrap() + .fields + .insert("hook_event_name".into(), string("PreCompact")); + assert!(matches!( + event_from_protocol(&mismatched), + Err(ProtocolHookError::MismatchedEvent) + )); + + let mut sourced = action(ProtocolEvent::Stop, [("turn_status", string("idle"))]); + sourced + .redacted_payload + .as_mut() + .unwrap() + .fields + .insert("hook_source".into(), string("project")); + assert!(matches!( + event_from_protocol(&sourced), + Err(ProtocolHookError::UnexpectedHookSource) + )); +} + +#[test] +fn oz_hooks_protocol_rejects_unsupported_schema_and_tool_use_mismatch() { + let mut unsupported = action(ProtocolEvent::Stop, [("turn_status", string("idle"))]); + unsupported.schema_version = "future".into(); + assert!(matches!( + event_from_protocol(&unsupported), + Err(ProtocolHookError::UnsupportedSchema) + )); + + let mut mismatched = action( + ProtocolEvent::PreToolUse, + [ + ("tool_name", string("run_shell_command")), + ("tool_use_id", string("different")), + ("tool_input", object([])), + ], + ); + mismatched.tool_use_id = "tool-use".into(); + assert!(matches!( + event_from_protocol(&mismatched), + Err(ProtocolHookError::MismatchedToolUseId) + )); +} + +fn diagnostic( + result: HookInvocationResult, + failure_category: Option, +) -> HookInvocationDiagnostic { + HookInvocationDiagnostic { + event: HookEventName::PreToolUse, + source: HookConfigSource::Project, + config_path: ".warp/hooks.json".into(), + definition_hash: "hash".into(), + matcher: None, + started_at: SystemTime::UNIX_EPOCH, + finished_at: SystemTime::UNIX_EPOCH, + duration: Duration::ZERO, + result, + exit_code: None, + output_truncated: false, + failure_category, + } +} + +#[test] +fn oz_hooks_protocol_maps_continue_deny_failed_and_cancelled_results() { + let action = action( + ProtocolEvent::PreToolUse, + [ + ("tool_name", string("run_shell_command")), + ("tool_use_id", string("tool-use")), + ("tool_input", object([])), + ], + ); + + assert!(matches!( + result_for_observation(&action, &[]).outcome, + Some(Outcome::Continue(_)) + )); + assert!(matches!( + result_for_pre_tool( + &action, + OzPreToolUseDecision::Deny { + reason: "no".into(), + source: HookConfigSource::Project, + diagnostics: vec![], + } + ) + .outcome, + Some(Outcome::Deny(_)) + )); + let failed = result_for_observation( + &action, + &[diagnostic( + HookInvocationResult::Continued, + Some(HookFailureCategory::Timeout), + )], + ); + assert!(matches!( + failed.outcome, + Some(Outcome::Failed(ref outcome)) + if outcome.resolved_action == ResolvedAction::Continue as i32 + )); + assert!(matches!( + result_for_observation( + &action, + &[diagnostic( + HookInvocationResult::Cancelled, + Some(HookFailureCategory::Cancelled), + )], + ) + .outcome, + Some(Outcome::Cancelled(_)) + )); +} diff --git a/app/src/ai/agent_sdk/hooks/redaction.rs b/app/src/ai/agent_sdk/hooks/redaction.rs index 1b5aad99987..f659b1a095f 100644 --- a/app/src/ai/agent_sdk/hooks/redaction.rs +++ b/app/src/ai/agent_sdk/hooks/redaction.rs @@ -1,10 +1,10 @@ use std::collections::BTreeMap; use regex::Regex; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use serde_json::Value; -#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(untagged)] pub(crate) enum RedactedValue { Null, @@ -67,7 +67,7 @@ impl From for RedactedValue { } } -#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] pub(crate) struct TruncationMetadata { pub(crate) truncated: bool, pub(crate) original_bytes: usize, @@ -140,26 +140,6 @@ impl HookRedactor { } } - pub(crate) fn redact_json_preview( - &self, - value: &RedactedValue, - maximum_bytes: usize, - ) -> RedactedValue { - if value.serialized_len() <= maximum_bytes { - return value.clone(); - } - RedactedValue::object([ - ("truncated", RedactedValue::Bool(true)), - ( - "original_bytes", - RedactedValue::Number(value.serialized_len().into()), - ), - ( - "preview", - RedactedValue::redacted("size_limit", value.serialized_len()), - ), - ]) - } } pub(crate) fn truncate_utf8(input: &str, maximum_bytes: usize) -> String { diff --git a/app/src/ai/agent_sdk/hooks/runtime.rs b/app/src/ai/agent_sdk/hooks/runtime.rs index cd79621c048..bd67a9fd619 100644 --- a/app/src/ai/agent_sdk/hooks/runtime.rs +++ b/app/src/ai/agent_sdk/hooks/runtime.rs @@ -29,6 +29,11 @@ pub(crate) struct OzHookEvent { pub(crate) payload: HookPayloadTemplate, } +enum ReaderFailure { + Overflow, + Io, +} + #[derive(Clone, Debug)] pub(crate) struct OzPreToolUseEvent(OzHookEvent); @@ -45,6 +50,7 @@ impl OzPreToolUseEvent { pub(crate) enum OzHookCancellationScope { Session, Invocation(String), + #[allow(dead_code)] Tool(String), } @@ -82,6 +88,7 @@ pub(crate) enum HookFailureCategory { Timeout, Cancelled, OutputOverflow, + OutputRead, InvalidUtf8, NonZeroExit, InvalidDecision, @@ -218,12 +225,15 @@ impl OzHookRuntimeService { Err(failure) => { diagnostic.exit_code = failure.exit_code; diagnostic.failure_category = Some(failure.category); + diagnostic.output_truncated = + failure.category == HookFailureCategory::OutputOverflow; diagnostic.result = match failure.category { HookFailureCategory::Timeout => HookInvocationResult::TimedOut, HookFailureCategory::Cancelled => HookInvocationResult::Cancelled, HookFailureCategory::Spawn | HookFailureCategory::Stdin | HookFailureCategory::OutputOverflow + | HookFailureCategory::OutputRead | HookFailureCategory::InvalidUtf8 | HookFailureCategory::NonZeroExit | HookFailureCategory::InvalidDecision @@ -248,6 +258,23 @@ impl OzHookRuntimeService { } } } + log::info!( + "Oz hook invocation: event={} source={} config_path={} definition_hash={} \ + matcher_present={} started_at={:?} finished_at={:?} duration_ms={} result={:?} \ + exit_code={:?} output_truncated={} failure_category={:?}", + diagnostic.event, + diagnostic.source.as_str(), + diagnostic.config_path.display(), + diagnostic.definition_hash, + diagnostic.matcher.is_some(), + diagnostic.started_at, + diagnostic.finished_at, + diagnostic.duration.as_millis(), + diagnostic.result, + diagnostic.exit_code, + diagnostic.output_truncated, + diagnostic.failure_category + ); outcome.diagnostics.push(diagnostic); } self.remove_pending(&event.invocation_id); @@ -410,11 +437,14 @@ async fn run_command( }); } }; - if stdin_task.await.is_err() { - return Err(CommandFailure { - category: HookFailureCategory::Stdin, - exit_code: status.code(), - }); + match stdin_task.await { + Ok(Ok(())) => {} + Ok(Err(_)) | Err(_) => { + return Err(CommandFailure { + category: HookFailureCategory::Stdin, + exit_code: status.code(), + }); + } } let stdout = join_output(stdout_task, status.code()).await?; let stderr = join_output(stderr_task, status.code()).await?; @@ -479,18 +509,21 @@ fn hook_environment(payload: &HookPayloadTemplate) -> HashMap, -) -> JoinHandle, ()>> { +) -> JoinHandle, ReaderFailure>> { tokio::spawn(async move { let mut output = Vec::new(); let mut buffer = [0_u8; 8192]; loop { - let read = reader.read(&mut buffer).await.map_err(|_| ())?; + let read = reader + .read(&mut buffer) + .await + .map_err(|_| ReaderFailure::Io)?; if read == 0 { return Ok(output); } if output.len() + read > MAX_OUTPUT_BYTES { let _ = overflow.send(()); - return Err(()); + return Err(ReaderFailure::Overflow); } output.extend_from_slice(&buffer[..read]); } @@ -498,7 +531,7 @@ fn spawn_bounded_reader( } async fn join_output( - task: JoinHandle, ()>>, + task: JoinHandle, ReaderFailure>>, exit_code: Option, ) -> Result { let bytes = task @@ -507,8 +540,11 @@ async fn join_output( category: HookFailureCategory::OutputOverflow, exit_code, })? - .map_err(|_| CommandFailure { - category: HookFailureCategory::OutputOverflow, + .map_err(|failure| CommandFailure { + category: match failure { + ReaderFailure::Overflow => HookFailureCategory::OutputOverflow, + ReaderFailure::Io => HookFailureCategory::OutputRead, + }, exit_code, })?; String::from_utf8(bytes).map_err(|_| CommandFailure { diff --git a/app/src/ai/blocklist/block/view_impl.rs b/app/src/ai/blocklist/block/view_impl.rs index 690b85ec7f7..24b69f8c426 100644 --- a/app/src/ai/blocklist/block/view_impl.rs +++ b/app/src/ai/blocklist/block/view_impl.rs @@ -1391,7 +1391,8 @@ impl AIAgentInput { | AIAgentInput::MessagesReceivedFromAgents { .. } | AIAgentInput::EventsFromAgents { .. } | AIAgentInput::PassiveSuggestionResult { .. } - | AIAgentInput::OrchestrationConfigUpdate { .. } => None, + | AIAgentInput::OrchestrationConfigUpdate { .. } + | AIAgentInput::OzHookResult(_) => None, } } } diff --git a/app/src/ai/blocklist/block/view_impl/common.rs b/app/src/ai/blocklist/block/view_impl/common.rs index 5dd5e1bde50..523c1ae4df9 100644 --- a/app/src/ai/blocklist/block/view_impl/common.rs +++ b/app/src/ai/blocklist/block/view_impl/common.rs @@ -3655,7 +3655,8 @@ pub(super) fn query_prefix_highlight_len( | AIAgentInput::MessagesReceivedFromAgents { .. } | AIAgentInput::EventsFromAgents { .. } | AIAgentInput::PassiveSuggestionResult { .. } - | AIAgentInput::OrchestrationConfigUpdate { .. } => None, + | AIAgentInput::OrchestrationConfigUpdate { .. } + | AIAgentInput::OzHookResult(_) => None, } } } diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index e115c10f2ce..d9594556844 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -57,6 +57,16 @@ use crate::ai::agent::{ use crate::ai::agent_events::AgentMessageEventMetadata; #[cfg(not(target_family = "wasm"))] use crate::ai::agent_sdk::ClaudeHarness; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::protocol::{ + event_from_protocol, failed_result, result_for_observation, result_for_pre_tool, +}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::runtime::OzHookCancellationScope; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::runtime::OzPreToolUseEvent; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::{HookEventName, OzHookSession, PAYLOAD_SCHEMA_VERSION}; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, @@ -367,6 +377,19 @@ pub struct BlocklistAIController { Option, )>, >, + #[cfg(not(target_family = "wasm"))] + oz_hook_session: Option, + #[cfg(not(target_family = "wasm"))] + pending_oz_hook_results: HashMap>, + #[cfg(not(target_family = "wasm"))] + oz_hook_compatible_streams: HashSet, + #[cfg(not(target_family = "wasm"))] + oz_hook_invocations: HashSet<(AIConversationId, String)>, + #[cfg(not(target_family = "wasm"))] + oz_hook_results_by_invocation: + HashMap<(AIConversationId, String), warp_multi_agent_api::OzHookResult>, + #[cfg(not(target_family = "wasm"))] + cancelled_oz_hook_invocations: HashSet<(AIConversationId, String)>, } enum InputQueryType { @@ -654,6 +677,18 @@ impl BlocklistAIController { pending_local_claude_wakes: HashMap::new(), pending_passive_follow_ups: HashSet::new(), pending_passive_suggestion_results: HashMap::new(), + #[cfg(not(target_family = "wasm"))] + oz_hook_session: None, + #[cfg(not(target_family = "wasm"))] + pending_oz_hook_results: HashMap::new(), + #[cfg(not(target_family = "wasm"))] + oz_hook_compatible_streams: HashSet::new(), + #[cfg(not(target_family = "wasm"))] + oz_hook_invocations: HashSet::new(), + #[cfg(not(target_family = "wasm"))] + oz_hook_results_by_invocation: HashMap::new(), + #[cfg(not(target_family = "wasm"))] + cancelled_oz_hook_invocations: HashSet::new(), } } @@ -2384,6 +2419,128 @@ impl BlocklistAIController { }); } + #[cfg(not(target_family = "wasm"))] + pub(crate) fn set_oz_hook_session(&mut self, session: Option) { + self.oz_hook_session = session; + } + + #[cfg(not(target_family = "wasm"))] + fn execute_protocol_oz_hook( + &mut self, + conversation_id: AIConversationId, + action: warp_multi_agent_api::RunOzHook, + ctx: &mut ModelContext, + ) { + let invocation_key = (conversation_id, action.invocation_id.clone()); + if !self.oz_hook_invocations.insert(invocation_key.clone()) { + if let Some(result) = self + .oz_hook_results_by_invocation + .get(&invocation_key) + .cloned() + { + self.pending_oz_hook_results + .entry(conversation_id) + .or_default() + .push(result); + self.send_pending_oz_hook_results(conversation_id, ctx); + } + return; + } + let Some(session) = &self.oz_hook_session else { + return; + }; + let event = match event_from_protocol(&action) { + Ok(event) => event, + Err(error) => { + let result = failed_result(&action, error); + self.oz_hook_results_by_invocation + .insert(invocation_key, result.clone()); + self.pending_oz_hook_results + .entry(conversation_id) + .or_default() + .push(result); + return; + } + }; + let runtime = Arc::clone(&session.runtime); + let event_name = event.payload.event_name(); + ctx.spawn( + async move { + if event_name == HookEventName::PreToolUse { + let event = OzPreToolUseEvent::new(event) + .expect("event name was validated before constructing pre-tool event"); + result_for_pre_tool(&action, runtime.pre_tool_use(event).await) + } else { + let observation = runtime.observe(event).await; + result_for_observation(&action, &observation.diagnostics) + } + }, + move |me, result, ctx| { + if me.cancelled_oz_hook_invocations.contains(&invocation_key) { + return; + } + me.oz_hook_results_by_invocation + .insert(invocation_key, result.clone()); + me.pending_oz_hook_results + .entry(conversation_id) + .or_default() + .push(result); + me.send_pending_oz_hook_results(conversation_id, ctx); + }, + ); + } + + #[cfg(not(target_family = "wasm"))] + fn send_pending_oz_hook_results( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if self + .in_flight_response_streams + .has_active_stream_for_conversation(conversation_id, ctx) + { + return; + } + let Some(results) = self.pending_oz_hook_results.remove(&conversation_id) else { + return; + }; + let Some(task_id) = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .map(|conversation| conversation.get_root_task_id().clone()) + else { + return; + }; + let scope = ResolvedTeamScope::from_scope(&self.team_context(ctx)); + let inputs = results + .iter() + .cloned() + .map(AIAgentInput::OzHookResult) + .collect(); + if let Err(error) = self.send_request_input( + RequestInput::for_task( + inputs, + task_id, + &self.active_session, + self.get_current_response_initiator(), + conversation_id, + self.terminal_surface_id, + &scope, + ctx, + ), + None, + RecoveryBudget::fresh(), + false, + ctx, + ) { + report_error!(error.context("Failed to submit Oz hook results")); + self.pending_oz_hook_results + .entry(conversation_id) + .or_default() + .extend(results); + } + } + #[cfg(test)] pub fn get_ambient_agent_task_id(&self) -> Option { self.ambient_agent_task_id @@ -2564,6 +2721,10 @@ impl BlocklistAIController { ); request_params.parent_agent_id = parent_agent_id; request_params.agent_name = agent_name; + #[cfg(not(target_family = "wasm"))] + if let Some(session) = &self.oz_hook_session { + request_params.oz_hook_context = Some(session.protocol_context.clone()); + } let server_conversation_token_for_identifiers = conversation_data.server_conversation_token.clone(); @@ -2977,6 +3138,27 @@ impl BlocklistAIController { }; match event { warp_multi_agent_api::response_event::Type::Init(init_event) => { + #[cfg(not(target_family = "wasm"))] + if self.oz_hook_session.is_some() { + if init_event + .supported_oz_hook_payload_schema_versions + .iter() + .any(|version| version == PAYLOAD_SCHEMA_VERSION) + { + self.oz_hook_compatible_streams.insert(stream_id.clone()); + } else { + report_error!( + "Oz lifecycle hook schema negotiation failed", + extra: { "stream_id" => ?stream_id } + ); + self.cancel_request( + &stream_id, + CancellationReason::ManuallyCancelled, + ctx, + ); + return; + } + } history_model.update(ctx, |history_model, ctx| { history_model.initialize_output_for_response_stream( &stream_id, @@ -3009,7 +3191,34 @@ impl BlocklistAIController { ); } warp_multi_agent_api::response_event::Type::ClientActions(actions) => { - let client_actions = actions.actions; + let mut client_actions = Vec::new(); + for mut client_action in actions.actions { + #[cfg(not(target_family = "wasm"))] + if let Some( + warp_multi_agent_api::client_action::Action::RunOzHook( + action, + ), + ) = client_action.action.take() + { + if !self.oz_hook_compatible_streams.contains(&stream_id) { + report_error!( + "Received Oz hook action before successful schema negotiation", + extra: { "stream_id" => ?stream_id } + ); + self.cancel_request( + &stream_id, + CancellationReason::ManuallyCancelled, + ctx, + ); + return; + } + self.execute_protocol_oz_hook(conversation_id, action, ctx); + continue; + } + if client_action.action.is_some() { + client_actions.push(client_action); + } + } let skill_path_origin = SessionContext::from_session( self.active_session.as_ref(ctx), ctx, @@ -3136,6 +3345,24 @@ impl BlocklistAIController { }; let history_model = BlocklistAIHistoryModel::handle(ctx); + #[cfg(not(target_family = "wasm"))] + if cancellation.is_some() { + if let Some(session) = &self.oz_hook_session { + for invocation_key in self + .oz_hook_invocations + .iter() + .filter(|(id, _)| *id == conversation_id) + { + session.runtime.cancel(OzHookCancellationScope::Invocation( + invocation_key.1.clone(), + )); + self.cancelled_oz_hook_invocations + .insert(invocation_key.clone()); + } + } + self.pending_oz_hook_results.remove(&conversation_id); + self.oz_hook_compatible_streams.remove(&stream_id); + } let Some(conversation) = history_model.as_ref(ctx).conversation(&conversation_id) else { log::warn!("Conversation not found."); @@ -3226,6 +3453,11 @@ impl BlocklistAIController { // Cancelled streams will handle pending_response_stream updates synchronously. if cancellation.is_none() { self.in_flight_response_streams.cleanup_stream(&stream_id); + #[cfg(not(target_family = "wasm"))] + { + self.oz_hook_compatible_streams.remove(&stream_id); + self.send_pending_oz_hook_results(conversation_id, ctx); + } // Now that the stream is cleaned up, re-check for pending // orchestration events that couldn't be drained earlier. diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 4e1a9753b60..ded5c0dc056 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -315,6 +315,7 @@ fn mock_response_stream_updates_history_through_controller() { request_id: "test-request".to_string(), conversation_id: "test-server-conversation".to_string(), run_id: String::new(), + supported_oz_hook_payload_schema_versions: vec![], })), }, ctx, diff --git a/app/src/ai/blocklist/history_model_tests.rs b/app/src/ai/blocklist/history_model_tests.rs index 337283a28c3..8471a815720 100644 --- a/app/src/ai/blocklist/history_model_tests.rs +++ b/app/src/ai/blocklist/history_model_tests.rs @@ -3095,6 +3095,7 @@ fn test_initialize_output_for_response_stream_persists_updated_conversation_stat request_id: "request-1".to_string(), conversation_id: server_token.clone(), run_id: run_id.clone(), + supported_oz_hook_payload_schema_versions: vec![], }, ctx, ); @@ -3361,6 +3362,7 @@ fn test_find_by_token_after_initialize_output_for_response_stream() { request_id: String::new(), conversation_id: server_token_str.clone(), run_id: String::new(), + supported_oz_hook_payload_schema_versions: vec![], }, ctx, ); diff --git a/app/src/ai/blocklist/permissions.rs b/app/src/ai/blocklist/permissions.rs index 0c23dcaefab..f39f84aa63e 100644 --- a/app/src/ai/blocklist/permissions.rs +++ b/app/src/ai/blocklist/permissions.rs @@ -1231,12 +1231,12 @@ fn command_for_execution_predicates(command: &str, escape_char: EscapeChar) -> S /// and must never be auto-written regardless of user autonomy settings. /// Returns `None` if no paths are protected. fn check_protected_write_paths(paths: &[PathBuf]) -> Option { - // MCP config files are always protected from auto-write to prevent security risks - // from injecting arbitrary context into the agent. - if paths - .iter() - .any(|p| mcp_provider_from_file_path(p).is_some()) - { + // Executable agent configuration is always protected from auto-write because it can inject + // arbitrary context or launch commands in a later request. + if paths.iter().any(|p| { + mcp_provider_from_file_path(p).is_some() + || p.ends_with(std::path::Path::new(".warp/hooks.json")) + }) { Some(FileWritePermission::Denied( FileWritePermissionDeniedReason::ProtectedPath, )) diff --git a/app/src/ai/blocklist/permissions_tests.rs b/app/src/ai/blocklist/permissions_tests.rs index 9d9f7dad608..404af0c5856 100644 --- a/app/src/ai/blocklist/permissions_tests.rs +++ b/app/src/ai/blocklist/permissions_tests.rs @@ -496,6 +496,7 @@ fn test_can_write_files_mcp_config_always_denied() { PathBuf::from("/project/.mcp.json"), PathBuf::from("/project/.warp/.mcp.json"), PathBuf::from("/project/.codex/config.toml"), + PathBuf::from("/project/.warp/hooks.json"), ]; for path in mcp_config_paths { diff --git a/app/src/ai/blocklist/persistence.rs b/app/src/ai/blocklist/persistence.rs index cc73d53ac54..9e190d71df7 100644 --- a/app/src/ai/blocklist/persistence.rs +++ b/app/src/ai/blocklist/persistence.rs @@ -102,7 +102,8 @@ impl TryFrom<&AIAgentInput> for PersistedAIInputType { | AIAgentInput::StartFromAmbientRunPrompt { .. } | AIAgentInput::MessagesReceivedFromAgents { .. } | AIAgentInput::EventsFromAgents { .. } - | AIAgentInput::OrchestrationConfigUpdate { .. } => Err(anyhow::anyhow!( + | AIAgentInput::OrchestrationConfigUpdate { .. } + | AIAgentInput::OzHookResult(_) => Err(anyhow::anyhow!( "This input type is not persisted. Only Query inputs are persisted for up-arrow history." )), } diff --git a/app/src/server/telemetry/events.rs b/app/src/server/telemetry/events.rs index 23b17483cf2..c1a7da872c3 100644 --- a/app/src/server/telemetry/events.rs +++ b/app/src/server/telemetry/events.rs @@ -998,6 +998,7 @@ pub enum AIAgentInput { EventsFromAgents { event_count: usize }, PassiveSuggestionResult, OrchestrationConfigUpdate, + OzHookResult, } impl From for AIAgentInput { @@ -1038,6 +1039,7 @@ impl From for AIAgentInput { }, FullAIAgentInput::PassiveSuggestionResult { .. } => Self::PassiveSuggestionResult, FullAIAgentInput::OrchestrationConfigUpdate { .. } => Self::OrchestrationConfigUpdate, + FullAIAgentInput::OzHookResult(_) => Self::OzHookResult, } } } diff --git a/app/src/terminal/shared_session/replay_agent_conversations.rs b/app/src/terminal/shared_session/replay_agent_conversations.rs index 866a377f59c..02c2c834b17 100644 --- a/app/src/terminal/shared_session/replay_agent_conversations.rs +++ b/app/src/terminal/shared_session/replay_agent_conversations.rs @@ -89,6 +89,7 @@ pub fn reconstruct_response_events_from_conversations( // Shared session replays don't need a run_id; the empty // string is filtered to None by initialize_output_for_response_stream. run_id: String::new(), + supported_oz_hook_payload_schema_versions: vec![], }, )), }); diff --git a/app/src/terminal/view/shared_session/view_impl_tests.rs b/app/src/terminal/view/shared_session/view_impl_tests.rs index 0397c74dac9..01ec4b59c01 100644 --- a/app/src/terminal/view/shared_session/view_impl_tests.rs +++ b/app/src/terminal/view/shared_session/view_impl_tests.rs @@ -2213,6 +2213,7 @@ fn test_shared_followup_on_existing_conversation_converts_user_query_input() { request_id: request_id.to_string(), conversation_id: conversation_token.to_string(), run_id: String::new(), + supported_oz_hook_payload_schema_versions: vec![], }, )), }; diff --git a/crates/warp_cli/src/agent.rs b/crates/warp_cli/src/agent.rs index 74135feb27c..9f6051052e8 100644 --- a/crates/warp_cli/src/agent.rs +++ b/crates/warp_cli/src/agent.rs @@ -129,6 +129,93 @@ impl FromStr for RepositoryHeadOverride { } } +const OZ_LIFECYCLE_HOOKS_CONTEXT_MAX_BYTES: usize = 64 * 1024; +const OZ_LIFECYCLE_HOOKS_MAX_TRUST_RECORDS: usize = 64; +const OZ_HOOK_PAYLOAD_SCHEMA_VERSION: &str = "warp.oz_hook.v1"; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OzLifecycleHooksContext { + pub required: bool, + pub supported_payload_schema_versions: Vec, + pub project_trust: Vec, +} + +impl OzLifecycleHooksContext { + fn validate(&self) -> Result<(), String> { + if !self.required { + return Err("required must be true".into()); + } + if self.supported_payload_schema_versions.is_empty() { + return Err("supported_payload_schema_versions must not be empty".into()); + } + if self + .supported_payload_schema_versions + .iter() + .any(|version| version != OZ_HOOK_PAYLOAD_SCHEMA_VERSION) + { + return Err(format!( + "unsupported Oz hook payload schema version; expected {OZ_HOOK_PAYLOAD_SCHEMA_VERSION}" + )); + } + if self.project_trust.len() > OZ_LIFECYCLE_HOOKS_MAX_TRUST_RECORDS { + return Err(format!( + "project_trust exceeds {OZ_LIFECYCLE_HOOKS_MAX_TRUST_RECORDS} records" + )); + } + for trust in &self.project_trust { + trust.validate()?; + } + Ok(()) + } +} + +impl FromStr for OzLifecycleHooksContext { + type Err = String; + + fn from_str(value: &str) -> Result { + if value.len() > OZ_LIFECYCLE_HOOKS_CONTEXT_MAX_BYTES { + return Err(format!( + "Oz lifecycle hooks context exceeds {OZ_LIFECYCLE_HOOKS_CONTEXT_MAX_BYTES} bytes" + )); + } + let context = serde_json::from_str::(value) + .map_err(|error| format!("invalid Oz lifecycle hooks context JSON: {error}"))?; + context.validate()?; + Ok(context) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct OzProjectHookTrust { + pub git_root: PathBuf, + pub config_path: PathBuf, + pub sha256: String, +} + +impl OzProjectHookTrust { + fn validate(&self) -> Result<(), String> { + if self.git_root.as_os_str().is_empty() { + return Err("project_trust git_root must not be empty".into()); + } + if self.config_path.as_os_str().is_empty() { + return Err("project_trust config_path must not be empty".into()); + } + if self.sha256.len() != 64 + || !self + .sha256 + .bytes() + .all(|byte| byte.is_ascii_digit() || matches!(byte, b'a'..=b'f')) + { + return Err( + "project_trust sha256 must be an exact 64-character lowercase hexadecimal hash" + .into(), + ); + } + Ok(()) + } +} #[derive(Debug, Clone, PartialEq, Eq)] pub enum Prompt { PlainText(String), @@ -596,6 +683,15 @@ pub struct RunAgentArgs { /// Remove the origin remote from environment repositories after setup. #[arg(long = "remove-repository-origins", requires = "task_id", hide = true)] pub remove_repository_origins: bool, + + /// Server-authenticated lifecycle hook capability and project trust for embedded Oz. + #[arg( + long = "oz-lifecycle-hooks-context", + value_name = "JSON", + requires = "task_id", + hide = true + )] + pub oz_lifecycle_hooks_context: Option, } impl RunAgentArgs { diff --git a/crates/warp_cli/src/lib_tests.rs b/crates/warp_cli/src/lib_tests.rs index c6547c08a90..25a71acfa87 100644 --- a/crates/warp_cli/src/lib_tests.rs +++ b/crates/warp_cli/src/lib_tests.rs @@ -1099,6 +1099,52 @@ fn agent_run_rejects_skip_initial_turn_without_task_id() { ); } +#[test] +fn agent_run_accepts_strict_oz_lifecycle_hooks_context() { + let context = r#"{"required":true,"supported_payload_schema_versions":["warp.oz_hook.v1"],"project_trust":[{"git_root":"/workspace/repo","config_path":"/workspace/repo/.warp/hooks.json","sha256":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}]}"#; + let args = Args::try_parse_from([ + "warp", + "agent", + "run", + "--task-id", + "abc", + "--oz-lifecycle-hooks-context", + context, + ]) + .unwrap(); + let Some(Command::CommandLine(boxed_cmd)) = args.command else { + panic!("Expected `warp agent run` command"); + }; + let CliCommand::Agent(AgentCommand::Run(run_args)) = boxed_cmd.as_ref() else { + panic!("Expected `warp agent run` command"); + }; + + let context = run_args.oz_lifecycle_hooks_context.as_ref().unwrap(); + assert!(context.required); + assert_eq!(context.project_trust.len(), 1); +} + +#[test] +fn agent_run_rejects_invalid_oz_lifecycle_hooks_context() { + for context in [ + r#"{"required":false,"supported_payload_schema_versions":["warp.oz_hook.v1"],"project_trust":[]}"#, + r#"{"required":true,"supported_payload_schema_versions":[],"project_trust":[]}"#, + r#"{"required":true,"supported_payload_schema_versions":["future"],"project_trust":[]}"#, + r#"{"required":true,"supported_payload_schema_versions":["warp.oz_hook.v1"],"project_trust":[],"unknown":true}"#, + r#"{"required":true,"supported_payload_schema_versions":["warp.oz_hook.v1"],"project_trust":[{"git_root":"/workspace/repo","config_path":"/workspace/repo/.warp/hooks.json","sha256":"bad"}]}"#, + ] { + Args::try_parse_from([ + "warp", + "agent", + "run", + "--task-id", + "abc", + "--oz-lifecycle-hooks-context", + context, + ]) + .expect_err("invalid hook context must fail parsing"); + } +} #[test] fn agent_run_accepts_snapshot_flags() { let args = Args::try_parse_from([ diff --git a/specs/APP-4344/TECH.md b/specs/APP-4344/TECH.md index 6a47c48da1c..3f667cb45f5 100644 --- a/specs/APP-4344/TECH.md +++ b/specs/APP-4344/TECH.md @@ -491,6 +491,33 @@ Required worker changes are limited to: - reject a hook-enabled task when a backend cannot preserve the embedded runtime contract - add backend tests that prove worker-control-plane credentials are not inherited +The assignment carries an optional, strict `oz_lifecycle_hooks` object: + +```json +{ + "required": true, + "supported_payload_schema_versions": ["warp.oz_hook.v1"], + "project_trust": [ + { + "git_root": "/workspace/repository", + "config_path": "/workspace/repository/.warp/hooks.json", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000" + } + ] +} +``` + +Unknown fields, `required: false`, empty or unsupported schema versions, more than 64 trust +records, and serialized objects over 64 KiB are rejected. The object is rejected for non-Oz +harnesses. The 64 KiB transport bound keeps the single argument below Linux `MAX_ARG_STRLEN`; the +256 KiB config and hook-stdin limits are separate. + +The worker passes this object to embedded Oz as one non-secret +`--oz-lifecycle-hooks-context ` argument pair. It is never inherited through the process +environment. The argument uses a strict `OzLifecycleHooksContext` CLI type and is validated before +the session starts. Backends that cannot preserve the argument, task cancellation, or sandbox +placement reject the hook-enabled task. + Direct: - The embedded runtime runs under the task process and task workspace. - It inherits only task environment into Oz. From 8c89457b1eb750dbd4748f6cd0fb564b032aa4ac Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:40:00 +0000 Subject: [PATCH 04/11] Fix Oz hook runtime cross-platform lint --- app/src/ai/agent_sdk/hooks/permissions.rs | 3 +++ app/src/ai/agent_sdk/hooks/redaction.rs | 4 +++- app/src/ai/agent_sdk/hooks/runtime.rs | 5 +++-- app/src/ai/agent_sdk/hooks/runtime_tests.rs | 1 + app/src/ai/agent_sdk/hooks/trust.rs | 2 ++ 5 files changed, 12 insertions(+), 3 deletions(-) diff --git a/app/src/ai/agent_sdk/hooks/permissions.rs b/app/src/ai/agent_sdk/hooks/permissions.rs index 2371574e302..35ba0b22563 100644 --- a/app/src/ai/agent_sdk/hooks/permissions.rs +++ b/app/src/ai/agent_sdk/hooks/permissions.rs @@ -7,12 +7,14 @@ pub(crate) enum NativePermission { } #[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[allow(dead_code)] pub(crate) enum HookPermission { Continue, Deny, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[allow(dead_code)] pub(crate) enum ComposedPermission { DeniedByWarp, DeniedByHook, @@ -20,6 +22,7 @@ pub(crate) enum ComposedPermission { Prompt, } +#[allow(dead_code)] pub(crate) fn compose_permission( native: NativePermission, hook: HookPermission, diff --git a/app/src/ai/agent_sdk/hooks/redaction.rs b/app/src/ai/agent_sdk/hooks/redaction.rs index f659b1a095f..86bf874b304 100644 --- a/app/src/ai/agent_sdk/hooks/redaction.rs +++ b/app/src/ai/agent_sdk/hooks/redaction.rs @@ -16,6 +16,7 @@ pub(crate) enum RedactedValue { } impl RedactedValue { + #[allow(dead_code)] pub(crate) fn object( fields: impl IntoIterator, RedactedValue)>, ) -> Self { @@ -27,6 +28,7 @@ impl RedactedValue { ) } + #[allow(dead_code)] pub(crate) fn redacted(reason: &str, byte_count: usize) -> Self { Self::object([ ("redacted", Self::Bool(true)), @@ -139,7 +141,6 @@ impl HookRedactor { }), } } - } pub(crate) fn truncate_utf8(input: &str, maximum_bytes: usize) -> String { @@ -155,6 +156,7 @@ fn floor_utf8_boundary(input: &str, maximum_bytes: usize) -> usize { boundary } +#[allow(dead_code)] pub(crate) fn contains_prohibited_payload_key(value: &Value) -> bool { const PROHIBITED_KEYS: [&str; 9] = [ "environment", diff --git a/app/src/ai/agent_sdk/hooks/runtime.rs b/app/src/ai/agent_sdk/hooks/runtime.rs index bd67a9fd619..630681daeb0 100644 --- a/app/src/ai/agent_sdk/hooks/runtime.rs +++ b/app/src/ai/agent_sdk/hooks/runtime.rs @@ -2,10 +2,11 @@ use std::collections::HashMap; use std::ffi::OsString; use std::path::Path; use std::process::Stdio; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant, SystemTime}; +use std::sync::Mutex; +use std::time::{Duration, SystemTime}; use async_trait::async_trait; +use instant::Instant; use serde::Deserialize; use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWriteExt as _}; use tokio::process::Command; diff --git a/app/src/ai/agent_sdk/hooks/runtime_tests.rs b/app/src/ai/agent_sdk/hooks/runtime_tests.rs index ad873cac553..ad5f9279f30 100644 --- a/app/src/ai/agent_sdk/hooks/runtime_tests.rs +++ b/app/src/ai/agent_sdk/hooks/runtime_tests.rs @@ -1,4 +1,5 @@ use std::fs; +use std::sync::Arc; use serde_json::json; diff --git a/app/src/ai/agent_sdk/hooks/trust.rs b/app/src/ai/agent_sdk/hooks/trust.rs index 035d0a452fe..df484c2145c 100644 --- a/app/src/ai/agent_sdk/hooks/trust.rs +++ b/app/src/ai/agent_sdk/hooks/trust.rs @@ -14,6 +14,7 @@ pub(crate) trait HookTrustStore: Send + Sync { } #[derive(Default)] +#[allow(dead_code)] pub(crate) struct DenyProjectHookTrust; impl HookTrustStore for DenyProjectHookTrust { @@ -32,6 +33,7 @@ impl ExactHookTrustStore { self.trusted.write().unwrap().insert(key); } + #[allow(dead_code)] pub(crate) fn revoke(&self, key: &HookTrustKey) { self.trusted.write().unwrap().remove(key); } From 20bdf611875183f71bb49301d87dea32f529ce84 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 14:13:54 +0000 Subject: [PATCH 05/11] Stage local Oz hooks around client actions --- app/Cargo.toml | 1 + app/src/ai/agent_sdk/driver.rs | 54 ++- app/src/ai/agent_sdk/hooks/adapters.rs | 296 +++++++++++++++ app/src/ai/agent_sdk/hooks/adapters_tests.rs | 62 +++ app/src/ai/agent_sdk/hooks/mod.rs | 3 + app/src/ai/agent_sdk/hooks/permissions.rs | 4 - app/src/ai/agent_sdk/hooks/protocol.rs | 43 ++- app/src/ai/agent_sdk/hooks/protocol_tests.rs | 55 +++ app/src/ai/agent_sdk/hooks/redaction.rs | 2 - app/src/ai/agent_sdk/hooks/runtime.rs | 178 +++++++-- app/src/ai/agent_sdk/hooks/runtime_tests.rs | 52 ++- app/src/ai/agent_sdk/hooks/trust.rs | 171 ++++++++- app/src/ai/agent_sdk/hooks/trust_tests.rs | 92 +++++ app/src/ai/blocklist/action_model.rs | 34 ++ app/src/ai/blocklist/action_model/execute.rs | 374 +++++++++++++++---- app/src/ai/blocklist/controller.rs | 11 +- app/src/ai/blocklist/permissions.rs | 2 + app/src/ai/blocklist/permissions_tests.rs | 1 + 18 files changed, 1308 insertions(+), 127 deletions(-) create mode 100644 app/src/ai/agent_sdk/hooks/adapters.rs create mode 100644 app/src/ai/agent_sdk/hooks/adapters_tests.rs create mode 100644 app/src/ai/agent_sdk/hooks/trust_tests.rs diff --git a/app/Cargo.toml b/app/Cargo.toml index 1f90fcf3142..9a47cc0e128 100644 --- a/app/Cargo.toml +++ b/app/Cargo.toml @@ -406,6 +406,7 @@ windows = { workspace = true, features = [ "Win32_System_Diagnostics_Debug", "Win32_System_Environment", "Win32_System_IO", + "Win32_System_JobObjects", "Win32_System_ProcessStatus", "Win32_System_SystemInformation", "Win32_System_Threading", diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index b86784b2033..1abda81c1b1 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -62,7 +62,10 @@ use crate::ai::agent_sdk::hooks::payload::{ }; use crate::ai::agent_sdk::hooks::redaction::HookRedactor; use crate::ai::agent_sdk::hooks::runtime::{OzHookEvent, OzHookRuntime, OzHookRuntimeService}; -use crate::ai::agent_sdk::hooks::trust::{ExactHookTrustStore, HookTrustKey}; +use crate::ai::agent_sdk::hooks::trust::{ + DenyProjectHookTrust, ExactHookTrustStore, HookTrustKey, HookTrustStore, + PersistentHookTrustStore, +}; use crate::ai::agent_sdk::hooks::{MAX_PROMPT_BYTES, OzHookSession, PAYLOAD_SCHEMA_VERSION}; use crate::ai::agent_sdk::setup_observability::{SetupClientEventReporter, SetupStep}; use crate::ai::ambient_agents::task::HarnessModelConfig; @@ -174,6 +177,16 @@ where } } +fn secret_values(secret: &ManagedSecretValue) -> Vec { + match secret { + ManagedSecretValue::RawValue { value } => vec![value.clone()], + secret => typed_secret_entries(secret) + .into_iter() + .map(|(_, value)| value.to_owned()) + .collect(), + } +} + const HARNESS_SAVE_INTERVAL: Duration = Duration::from_secs(30); /// Delay after the initial exit request before retrying with the harness's /// follow-up input (e.g. Claude's confirmation-dialog dismissal). Sent @@ -1009,12 +1022,13 @@ impl AgentDriver { async fn initialize_oz_hook_runtime( foreground: &ModelSpawner, ) -> Result>, AgentDriverError> { - let (context, cwd, task_id, hooks_enabled) = foreground + let (context, cwd, task_id, secrets, hooks_enabled) = foreground .spawn(|me, _| { ( me.oz_lifecycle_hooks_context.clone(), me.harness_working_dir.clone(), me.task_id, + Arc::clone(&me.secrets), FeatureFlag::OzLifecycleHooks.is_enabled(), ) }) @@ -1023,8 +1037,8 @@ impl AgentDriver { return Ok(None); } - let trust_store = ExactHookTrustStore::default(); - if let Some(context) = context { + let trust_store: Arc = if let Some(context) = context { + let trust_store = ExactHookTrustStore::default(); for trust in context.project_trust { let (Ok(git_root), Ok(config_path)) = ( std::fs::canonicalize(trust.git_root), @@ -1038,8 +1052,17 @@ impl AgentDriver { definition_hash: trust.sha256, }); } - } - let config = discover_hook_config(&cwd, &trust_store); + Arc::new(trust_store) + } else { + match PersistentHookTrustStore::load_default() { + Ok(trust_store) => Arc::new(trust_store), + Err(error) => { + log::warn!("Failed to load Oz hook trust store: {error}"); + Arc::new(DenyProjectHookTrust) + } + } + }; + let config = discover_hook_config(&cwd, trust_store.as_ref()); for diagnostic in config.diagnostics.iter() { log::warn!( "Oz hook configuration diagnostic: kind={:?} path={} hash_present={}", @@ -1052,6 +1075,9 @@ impl AgentDriver { .enabled_events() .map(|event| event.protocol_value().into()) .collect(); + let run_id = task_id + .map(|id| id.to_string()) + .unwrap_or_else(|| Uuid::new_v4().to_string()); let runtime: Arc = Arc::new(OzHookRuntimeService::new(config)); let session = OzHookSession { runtime: Arc::clone(&runtime), @@ -1059,22 +1085,28 @@ impl AgentDriver { enabled_events, supported_payload_schema_versions: vec![PAYLOAD_SCHEMA_VERSION.into()], }, + payload_context: HookPayloadContext { + session_id: run_id.clone(), + run_id: run_id.clone(), + conversation_id: String::new(), + cwd: cwd.to_string_lossy().into_owned(), + model: String::new(), + permission_mode: "supervised".into(), + }, + redactor: HookRedactor::new(secrets.values().flat_map(secret_values)), }; foreground .spawn(move |me, ctx| { me.terminal_driver.update(ctx, |driver, ctx| { driver.with_terminal_view(ctx, |terminal, ctx| { - terminal.ai_controller().update(ctx, |controller, _| { - controller.set_oz_hook_session(Some(session)); + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.set_oz_hook_session(Some(session), ctx); }); }); }); }) .await?; - let run_id = task_id - .map(|id| id.to_string()) - .unwrap_or_else(|| Uuid::new_v4().to_string()); runtime .observe(OzHookEvent { invocation_id: Uuid::new_v4().to_string(), diff --git a/app/src/ai/agent_sdk/hooks/adapters.rs b/app/src/ai/agent_sdk/hooks/adapters.rs new file mode 100644 index 00000000000..9717613c61b --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/adapters.rs @@ -0,0 +1,296 @@ +use super::MAX_TOOL_INPUT_BYTES; +use super::redaction::{HookRedactor, RedactedText, RedactedValue}; +use crate::ai::agent::{ + AIAgentActionResultType, AIAgentActionType, AIAgentActionTypeDiscriminants, +}; +const MAX_METADATA_TEXT_BYTES: usize = 8 * 1024; +const MAX_METADATA_ITEMS: usize = 128; + +pub(crate) fn local_action_payload( + action: &AIAgentActionType, + redactor: &HookRedactor, +) -> (&'static str, RedactedValue) { + let tool_name = match AIAgentActionTypeDiscriminants::from(action) { + AIAgentActionTypeDiscriminants::RequestCommandOutput => "run_shell_command", + AIAgentActionTypeDiscriminants::WriteToLongRunningShellCommand => { + "write_to_long_running_shell_command" + } + AIAgentActionTypeDiscriminants::ReadFiles => "read_files", + AIAgentActionTypeDiscriminants::UploadArtifact => "upload_artifact", + AIAgentActionTypeDiscriminants::SearchCodebase => "search_codebase", + AIAgentActionTypeDiscriminants::RequestFileEdits => "request_file_edits", + AIAgentActionTypeDiscriminants::Grep => "grep", + AIAgentActionTypeDiscriminants::FileGlob => "file_glob", + AIAgentActionTypeDiscriminants::FileGlobV2 => "file_glob", + AIAgentActionTypeDiscriminants::ReadMCPResource => "read_mcp_resource", + AIAgentActionTypeDiscriminants::CallMCPTool => "call_mcp_tool", + AIAgentActionTypeDiscriminants::SuggestNewConversation => "suggest_new_conversation", + AIAgentActionTypeDiscriminants::SuggestPrompt => "suggest_prompt", + AIAgentActionTypeDiscriminants::InitProject => "init_project", + AIAgentActionTypeDiscriminants::OpenCodeReview => "open_code_review", + AIAgentActionTypeDiscriminants::ReadDocuments => "read_documents", + AIAgentActionTypeDiscriminants::EditDocuments => "edit_documents", + AIAgentActionTypeDiscriminants::CreateDocuments => "create_documents", + AIAgentActionTypeDiscriminants::ReadShellCommandOutput => "read_shell_command_output", + AIAgentActionTypeDiscriminants::UseComputer => "use_computer", + AIAgentActionTypeDiscriminants::InsertCodeReviewComments => "insert_code_review_comments", + AIAgentActionTypeDiscriminants::RequestComputerUse => "request_computer_use", + AIAgentActionTypeDiscriminants::StartRecording => "start_recording", + AIAgentActionTypeDiscriminants::StopRecording => "stop_recording", + AIAgentActionTypeDiscriminants::ReadSkill => "read_skill", + AIAgentActionTypeDiscriminants::FetchConversation => "fetch_conversation", + AIAgentActionTypeDiscriminants::SendMessageToAgent => "send_message_to_agent", + AIAgentActionTypeDiscriminants::TransferShellCommandControlToUser => { + "transfer_shell_command_control_to_user" + } + AIAgentActionTypeDiscriminants::AskUserQuestion => "ask_user_question", + AIAgentActionTypeDiscriminants::RunAgents => "run_agents", + AIAgentActionTypeDiscriminants::WaitForEvents => "wait_for_events", + }; + let input = match action { + AIAgentActionType::RequestCommandOutput { + command, + is_read_only, + is_risky, + wait_until_completion, + uses_pager, + .. + } => RedactedValue::object([ + ("command", redacted_text(redactor, command)), + ( + "is_read_only", + is_read_only.map_or(RedactedValue::Null, RedactedValue::Bool), + ), + ( + "is_risky", + is_risky.map_or(RedactedValue::Null, RedactedValue::Bool), + ), + ( + "wait_until_completion", + RedactedValue::Bool(*wait_until_completion), + ), + ( + "uses_pager", + uses_pager.map_or(RedactedValue::Null, RedactedValue::Bool), + ), + ]), + AIAgentActionType::WriteToLongRunningShellCommand { input, mode, .. } => { + RedactedValue::object([ + ( + "input", + redacted_text(redactor, &String::from_utf8_lossy(input)), + ), + ( + "mode", + RedactedValue::from(format!("{mode:?}").to_lowercase()), + ), + ]) + } + AIAgentActionType::ReadFiles(request) => RedactedValue::object([ + ( + "paths", + string_array( + request + .locations + .iter() + .map(|location| location.name.as_str()), + redactor, + ), + ), + ( + "path_count", + RedactedValue::from(request.locations.len() as u64), + ), + ]), + AIAgentActionType::UploadArtifact(request) => RedactedValue::object([ + ("path", redacted_text(redactor, &request.file_path)), + ( + "description", + request + .description + .as_deref() + .map_or(RedactedValue::Null, |value| redacted_text(redactor, value)), + ), + ]), + AIAgentActionType::SearchCodebase(request) => RedactedValue::object([ + ("query", redacted_text(redactor, &request.query)), + ( + "paths", + request + .partial_paths + .as_ref() + .map_or(RedactedValue::Null, |paths| { + string_array(paths.iter().map(String::as_str), redactor) + }), + ), + ( + "codebase_path", + request + .codebase_path + .as_deref() + .map_or(RedactedValue::Null, |value| redacted_text(redactor, value)), + ), + ]), + AIAgentActionType::RequestFileEdits { file_edits, .. } => RedactedValue::object([ + ( + "paths", + string_array(file_edits.iter().filter_map(|edit| edit.file()), redactor), + ), + ("file_count", RedactedValue::from(file_edits.len() as u64)), + ]), + AIAgentActionType::Grep { queries, path } => RedactedValue::object([ + ( + "queries", + string_array(queries.iter().map(String::as_str), redactor), + ), + ("path", redacted_text(redactor, path)), + ("query_count", RedactedValue::from(queries.len() as u64)), + ]), + AIAgentActionType::FileGlob { patterns, path } => RedactedValue::object([ + ( + "patterns", + string_array(patterns.iter().map(String::as_str), redactor), + ), + ( + "path", + path.as_deref() + .map_or(RedactedValue::Null, |value| redacted_text(redactor, value)), + ), + ]), + AIAgentActionType::FileGlobV2 { + patterns, + search_dir, + } => RedactedValue::object([ + ( + "patterns", + string_array(patterns.iter().map(String::as_str), redactor), + ), + ( + "search_dir", + search_dir + .as_deref() + .map_or(RedactedValue::Null, |value| redacted_text(redactor, value)), + ), + ]), + AIAgentActionType::CallMCPTool { name, input, .. } => { + let argument_keys = input + .as_object() + .map(|object| object.keys().map(String::as_str)) + .map_or_else( + || RedactedValue::Array(Vec::new()), + |keys| string_array(keys, redactor), + ); + RedactedValue::object([ + ("name", redacted_text(redactor, name)), + ("argument_keys", argument_keys), + ]) + } + AIAgentActionType::ReadMCPResource { name, uri, .. } => RedactedValue::object([ + ("name", redacted_text(redactor, name)), + ( + "uri", + uri.as_deref() + .map_or(RedactedValue::Null, |value| redacted_text(redactor, value)), + ), + ]), + AIAgentActionType::InsertCodeReviewComments { comments, .. } => { + RedactedValue::object([("comment_count", RedactedValue::from(comments.len() as u64))]) + } + AIAgentActionType::AskUserQuestion { questions } => RedactedValue::object([( + "question_count", + RedactedValue::from(questions.len() as u64), + )]), + AIAgentActionType::RunAgents(request) => RedactedValue::object([ + ( + "agent_count", + RedactedValue::from(request.agent_run_configs.len() as u64), + ), + ("summary", redacted_text(redactor, &request.summary)), + ("model", redacted_text(redactor, &request.model_id)), + ("harness", redacted_text(redactor, &request.harness_type)), + ( + "execution_mode", + RedactedValue::from(if request.execution_mode.is_remote() { + "remote" + } else { + "local" + }), + ), + ]), + AIAgentActionType::SuggestNewConversation { .. } + | AIAgentActionType::SuggestPrompt(_) + | AIAgentActionType::InitProject + | AIAgentActionType::OpenCodeReview + | AIAgentActionType::ReadDocuments(_) + | AIAgentActionType::EditDocuments(_) + | AIAgentActionType::CreateDocuments(_) + | AIAgentActionType::ReadShellCommandOutput { .. } + | AIAgentActionType::UseComputer(_) + | AIAgentActionType::RequestComputerUse(_) + | AIAgentActionType::StartRecording { .. } + | AIAgentActionType::StopRecording { .. } + | AIAgentActionType::ReadSkill(_) + | AIAgentActionType::FetchConversation { .. } + | AIAgentActionType::SendMessageToAgent { .. } + | AIAgentActionType::TransferShellCommandControlToUser { .. } + | AIAgentActionType::WaitForEvents { .. } => { + RedactedValue::redacted("sensitive_tool_input", 0) + } + }; + let input_bytes = input.serialized_len(); + let input = if input_bytes > MAX_TOOL_INPUT_BYTES { + RedactedValue::redacted("tool_input_size_limit", input_bytes) + } else { + input + }; + (tool_name, input) +} + +fn redacted_text(redactor: &HookRedactor, value: &str) -> RedactedValue { + let RedactedText { value, truncation } = redactor.redact_text(value, MAX_METADATA_TEXT_BYTES); + match truncation { + Some(truncation) => RedactedValue::object([ + ("value", RedactedValue::from(value)), + ("truncated", RedactedValue::Bool(truncation.truncated)), + ( + "original_bytes", + RedactedValue::from(truncation.original_bytes as u64), + ), + ( + "included_bytes", + RedactedValue::from(truncation.included_bytes as u64), + ), + ]), + None => RedactedValue::from(value), + } +} + +fn string_array<'a>( + values: impl IntoIterator, + redactor: &HookRedactor, +) -> RedactedValue { + RedactedValue::Array( + values + .into_iter() + .take(MAX_METADATA_ITEMS) + .map(|value| redacted_text(redactor, value)) + .collect(), + ) +} + +pub(crate) fn local_action_result_payload(result: &AIAgentActionResultType) -> RedactedValue { + let status = if result.is_cancelled() { + "cancelled" + } else if result.is_failed() { + "failed" + } else if result.is_successful() { + "succeeded" + } else { + "completed" + }; + RedactedValue::object([("status", RedactedValue::from(status))]) +} + +#[cfg(test)] +#[path = "adapters_tests.rs"] +mod tests; diff --git a/app/src/ai/agent_sdk/hooks/adapters_tests.rs b/app/src/ai/agent_sdk/hooks/adapters_tests.rs new file mode 100644 index 00000000000..43436dec6f2 --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/adapters_tests.rs @@ -0,0 +1,62 @@ +use ai::agent::action::{ReadFilesRequest, RunAgentsExecutionMode, RunAgentsRequest}; + +use super::*; + +#[test] +fn oz_hooks_adapters_cover_side_effect_categories_without_sensitive_content() { + let redactor = HookRedactor::new(["secret".into()]); + let cases = [ + ( + AIAgentActionType::RequestCommandOutput { + command: "printf secret".into(), + is_read_only: Some(false), + is_risky: Some(true), + wait_until_completion: true, + uses_pager: Some(false), + rationale: Some("secret rationale".into()), + citations: vec![], + }, + "run_shell_command", + "secret", + ), + ( + AIAgentActionType::ReadFiles(ReadFilesRequest { locations: vec![] }), + "read_files", + "file contents", + ), + ( + AIAgentActionType::CallMCPTool { + server_id: None, + name: "safe-tool-name".into(), + input: serde_json::json!({"token": "secret"}), + }, + "call_mcp_tool", + "secret", + ), + ( + AIAgentActionType::RunAgents(RunAgentsRequest { + summary: "safe summary".into(), + base_prompt: "secret child prompt".into(), + skills: vec![], + model_id: "model".into(), + harness_type: "oz".into(), + execution_mode: RunAgentsExecutionMode::Local, + agent_run_configs: vec![], + plan_id: String::new(), + harness_auth_secret_name: Some("secret-name".into()), + }), + "run_agents", + "secret", + ), + ]; + + for (action, expected_name, prohibited) in cases { + let (name, payload) = local_action_payload(&action, &redactor); + let serialized = serde_json::to_string(&payload).unwrap(); + assert_eq!(name, expected_name); + assert!(!serialized.contains(prohibited)); + assert!(!super::super::redaction::contains_prohibited_payload_key( + &serde_json::to_value(payload).unwrap() + )); + } +} diff --git a/app/src/ai/agent_sdk/hooks/mod.rs b/app/src/ai/agent_sdk/hooks/mod.rs index fe6972c62a0..f7c22fb1cc5 100644 --- a/app/src/ai/agent_sdk/hooks/mod.rs +++ b/app/src/ai/agent_sdk/hooks/mod.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; +pub(crate) mod adapters; pub(crate) mod config; pub(crate) mod payload; pub(crate) mod permissions; @@ -26,6 +27,8 @@ pub(crate) const MAX_DENIAL_REASON_BYTES: usize = 4 * 1024; pub(crate) struct OzHookSession { pub(crate) runtime: Arc, pub(crate) protocol_context: warp_multi_agent_api::OzHookContext, + pub(crate) payload_context: payload::HookPayloadContext, + pub(crate) redactor: redaction::HookRedactor, } #[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] diff --git a/app/src/ai/agent_sdk/hooks/permissions.rs b/app/src/ai/agent_sdk/hooks/permissions.rs index 35ba0b22563..d16e469f116 100644 --- a/app/src/ai/agent_sdk/hooks/permissions.rs +++ b/app/src/ai/agent_sdk/hooks/permissions.rs @@ -1,5 +1,4 @@ #[derive(Clone, Copy, Debug, Eq, PartialEq)] -#[allow(dead_code)] pub(crate) enum NativePermission { Deny, Allow, @@ -7,14 +6,12 @@ pub(crate) enum NativePermission { } #[derive(Clone, Copy, Debug, Eq, PartialEq)] -#[allow(dead_code)] pub(crate) enum HookPermission { Continue, Deny, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] -#[allow(dead_code)] pub(crate) enum ComposedPermission { DeniedByWarp, DeniedByHook, @@ -22,7 +19,6 @@ pub(crate) enum ComposedPermission { Prompt, } -#[allow(dead_code)] pub(crate) fn compose_permission( native: NativePermission, hook: HookPermission, diff --git a/app/src/ai/agent_sdk/hooks/protocol.rs b/app/src/ai/agent_sdk/hooks/protocol.rs index f87cc09e094..a6a75c6750c 100644 --- a/app/src/ai/agent_sdk/hooks/protocol.rs +++ b/app/src/ai/agent_sdk/hooks/protocol.rs @@ -105,6 +105,13 @@ pub(crate) fn event_from_protocol(action: &RunOzHook) -> Result diagnostics - .iter() - .rev() - .find(|diagnostic| diagnostic.failure_category.is_some()) - .map(|diagnostic| Some(failed_outcome(diagnostic, ResolvedAction::Deny))) - .unwrap_or_else(|| { + } => { + if let Some(diagnostic) = diagnostics.last() + && diagnostic.result == HookInvocationResult::Denied + && diagnostic.failure_category.is_some() + { + Some(failed_outcome(diagnostic, ResolvedAction::Deny)) + } else { Some(oz_hook_result::Outcome::Deny(oz_hook_result::Deny { reason: super::redaction::truncate_utf8(&reason, MAX_DENIAL_REASON_BYTES), source: source.as_str().into(), })) - }), + } + } + OzPreToolUseDecision::Cancelled { diagnostics } => { + let _ = diagnostics; + Some(oz_hook_result::Outcome::Cancelled( + oz_hook_result::Cancelled {}, + )) + } }; OzHookResult { invocation_id: action.invocation_id.clone(), @@ -239,14 +254,16 @@ fn validate_envelope_fields( if fields.contains_key("hook_source") { return Err(ProtocolHookError::UnexpectedHookSource); } - if let Some(schema_version) = fields.remove("schema_version") - && schema_version != RedactedValue::String(PAYLOAD_SCHEMA_VERSION.into()) - { + let schema_version = fields + .remove("schema_version") + .ok_or(ProtocolHookError::MissingPayloadField("schema_version"))?; + if schema_version != RedactedValue::String(PAYLOAD_SCHEMA_VERSION.into()) { return Err(ProtocolHookError::UnsupportedSchema); } - if let Some(hook_event_name) = fields.remove("hook_event_name") - && hook_event_name != RedactedValue::String(event.as_str().into()) - { + let hook_event_name = fields + .remove("hook_event_name") + .ok_or(ProtocolHookError::MissingPayloadField("hook_event_name"))?; + if hook_event_name != RedactedValue::String(event.as_str().into()) { return Err(ProtocolHookError::MismatchedEvent); } Ok(()) diff --git a/app/src/ai/agent_sdk/hooks/protocol_tests.rs b/app/src/ai/agent_sdk/hooks/protocol_tests.rs index eac1a3de80b..973e7348df4 100644 --- a/app/src/ai/agent_sdk/hooks/protocol_tests.rs +++ b/app/src/ai/agent_sdk/hooks/protocol_tests.rs @@ -31,6 +31,7 @@ fn object(fields: impl IntoIterator) -> Value { fn common_fields() -> BTreeMap { [ + ("schema_version", string(PAYLOAD_SCHEMA_VERSION)), ("session_id", string("session")), ("run_id", string("run")), ("conversation_id", string("conversation")), @@ -48,6 +49,17 @@ fn action( event_fields: impl IntoIterator, ) -> RunOzHook { let mut fields = common_fields(); + let hook_event_name = match event { + ProtocolEvent::SessionStart => "SessionStart", + ProtocolEvent::SessionEnd => "SessionEnd", + ProtocolEvent::UserPromptSubmit => "UserPromptSubmit", + ProtocolEvent::Stop => "Stop", + ProtocolEvent::PreToolUse => "PreToolUse", + ProtocolEvent::PostToolUse => "PostToolUse", + ProtocolEvent::PreCompact => "PreCompact", + ProtocolEvent::Unspecified => "Unspecified", + }; + fields.insert("hook_event_name".into(), string(hook_event_name)); fields.extend( event_fields .into_iter() @@ -68,6 +80,31 @@ fn action( } } +#[test] +fn oz_hooks_protocol_requires_envelope_fields_and_rejects_non_tool_ids() { + for field in ["schema_version", "hook_event_name"] { + let mut missing = action(ProtocolEvent::Stop, [("turn_status", string("idle"))]); + missing + .redacted_payload + .as_mut() + .unwrap() + .fields + .remove(field); + assert!(matches!( + event_from_protocol(&missing), + Err(ProtocolHookError::MissingPayloadField(missing_field)) + if missing_field == field + )); + } + + let mut non_tool = action(ProtocolEvent::Stop, [("turn_status", string("idle"))]); + non_tool.tool_use_id = "unexpected".into(); + assert!(matches!( + event_from_protocol(&non_tool), + Err(ProtocolHookError::InvalidToolUseId) + )); +} + #[test] fn oz_hooks_protocol_parses_all_seven_protocol_events() { let cases = [ @@ -258,4 +295,22 @@ fn oz_hooks_protocol_maps_continue_deny_failed_and_cancelled_results() { .outcome, Some(Outcome::Cancelled(_)) )); + assert!(matches!( + result_for_pre_tool( + &action, + OzPreToolUseDecision::Deny { + reason: "explicit".into(), + source: HookConfigSource::Project, + diagnostics: vec![ + diagnostic( + HookInvocationResult::Continued, + Some(HookFailureCategory::Timeout), + ), + diagnostic(HookInvocationResult::Denied, None), + ], + }, + ) + .outcome, + Some(Outcome::Deny(ref deny)) if deny.reason == "explicit" + )); } diff --git a/app/src/ai/agent_sdk/hooks/redaction.rs b/app/src/ai/agent_sdk/hooks/redaction.rs index 86bf874b304..54e89c3ce23 100644 --- a/app/src/ai/agent_sdk/hooks/redaction.rs +++ b/app/src/ai/agent_sdk/hooks/redaction.rs @@ -16,7 +16,6 @@ pub(crate) enum RedactedValue { } impl RedactedValue { - #[allow(dead_code)] pub(crate) fn object( fields: impl IntoIterator, RedactedValue)>, ) -> Self { @@ -28,7 +27,6 @@ impl RedactedValue { ) } - #[allow(dead_code)] pub(crate) fn redacted(reason: &str, byte_count: usize) -> Self { Self::object([ ("redacted", Self::Bool(true)), diff --git a/app/src/ai/agent_sdk/hooks/runtime.rs b/app/src/ai/agent_sdk/hooks/runtime.rs index 630681daeb0..7af57d2b7a5 100644 --- a/app/src/ai/agent_sdk/hooks/runtime.rs +++ b/app/src/ai/agent_sdk/hooks/runtime.rs @@ -70,6 +70,9 @@ pub(crate) enum OzPreToolUseDecision { source: HookConfigSource, diagnostics: Vec, }, + Cancelled { + diagnostics: Vec, + }, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -164,7 +167,10 @@ impl OzHookRuntimeService { guard = self.queue.lock() => guard, () = token.cancelled() => { self.remove_pending(&event.invocation_id); - return EventOutcome::default(); + return EventOutcome { + cancelled: true, + ..Default::default() + }; } }; @@ -207,6 +213,7 @@ impl OzHookRuntimeService { output_truncated: false, failure_category: None, }; + let mut stop_handlers = false; match result { Ok(CommandOutcome::Continue { exit_code }) => { diagnostic.exit_code = exit_code; @@ -214,9 +221,8 @@ impl OzHookRuntimeService { Ok(CommandOutcome::Deny { reason, exit_code }) if pre_tool => { diagnostic.result = HookInvocationResult::Denied; diagnostic.exit_code = exit_code; - outcome.diagnostics.push(diagnostic); outcome.denial = Some((reason, handler.source)); - break; + stop_handlers = true; } Ok(CommandOutcome::Deny { exit_code, .. }) => { diagnostic.result = HookInvocationResult::Failed; @@ -250,12 +256,11 @@ impl OzHookRuntimeService { && handler.on_failure == FailureMode::Deny && failure.category != HookFailureCategory::Cancelled { - outcome.diagnostics.push(diagnostic); outcome.denial = Some(( "An Oz hook failed closed and denied this tool.".into(), handler.source, )); - break; + stop_handlers = true; } } } @@ -277,6 +282,9 @@ impl OzHookRuntimeService { diagnostic.failure_category ); outcome.diagnostics.push(diagnostic); + if stop_handlers { + break; + } } self.remove_pending(&event.invocation_id); outcome @@ -307,6 +315,9 @@ impl OzHookRuntime for OzHookRuntimeService { source, diagnostics: outcome.diagnostics, }, + None if outcome.cancelled => OzPreToolUseDecision::Cancelled { + diagnostics: outcome.diagnostics, + }, None => OzPreToolUseDecision::Continue { diagnostics: outcome.diagnostics, }, @@ -337,6 +348,7 @@ impl OzHookRuntime for OzHookRuntimeService { struct EventOutcome { diagnostics: Vec, denial: Option<(String, HookConfigSource)>, + cancelled: bool, } fn effective_timeout( @@ -389,6 +401,10 @@ async fn run_command( exit_code: None, })?; let process_id = child.id(); + let process_tree = HookProcessTree::attach(process_id).map_err(|_| CommandFailure { + category: HookFailureCategory::Spawn, + exit_code: None, + })?; let mut stdin = child.stdin.take().unwrap(); let stdin_payload = stdin_payload.to_vec(); let stdin_task = tokio::spawn(async move { @@ -405,9 +421,10 @@ async fn run_command( Cancelled, Overflow, } + let deadline = tokio::time::Instant::now() + timeout; let completion = tokio::select! { status = child.wait() => Completion::Exited(status), - () = tokio::time::sleep(timeout) => Completion::Timeout, + () = tokio::time::sleep_until(deadline) => Completion::Timeout, () = cancellation.cancelled() => Completion::Cancelled, Some(()) = overflow_rx.recv() => Completion::Overflow, }; @@ -417,38 +434,64 @@ async fn run_command( exit_code: None, })?, Completion::Timeout => { - kill_process_tree(process_id, &mut child).await; + process_tree.kill(process_id, &mut child).await; return Err(CommandFailure { category: HookFailureCategory::Timeout, exit_code: None, }); } Completion::Cancelled => { - kill_process_tree(process_id, &mut child).await; + process_tree.kill(process_id, &mut child).await; return Err(CommandFailure { category: HookFailureCategory::Cancelled, exit_code: None, }); } Completion::Overflow => { - kill_process_tree(process_id, &mut child).await; + process_tree.kill(process_id, &mut child).await; return Err(CommandFailure { category: HookFailureCategory::OutputOverflow, exit_code: None, }); } }; - match stdin_task.await { - Ok(Ok(())) => {} - Ok(Err(_)) | Err(_) => { + let exit_code = status.code(); + let outputs = async move { + let stdin_failed = !matches!(stdin_task.await, Ok(Ok(()))); + let stdout = join_output(stdout_task, exit_code).await?; + let stderr = join_output(stderr_task, exit_code).await?; + Ok((stdout, stderr, stdin_failed)) + }; + let (stdout, stderr, stdin_failed) = tokio::select! { + outputs = outputs => outputs?, + () = tokio::time::sleep_until(deadline) => { + process_tree.kill(process_id, &mut child).await; return Err(CommandFailure { - category: HookFailureCategory::Stdin, - exit_code: status.code(), + category: HookFailureCategory::Timeout, + exit_code, }); } + () = cancellation.cancelled() => { + process_tree.kill(process_id, &mut child).await; + return Err(CommandFailure { + category: HookFailureCategory::Cancelled, + exit_code, + }); + } + Some(()) = overflow_rx.recv() => { + process_tree.kill(process_id, &mut child).await; + return Err(CommandFailure { + category: HookFailureCategory::OutputOverflow, + exit_code, + }); + } + }; + if stdin_failed && status.success() { + return Err(CommandFailure { + category: HookFailureCategory::Stdin, + exit_code, + }); } - let stdout = join_output(stdout_task, status.code()).await?; - let stderr = join_output(stderr_task, status.code()).await?; parse_command_result(payload.event_name(), status.code(), stdout, stderr) } @@ -646,22 +689,101 @@ fn configure_process_group(command: &mut Command) { command.creation_flags(windows::Win32::System::Threading::CREATE_NEW_PROCESS_GROUP.0); } -#[cfg(unix)] -async fn kill_process_tree(process_id: Option, child: &mut tokio::process::Child) { - if let Some(process_id) = process_id { - let _ = nix::sys::signal::killpg( - nix::unistd::Pid::from_raw(process_id as i32), - nix::sys::signal::Signal::SIGKILL, - ); +struct HookProcessTree { + #[cfg(windows)] + job: WindowsJob, +} + +impl HookProcessTree { + #[cfg(unix)] + fn attach(_process_id: Option) -> Result { + Ok(Self {}) + } + + #[cfg(windows)] + fn attach(process_id: Option) -> Result { + Ok(Self { + job: WindowsJob::attach(process_id.ok_or(())?)?, + }) + } + + #[cfg(unix)] + async fn kill(&self, process_id: Option, child: &mut tokio::process::Child) { + if let Some(process_id) = process_id { + let _ = nix::sys::signal::killpg( + nix::unistd::Pid::from_raw(process_id as i32), + nix::sys::signal::Signal::SIGKILL, + ); + } + let _ = child.start_kill(); + let _ = child.wait().await; + } + + #[cfg(windows)] + async fn kill(&self, _process_id: Option, child: &mut tokio::process::Child) { + self.job.terminate(); + let _ = child.start_kill(); + let _ = child.wait().await; } - let _ = child.start_kill(); - let _ = child.wait().await; } #[cfg(windows)] -async fn kill_process_tree(_process_id: Option, child: &mut tokio::process::Child) { - let _ = child.start_kill(); - let _ = child.wait().await; +struct WindowsJob(windows::Win32::Foundation::HANDLE); + +#[cfg(windows)] +impl WindowsJob { + fn attach(process_id: u32) -> Result { + use windows::Win32::Foundation::CloseHandle; + use windows::Win32::System::JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, + SetInformationJobObject, + }; + use windows::Win32::System::Threading::{ + OpenProcess, PROCESS_SET_QUOTA, PROCESS_TERMINATE, + }; + + unsafe { + let job = CreateJobObjectW(None, windows::core::PCWSTR::null()).map_err(|_| ())?; + let mut limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); + limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + if SetInformationJobObject( + job, + JobObjectExtendedLimitInformation, + std::ptr::from_ref(&limits).cast(), + std::mem::size_of_val(&limits) as u32, + ) + .is_err() + { + let _ = CloseHandle(job); + return Err(()); + } + let process = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, false, process_id) + .map_err(|_| ())?; + let assigned = AssignProcessToJobObject(job, process); + let _ = CloseHandle(process); + if assigned.is_err() { + let _ = CloseHandle(job); + return Err(()); + } + Ok(Self(job)) + } + } + + fn terminate(&self) { + unsafe { + let _ = windows::Win32::System::JobObjects::TerminateJobObject(self.0, 1); + } + } +} + +#[cfg(windows)] +impl Drop for WindowsJob { + fn drop(&mut self) { + unsafe { + let _ = windows::Win32::Foundation::CloseHandle(self.0); + } + } } #[derive(Debug, thiserror::Error)] diff --git a/app/src/ai/agent_sdk/hooks/runtime_tests.rs b/app/src/ai/agent_sdk/hooks/runtime_tests.rs index ad5f9279f30..5dbb0801f25 100644 --- a/app/src/ai/agent_sdk/hooks/runtime_tests.rs +++ b/app/src/ai/agent_sdk/hooks/runtime_tests.rs @@ -156,10 +156,15 @@ async fn oz_hooks_runtime_invalid_allow_output_fails_open_or_closed() { #[tokio::test] async fn oz_hooks_runtime_timeout_kills_and_resolves_failure_mode() { + let temp = tempfile::tempdir().unwrap(); + let sentinel = temp.path().join("descendant-survived"); let runtime = runtime_with_hooks(json!({ "PreToolUse": [{"hooks": [{ "type": "command", - "command": "sleep 30", + "command": format!( + "(sleep 2; printf survived > '{}') & sleep 30", + sentinel.display() + ), "timeout": 1, "on_failure": "deny" }]}] @@ -170,6 +175,8 @@ async fn oz_hooks_runtime_timeout_kills_and_resolves_failure_mode() { assert!(started.elapsed() < Duration::from_secs(5)); assert!(matches!(decision, OzPreToolUseDecision::Deny { .. })); + tokio::time::sleep(Duration::from_secs(2)).await; + assert!(!sentinel.exists()); } #[tokio::test] @@ -227,3 +234,46 @@ async fn oz_hooks_runtime_cancellation_removes_pending_event() { HookInvocationResult::Cancelled ); } + +#[tokio::test] +async fn oz_hooks_runtime_cancellation_while_queued_is_not_continue() { + let runtime = Arc::new(runtime_with_hooks(json!({ + "SessionStart": [{"hooks": [{ + "type": "command", + "command": "sleep 30" + }]}], + "PreToolUse": [{"hooks": [{ + "type": "command", + "command": "exit 0" + }]}] + }))); + let blocker = { + let runtime = Arc::clone(&runtime); + tokio::spawn(async move { + runtime + .observe(event( + "blocker", + HookEventFields::SessionStart { + source: SessionStartSource::Startup, + }, + )) + .await + }) + }; + tokio::time::sleep(Duration::from_millis(100)).await; + let queued = { + let runtime = Arc::clone(&runtime); + tokio::spawn(async move { runtime.pre_tool_use(pre_tool_event("queued")).await }) + }; + tokio::task::yield_now().await; + + runtime.cancel(OzHookCancellationScope::Invocation("queued".into())); + let decision = tokio::time::timeout(Duration::from_secs(5), queued) + .await + .unwrap() + .unwrap(); + assert!(matches!(decision, OzPreToolUseDecision::Cancelled { .. })); + + runtime.cancel(OzHookCancellationScope::Invocation("blocker".into())); + blocker.await.unwrap(); +} diff --git a/app/src/ai/agent_sdk/hooks/trust.rs b/app/src/ai/agent_sdk/hooks/trust.rs index df484c2145c..30e0063775c 100644 --- a/app/src/ai/agent_sdk/hooks/trust.rs +++ b/app/src/ai/agent_sdk/hooks/trust.rs @@ -1,20 +1,185 @@ use std::collections::HashSet; -use std::path::PathBuf; +use std::fs; +use std::io::Write as _; +use std::path::{Path, PathBuf}; use std::sync::RwLock; -#[derive(Clone, Debug, Eq, Hash, PartialEq)] +use serde::{Deserialize, Serialize}; + +const TRUST_STORE_RELATIVE_PATH: &str = ".warp/oz-hook-trust.json"; +const TRUST_STORE_SCHEMA_VERSION: &str = "warp.oz_hook_trust.v1"; +const MAX_TRUST_RECORDS: usize = 1024; +const MAX_TRUST_STORE_BYTES: usize = 256 * 1024; + +#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] pub(crate) struct HookTrustKey { pub(crate) git_root: PathBuf, pub(crate) config_path: PathBuf, pub(crate) definition_hash: String, } +pub(crate) struct PersistentHookTrustStore { + path: PathBuf, + trusted: RwLock>, +} + +impl PersistentHookTrustStore { + pub(crate) fn load_default() -> Result { + let path = dirs::home_dir() + .ok_or(PersistentTrustError::MissingHome)? + .join(TRUST_STORE_RELATIVE_PATH); + Self::load(path) + } + + pub(crate) fn load(path: PathBuf) -> Result { + let trusted = match fs::read(&path) { + Ok(bytes) => { + if bytes.len() > MAX_TRUST_STORE_BYTES { + return Err(PersistentTrustError::Oversized); + } + let file: PersistentTrustFile = serde_json::from_slice(&bytes)?; + if file.schema_version != TRUST_STORE_SCHEMA_VERSION { + return Err(PersistentTrustError::UnsupportedSchema); + } + if file.records.len() > MAX_TRUST_RECORDS + || file.records.iter().any(|record| !valid_trust_key(record)) + { + return Err(PersistentTrustError::InvalidRecord); + } + file.records.into_iter().collect() + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => HashSet::new(), + Err(error) => return Err(error.into()), + }; + Ok(Self { + path, + trusted: RwLock::new(trusted), + }) + } + + #[allow(dead_code)] + pub(crate) fn trust(&self, key: HookTrustKey) -> Result<(), PersistentTrustError> { + let key = canonical_trust_key(key)?; + let mut trusted = self.trusted.write().unwrap(); + if trusted.len() >= MAX_TRUST_RECORDS && !trusted.contains(&key) { + return Err(PersistentTrustError::TooManyRecords); + } + let mut updated = trusted.clone(); + updated.insert(key); + self.persist(&updated)?; + *trusted = updated; + Ok(()) + } + + #[allow(dead_code)] + pub(crate) fn revoke(&self, key: &HookTrustKey) -> Result<(), PersistentTrustError> { + let key = canonical_trust_key(key.clone())?; + let mut trusted = self.trusted.write().unwrap(); + let mut updated = trusted.clone(); + updated.remove(&key); + self.persist(&updated)?; + *trusted = updated; + Ok(()) + } + + fn persist(&self, trusted: &HashSet) -> Result<(), PersistentTrustError> { + let parent = self + .path + .parent() + .ok_or(PersistentTrustError::InvalidPath)?; + fs::create_dir_all(parent)?; + let mut records = trusted.iter().cloned().collect::>(); + records.sort_by(|left, right| { + (&left.git_root, &left.config_path, &left.definition_hash).cmp(&( + &right.git_root, + &right.config_path, + &right.definition_hash, + )) + }); + let bytes = serde_json::to_vec_pretty(&PersistentTrustFile { + schema_version: TRUST_STORE_SCHEMA_VERSION.into(), + records, + })?; + if bytes.len() > MAX_TRUST_STORE_BYTES { + return Err(PersistentTrustError::Oversized); + } + let mut temporary = tempfile::NamedTempFile::new_in(parent)?; + temporary.write_all(&bytes)?; + temporary.flush()?; + temporary + .persist(&self.path) + .map_err(|error| PersistentTrustError::Io(error.error))?; + Ok(()) + } +} + +#[cfg(test)] +#[path = "trust_tests.rs"] +mod tests; + +impl HookTrustStore for PersistentHookTrustStore { + fn is_trusted(&self, key: &HookTrustKey) -> bool { + self.trusted.read().unwrap().contains(key) + } +} + +#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct PersistentTrustFile { + schema_version: String, + records: Vec, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum PersistentTrustError { + #[error("home directory is unavailable")] + MissingHome, + #[error("trust store path has no parent")] + InvalidPath, + #[error("trust store exceeds its size limit")] + Oversized, + #[error("trust store schema version is unsupported")] + UnsupportedSchema, + #[error("trust store contains an invalid record")] + InvalidRecord, + #[error("trust store contains too many records")] + TooManyRecords, + #[error("failed to access trust store: {0}")] + Io(#[from] std::io::Error), + #[error("failed to parse trust store: {0}")] + Json(#[from] serde_json::Error), +} + +pub(crate) fn is_hook_trust_store_path(path: &Path) -> bool { + path.ends_with(Path::new(TRUST_STORE_RELATIVE_PATH)) +} +fn canonical_trust_key(mut key: HookTrustKey) -> Result { + key.git_root = fs::canonicalize(key.git_root)?; + key.config_path = fs::canonicalize(key.config_path)?; + if valid_trust_key(&key) { + Ok(key) + } else { + Err(PersistentTrustError::InvalidRecord) + } +} + +fn valid_trust_key(key: &HookTrustKey) -> bool { + key.git_root.is_absolute() + && key.config_path.is_absolute() + && key.config_path.starts_with(&key.git_root) + && key.definition_hash.len() == 64 + && key + .definition_hash + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} + pub(crate) trait HookTrustStore: Send + Sync { fn is_trusted(&self, key: &HookTrustKey) -> bool; } #[derive(Default)] -#[allow(dead_code)] pub(crate) struct DenyProjectHookTrust; impl HookTrustStore for DenyProjectHookTrust { diff --git a/app/src/ai/agent_sdk/hooks/trust_tests.rs b/app/src/ai/agent_sdk/hooks/trust_tests.rs new file mode 100644 index 00000000000..9fc36423e3d --- /dev/null +++ b/app/src/ai/agent_sdk/hooks/trust_tests.rs @@ -0,0 +1,92 @@ +use std::fs; + +use sha2::{Digest as _, Sha256}; + +use super::*; + +fn trust_key(root: &Path, config_path: &Path) -> HookTrustKey { + HookTrustKey { + git_root: fs::canonicalize(root).unwrap(), + config_path: fs::canonicalize(config_path).unwrap(), + definition_hash: hex::encode(Sha256::digest(fs::read(config_path).unwrap())), + } +} + +#[test] +fn oz_hooks_persistent_trust_round_trips_and_revokes_exact_definition() { + let temp = tempfile::tempdir().unwrap(); + let project = temp.path().join("project"); + let config_dir = project.join(".warp"); + let config_path = config_dir.join("hooks.json"); + let store_path = temp.path().join("trust.json"); + fs::create_dir_all(&config_dir).unwrap(); + fs::write(&config_path, b"trusted definition").unwrap(); + let key = trust_key(&project, &config_path); + let store = PersistentHookTrustStore::load(store_path.clone()).unwrap(); + + store.trust(key.clone()).unwrap(); + assert!(store.is_trusted(&key)); + assert!( + PersistentHookTrustStore::load(store_path.clone()) + .unwrap() + .is_trusted(&key) + ); + + fs::write(&config_path, b"trusted definition\n").unwrap(); + let changed_key = trust_key(&project, &config_path); + assert!( + !PersistentHookTrustStore::load(store_path.clone()) + .unwrap() + .is_trusted(&changed_key) + ); + + store.revoke(&key).unwrap(); + assert!( + !PersistentHookTrustStore::load(store_path) + .unwrap() + .is_trusted(&key) + ); +} + +#[test] +fn oz_hooks_persistent_trust_rejects_unknown_schema_fields_and_invalid_hashes() { + let temp = tempfile::tempdir().unwrap(); + let store_path = temp.path().join("trust.json"); + let cases = [ + serde_json::json!({ + "schema_version": TRUST_STORE_SCHEMA_VERSION, + "records": [], + "unknown": true + }), + serde_json::json!({ + "schema_version": "future", + "records": [] + }), + serde_json::json!({ + "schema_version": TRUST_STORE_SCHEMA_VERSION, + "records": [{ + "git_root": "/project", + "config_path": "/project/.warp/hooks.json", + "definition_hash": "ABC" + }] + }), + ]; + + for contents in cases { + fs::write(&store_path, serde_json::to_vec(&contents).unwrap()).unwrap(); + assert!(PersistentHookTrustStore::load(store_path.clone()).is_err()); + } +} + +#[test] +fn oz_hooks_identifies_only_the_host_trust_store_path() { + assert!(is_hook_trust_store_path(Path::new( + "/home/user/.warp/oz-hook-trust.json" + ))); + assert!(!is_hook_trust_store_path(Path::new( + "/project/.warp/hooks.json" + ))); + assert!(!is_hook_trust_store_path(Path::new( + "/project/oz-hook-trust.json" + ))); +} diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index dc224d324c1..88ee53c453d 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -63,6 +63,8 @@ use crate::ai::agent::{ CancellationOutcome, CancellationReason, CreateDocumentsResult, EditDocumentsResult, RequestCommandOutputResult, }; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::OzHookSession; use crate::ai::blocklist::action_model::execute::suggest_new_conversation::SuggestNewConversationExecutor; use crate::ai::blocklist::telemetry::send_run_agents_completed_telemetry; use crate::ai::document::ai_document_model::AIDocumentModel; @@ -289,6 +291,27 @@ impl BlocklistAIActionModel { } => { me.handle_action_result(*conversation_id, result.clone(), *cancellation_reason, ctx) } + BlocklistAIActionExecutorEvent::OzPreflightNotExecuted { + action, + conversation_id, + reason, + } => { + let should_remove_entry = + me.running_actions + .get_mut(conversation_id) + .is_some_and(|running| { + running.remove_action(&action.id); + running.is_empty() + }); + if should_remove_entry { + me.running_actions.remove(conversation_id); + } + me.pending_actions + .entry(*conversation_id) + .or_default() + .push_front((**action).clone()); + me.handle_not_executed_action(action, *reason, *conversation_id, ctx); + } BlocklistAIActionExecutorEvent::InitProject(id) => { ctx.emit(BlocklistAIActionEvent::InitProject(id.clone())) } @@ -430,6 +453,17 @@ impl BlocklistAIActionModel { }); } + #[cfg(not(target_family = "wasm"))] + pub(crate) fn set_oz_hook_session( + &mut self, + session: Option, + ctx: &mut ModelContext, + ) { + self.executor.update(ctx, |executor, _| { + executor.set_oz_hook_session(session); + }); + } + fn blocked_action_for_conversation( &self, conversation_id: &AIConversationId, diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 64c2c47e598..31c181c7b79 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -92,6 +92,20 @@ use crate::ai::agent::{ AIAgentActionType, AIAgentActionTypeDiscriminants, CancellationReason, FileContext, FileLocations, ReadFilesFailedFile, ServerOutputId, }; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::OzHookSession; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::adapters::{local_action_payload, local_action_result_payload}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::payload::{HookEventFields, HookPayloadTemplate}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::permissions::{ + ComposedPermission, HookPermission, NativePermission, compose_permission, +}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::runtime::{ + OzHookCancellationScope, OzHookEvent, OzPreToolUseDecision, OzPreToolUseEvent, +}; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::action_model::recording_controller::RecordingController; use crate::ai::blocklist::telemetry::send_run_agents_completed_telemetry; @@ -120,6 +134,12 @@ pub(super) enum ParallelExecutionPolicy { /// same execution phase when the underlying runtime supports it. ReadOnlyLocalContext, } +#[derive(Clone, Copy, Eq, PartialEq)] +enum LocalActionHookStage { + Preflight, + Executing, + Postflight, +} /// Whether an action is running serially or in parallel with other actions. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -239,6 +259,7 @@ struct AsyncExecutingAction { /// The conversation this action belongs to so cancellation and follow-up scheduling remain /// scoped even when several conversations have async actions in flight. conversation_id: AIConversationId, + hook_stage: LocalActionHookStage, } impl AsyncExecutingAction { @@ -286,6 +307,10 @@ pub struct BlocklistAIActionExecutor { /// Reference to the terminal model for checking session sharing state. terminal_model: Arc>, team_context_resolver: TeamContextResolver, + #[cfg(not(target_family = "wasm"))] + oz_hook_session: Option, + #[cfg(not(target_family = "wasm"))] + oz_hook_approved_actions: std::collections::HashSet, } impl BlocklistAIActionExecutor { @@ -381,9 +406,18 @@ impl BlocklistAIActionExecutor { send_message_executor, ask_user_question_executor, wait_for_events_executor, + #[cfg(not(target_family = "wasm"))] + oz_hook_session: None, + #[cfg(not(target_family = "wasm"))] + oz_hook_approved_actions: Default::default(), } } + #[cfg(not(target_family = "wasm"))] + pub(crate) fn set_oz_hook_session(&mut self, session: Option) { + self.oz_hook_session = session; + } + pub fn async_executing_action(&self, action_id: &AIAgentActionId) -> Option<&AIAgentAction> { self.async_executing_actions .get(action_id) @@ -596,6 +630,27 @@ impl BlocklistAIActionExecutor { conversation_id: AIConversationId, is_user_initiated: bool, ctx: &mut ModelContext, + ) -> TryExecuteResult { + #[cfg(not(target_family = "wasm"))] + let oz_preflight_complete = self.oz_hook_approved_actions.remove(&action.id); + #[cfg(target_family = "wasm")] + let oz_preflight_complete = false; + self.try_to_execute_action_inner( + action, + conversation_id, + is_user_initiated, + oz_preflight_complete, + ctx, + ) + } + + fn try_to_execute_action_inner( + &mut self, + action: AIAgentAction, + conversation_id: AIConversationId, + is_user_initiated: bool, + oz_preflight_complete: bool, + ctx: &mut ModelContext, ) -> TryExecuteResult { // We should never actually execute actions in view-only mode. if self.is_shared_session_viewer() { @@ -618,12 +673,14 @@ impl BlocklistAIActionExecutor { let needs_confirmation = !(is_user_initiated || can_auto_execute || (is_agent_autonomous && action.action.is_request_command_output())); - if needs_confirmation { - return TryExecuteResult::NotExecuted { - action: Box::new(action), - reason: NotExecutedReason::NeedsConfirmation, - }; - } else if !is_user_initiated && !can_auto_execute && is_agent_autonomous { + let native_permission = if !is_user_initiated && !can_auto_execute && is_agent_autonomous { + NativePermission::Deny + } else if needs_confirmation { + NativePermission::Prompt + } else { + NativePermission::Allow + }; + if native_permission == NativePermission::Deny { // It must be the case that the autonomous agent is requesting a denylisted command. if let AIAgentActionType::RequestCommandOutput { command, .. } = &action.action { let action_id = action.id.clone(); @@ -650,6 +707,121 @@ impl BlocklistAIActionExecutor { } } + #[cfg(not(target_family = "wasm"))] + if !oz_preflight_complete && let Some(session) = self.oz_hook_session.clone() { + let (tool_name, tool_input) = local_action_payload(&action.action, &session.redactor); + let tool_use_id = action.id.to_string(); + let mut payload_context = session.payload_context.clone(); + payload_context.conversation_id = conversation_id.to_string(); + let event = OzPreToolUseEvent::new(OzHookEvent { + invocation_id: uuid::Uuid::new_v4().to_string(), + tool_use_id: Some(tool_use_id.clone()), + payload: HookPayloadTemplate { + context: payload_context, + event: HookEventFields::PreToolUse { + tool_name: tool_name.into(), + tool_use_id, + tool_input, + }, + }, + }) + .expect("local pre-tool event uses the pre-tool payload"); + let action_id = action.id.clone(); + self.async_executing_actions.insert( + action_id.clone(), + AsyncExecutingAction { + action: action.clone(), + conversation_id, + hook_stage: LocalActionHookStage::Preflight, + }, + ); + ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { + action_id: action_id.clone(), + }); + ctx.spawn( + async move { session.runtime.pre_tool_use(event).await }, + move |me, decision, ctx| { + if !me.async_executing_actions.contains_key(&action_id) { + return; + } + if matches!(decision, OzPreToolUseDecision::Cancelled { .. }) { + let running = me + .async_executing_actions + .remove(&action_id) + .expect("preflight action was checked above"); + let result = running.action.action.cancelled_result(); + Self::emit_finished_action(running, result, None, ctx); + return; + } + let hook_permission = match decision { + OzPreToolUseDecision::Continue { .. } => HookPermission::Continue, + OzPreToolUseDecision::Deny { .. } => HookPermission::Deny, + OzPreToolUseDecision::Cancelled { .. } => unreachable!(), + }; + match compose_permission(native_permission, hook_permission) { + ComposedPermission::Allow => { + if let Some(running) = me.async_executing_actions.get_mut(&action_id) { + running.hook_stage = LocalActionHookStage::Executing; + } + let result = me.try_to_execute_action_inner( + action, + conversation_id, + is_user_initiated, + true, + ctx, + ); + if let TryExecuteResult::NotExecuted { action, reason } = result { + me.async_executing_actions.remove(&action_id); + me.oz_hook_approved_actions.insert(action_id.clone()); + ctx.emit(BlocklistAIActionExecutorEvent::OzPreflightNotExecuted { + action, + conversation_id, + reason, + }); + } + } + ComposedPermission::DeniedByHook => { + let Some(running) = me.async_executing_actions.remove(&action_id) + else { + return; + }; + ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction { + result: Arc::new(AIAgentActionResult { + id: running.action.id, + task_id: running.action.task_id, + result: running.action.action.cancelled_result(), + }), + conversation_id: running.conversation_id, + cancellation_reason: None, + }); + } + ComposedPermission::Prompt => { + let Some(running) = me.async_executing_actions.remove(&action_id) + else { + return; + }; + me.oz_hook_approved_actions.insert(action_id.clone()); + ctx.emit(BlocklistAIActionExecutorEvent::OzPreflightNotExecuted { + action: Box::new(running.action), + conversation_id: running.conversation_id, + reason: NotExecutedReason::NeedsConfirmation, + }); + } + ComposedPermission::DeniedByWarp => { + unreachable!("native denials return before hook preflight") + } + } + }, + ); + return TryExecuteResult::ExecutedAsync; + } + if needs_confirmation { + return TryExecuteResult::NotExecuted { + action: Box::new(action), + reason: NotExecutedReason::NeedsConfirmation, + }; + } + let action_clone = action.clone(); let execution = match &action.action { AIAgentActionType::RequestCommandOutput { .. } @@ -807,25 +979,18 @@ impl BlocklistAIActionExecutor { AsyncExecutingAction { action: action_clone, conversation_id, + hook_stage: LocalActionHookStage::Executing, }, ); ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { action_id: action_id.clone(), }); ctx.spawn(execute_future, move |me, result, ctx| { - let Some(running) = me.async_executing_actions.remove(&action_id) else { + let Some(running) = me.async_executing_actions.get(&action_id).cloned() else { return; }; let result = on_complete(result, ctx); - ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction { - result: Arc::new(AIAgentActionResult { - id: action_id, - task_id: running.action.task_id, - result, - }), - conversation_id: running.conversation_id, - cancellation_reason: None, - }); + me.finish_action(running, result, None, ctx); }); TryExecuteResult::ExecutedAsync } @@ -833,18 +998,93 @@ impl BlocklistAIActionExecutor { ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction { action_id: action_id.clone(), }); - ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction { - result: Arc::new(AIAgentActionResult { - id: action_id, - task_id: action.task_id, - result: action_result, - }), - conversation_id, - cancellation_reason: None, - }); - TryExecuteResult::ExecutedSync + let running = self + .async_executing_actions + .get(&action_id) + .cloned() + .unwrap_or(AsyncExecutingAction { + action: action_clone, + conversation_id, + hook_stage: LocalActionHookStage::Executing, + }); + self.finish_action(running, action_result, None, ctx); + if oz_preflight_complete { + TryExecuteResult::ExecutedAsync + } else { + TryExecuteResult::ExecutedSync + } + } + } + } + + fn finish_action( + &mut self, + running: AsyncExecutingAction, + result: AIAgentActionResultType, + cancellation_reason: Option, + ctx: &mut ModelContext, + ) { + #[cfg(not(target_family = "wasm"))] + if let Some(session) = self.oz_hook_session.clone() { + let action_id = running.action.id.clone(); + let tool_use_id = action_id.to_string(); + let (tool_name, tool_input) = + local_action_payload(&running.action.action, &session.redactor); + let mut payload_context = session.payload_context.clone(); + payload_context.conversation_id = running.conversation_id.to_string(); + if let Some(state) = self.async_executing_actions.get_mut(&action_id) { + state.hook_stage = LocalActionHookStage::Postflight; + } else { + self.async_executing_actions + .insert(action_id.clone(), running.clone()); + self.async_executing_actions + .get_mut(&action_id) + .expect("action was inserted") + .hook_stage = LocalActionHookStage::Postflight; } + let event = OzHookEvent { + invocation_id: uuid::Uuid::new_v4().to_string(), + tool_use_id: Some(tool_use_id.clone()), + payload: HookPayloadTemplate { + context: payload_context, + event: HookEventFields::PostToolUse { + tool_name: tool_name.into(), + tool_use_id, + tool_input, + tool_response: local_action_result_payload(&result), + }, + }, + }; + ctx.spawn( + async move { session.runtime.observe(event).await }, + move |me, _, ctx| { + if me.async_executing_actions.remove(&action_id).is_none() { + return; + } + Self::emit_finished_action(running, result, cancellation_reason, ctx); + }, + ); + return; } + self.async_executing_actions.remove(&running.action.id); + Self::emit_finished_action(running, result, cancellation_reason, ctx); + } + + fn emit_finished_action( + running: AsyncExecutingAction, + result: AIAgentActionResultType, + cancellation_reason: Option, + ctx: &mut ModelContext, + ) { + ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction { + result: Arc::new(AIAgentActionResult { + id: running.action.id, + task_id: running.action.task_id, + result, + }), + conversation_id: running.conversation_id, + cancellation_reason, + }); } pub fn can_autoexecute_action( @@ -873,39 +1113,45 @@ impl BlocklistAIActionExecutor { return; } if let Some(running) = self.async_executing_actions.remove(action_id) { + #[cfg(not(target_family = "wasm"))] + if let Some(session) = &self.oz_hook_session { + session + .runtime + .cancel(OzHookCancellationScope::Tool(action_id.to_string())); + } let action_kind = AIAgentActionTypeDiscriminants::from(&running.action.action); log::info!( "Canceling running async action of type {action_kind:?} action_id={action_id:?}, reason={reason:?}, backtrace=\n{}", std::backtrace::Backtrace::force_capture() ); - if running.is_shell_command_action() { - self.shell_command_executor.update(ctx, |executor, ctx| { - executor.cancel_execution(&running.action.id, ctx); - }); - } else if matches!(running.action.action, AIAgentActionType::SearchCodebase(..)) { - self.search_codebase_executor.update(ctx, |executor, ctx| { - executor.cancel_execution(&running.action.id, ctx); - }); - } else if matches!(running.action.action, AIAgentActionType::RunAgents(..)) { - self.run_agents_executor.update(ctx, |executor, ctx| { - executor.cancel_execution(&running.action.id, ctx); - }); - } else if matches!( - running.action.action, - AIAgentActionType::StartRecording { .. } - ) { - RecordingController::handle(ctx).update(ctx, |controller, _| { - controller.abort_start(running.conversation_id); - }); - } else if let AIAgentActionType::WaitForEvents { tool_call_id, .. } = - &running.action.action - { - // Drop the executor's pending entry; the shared cancel - // path emits FinishedAction(Cancelled). - let tool_call_id = tool_call_id.clone(); - self.wait_for_events_executor.update(ctx, |executor, _| { - executor.cancel_execution(&tool_call_id); - }); + if running.hook_stage != LocalActionHookStage::Preflight { + if running.is_shell_command_action() { + self.shell_command_executor.update(ctx, |executor, ctx| { + executor.cancel_execution(&running.action.id, ctx); + }); + } else if matches!(running.action.action, AIAgentActionType::SearchCodebase(..)) { + self.search_codebase_executor.update(ctx, |executor, ctx| { + executor.cancel_execution(&running.action.id, ctx); + }); + } else if matches!(running.action.action, AIAgentActionType::RunAgents(..)) { + self.run_agents_executor.update(ctx, |executor, ctx| { + executor.cancel_execution(&running.action.id, ctx); + }); + } else if matches!( + running.action.action, + AIAgentActionType::StartRecording { .. } + ) { + RecordingController::handle(ctx).update(ctx, |controller, _| { + controller.abort_start(running.conversation_id); + }); + } else if let AIAgentActionType::WaitForEvents { tool_call_id, .. } = + &running.action.action + { + let tool_call_id = tool_call_id.clone(); + self.wait_for_events_executor.update(ctx, |executor, _| { + executor.cancel_execution(&tool_call_id); + }); + } } let result = running.action.action.cancelled_result(); send_run_agents_completed_telemetry( @@ -914,15 +1160,11 @@ impl BlocklistAIActionExecutor { &result, ctx, ); - ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction { - result: Arc::new(AIAgentActionResult { - id: running.action.id.clone(), - task_id: running.action.task_id, - result, - }), - conversation_id: running.conversation_id, - cancellation_reason: reason, - }); + if running.hook_stage == LocalActionHookStage::Executing { + self.finish_action(running, result, reason, ctx); + } else { + Self::emit_finished_action(running, result, reason, ctx); + } } } @@ -1088,6 +1330,12 @@ pub enum BlocklistAIActionExecutorEvent { cancellation_reason: Option, }, + OzPreflightNotExecuted { + action: Box, + conversation_id: AIConversationId, + reason: NotExecutedReason, + }, + InitProject(AIAgentActionId), OpenCodeReview(AIAgentActionId), InsertCodeReviewComments { diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index d9594556844..93a005f662c 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -2420,8 +2420,15 @@ impl BlocklistAIController { } #[cfg(not(target_family = "wasm"))] - pub(crate) fn set_oz_hook_session(&mut self, session: Option) { - self.oz_hook_session = session; + pub(crate) fn set_oz_hook_session( + &mut self, + session: Option, + ctx: &mut ModelContext, + ) { + self.oz_hook_session = session.clone(); + self.action_model.update(ctx, |model, ctx| { + model.set_oz_hook_session(session, ctx); + }); } #[cfg(not(target_family = "wasm"))] diff --git a/app/src/ai/blocklist/permissions.rs b/app/src/ai/blocklist/permissions.rs index f39f84aa63e..7684fe33938 100644 --- a/app/src/ai/blocklist/permissions.rs +++ b/app/src/ai/blocklist/permissions.rs @@ -14,6 +14,7 @@ use warpui::{AppContext, Entity, EntityId, ModelContext, SingletonEntity}; use super::BlocklistAIHistoryModel; use crate::ai::agent::conversation::AIConversationId; +use crate::ai::agent_sdk::hooks::trust::is_hook_trust_store_path; use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; use crate::ai::execution_profiles::{ AIExecutionProfile, ActionPermission, AskUserQuestionPermission, ExecutionProfileId, @@ -1236,6 +1237,7 @@ fn check_protected_write_paths(paths: &[PathBuf]) -> Option if paths.iter().any(|p| { mcp_provider_from_file_path(p).is_some() || p.ends_with(std::path::Path::new(".warp/hooks.json")) + || is_hook_trust_store_path(p) }) { Some(FileWritePermission::Denied( FileWritePermissionDeniedReason::ProtectedPath, diff --git a/app/src/ai/blocklist/permissions_tests.rs b/app/src/ai/blocklist/permissions_tests.rs index 404af0c5856..78264c46074 100644 --- a/app/src/ai/blocklist/permissions_tests.rs +++ b/app/src/ai/blocklist/permissions_tests.rs @@ -497,6 +497,7 @@ fn test_can_write_files_mcp_config_always_denied() { PathBuf::from("/project/.warp/.mcp.json"), PathBuf::from("/project/.codex/config.toml"), PathBuf::from("/project/.warp/hooks.json"), + PathBuf::from("/home/user/.warp/oz-hook-trust.json"), ]; for path in mcp_config_paths { From ae3e11fe810bfac7ee77d0fd8c9ce8ed3b50fa94 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 14:49:15 +0000 Subject: [PATCH 06/11] Wire Oz hooks to conversation lifecycle --- app/src/ai/agent_sdk/driver.rs | 211 +++++------ app/src/ai/agent_sdk/driver/harness/mod.rs | 6 + app/src/ai/agent_sdk/driver_tests.rs | 17 + app/src/ai/agent_sdk/hooks/mod.rs | 1 + app/src/ai/blocklist/action_model.rs | 3 +- app/src/ai/blocklist/action_model/execute.rs | 24 +- app/src/ai/blocklist/controller.rs | 338 ++++++++++++++++-- .../blocklist/controller/response_stream.rs | 96 +++++ .../controller/response_stream_tests.rs | 37 ++ 9 files changed, 602 insertions(+), 131 deletions(-) diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index 1abda81c1b1..e910870d0d9 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -66,7 +66,7 @@ use crate::ai::agent_sdk::hooks::trust::{ DenyProjectHookTrust, ExactHookTrustStore, HookTrustKey, HookTrustStore, PersistentHookTrustStore, }; -use crate::ai::agent_sdk::hooks::{MAX_PROMPT_BYTES, OzHookSession, PAYLOAD_SCHEMA_VERSION}; +use crate::ai::agent_sdk::hooks::{OzHookSession, PAYLOAD_SCHEMA_VERSION}; use crate::ai::agent_sdk::setup_observability::{SetupClientEventReporter, SetupStep}; use crate::ai::ambient_agents::task::HarnessModelConfig; use crate::ai::ambient_agents::{ @@ -1019,23 +1019,57 @@ impl From for AgentDriverError { } impl AgentDriver { - async fn initialize_oz_hook_runtime( + async fn prepare_oz_hook_conversation( foreground: &ModelSpawner, - ) -> Result>, AgentDriverError> { - let (context, cwd, task_id, secrets, hooks_enabled) = foreground - .spawn(|me, _| { - ( - me.oz_lifecycle_hooks_context.clone(), - me.harness_working_dir.clone(), - me.task_id, - Arc::clone(&me.secrets), - FeatureFlag::OzLifecycleHooks.is_enabled(), - ) + ) -> Result<(), AgentDriverError> { + foreground + .spawn(|me, ctx| { + if !FeatureFlag::OzLifecycleHooks.is_enabled() || me.run_conversation_id.is_some() { + return; + } + let mut conversation_id = None; + me.terminal_driver.update(ctx, |driver, ctx| { + driver.with_terminal_view(ctx, |terminal, ctx| { + terminal.ai_controller().update(ctx, |controller, ctx| { + conversation_id = Some(controller.start_oz_hook_conversation(ctx)); + }); + }); + }); + if let Some(conversation_id) = conversation_id { + me.run_conversation_id = Some(conversation_id); + stamp_parent_agent_id_if_some( + conversation_id, + me.parent_run_id.as_deref(), + ctx, + ); + register_agent_event_consumer(conversation_id, ctx.model_id(), ctx); + } }) .await?; + Ok(()) + } + async fn initialize_oz_hook_runtime( + model: String, + foreground: &ModelSpawner, + ) -> Result>, AgentDriverError> { + let (context, cwd, task_id, conversation_id, is_resume, secrets, hooks_enabled) = + foreground + .spawn(|me, _| { + ( + me.oz_lifecycle_hooks_context.clone(), + me.harness_working_dir.clone(), + me.task_id, + me.run_conversation_id, + me.restored_conversation_id.is_some(), + Arc::clone(&me.secrets), + FeatureFlag::OzLifecycleHooks.is_enabled(), + ) + }) + .await?; if !hooks_enabled { return Ok(None); } + let conversation_id = conversation_id.ok_or(AgentDriverError::InvalidRuntimeState)?; let trust_store: Arc = if let Some(context) = context { let trust_store = ExactHookTrustStore::default(); @@ -1078,6 +1112,15 @@ impl AgentDriver { let run_id = task_id .map(|id| id.to_string()) .unwrap_or_else(|| Uuid::new_v4().to_string()); + let session_id = Uuid::new_v4().to_string(); + let payload_context = HookPayloadContext { + session_id, + run_id, + conversation_id: conversation_id.to_string(), + cwd: cwd.to_string_lossy().into_owned(), + model, + permission_mode: "supervised".into(), + }; let runtime: Arc = Arc::new(OzHookRuntimeService::new(config)); let session = OzHookSession { runtime: Arc::clone(&runtime), @@ -1085,22 +1128,16 @@ impl AgentDriver { enabled_events, supported_payload_schema_versions: vec![PAYLOAD_SCHEMA_VERSION.into()], }, - payload_context: HookPayloadContext { - session_id: run_id.clone(), - run_id: run_id.clone(), - conversation_id: String::new(), - cwd: cwd.to_string_lossy().into_owned(), - model: String::new(), - permission_mode: "supervised".into(), - }, + payload_context: payload_context.clone(), redactor: HookRedactor::new(secrets.values().flat_map(secret_values)), + is_driver_owned: true, }; foreground .spawn(move |me, ctx| { me.terminal_driver.update(ctx, |driver, ctx| { driver.with_terminal_view(ctx, |terminal, ctx| { terminal.ai_controller().update(ctx, |controller, ctx| { - controller.set_oz_hook_session(Some(session), ctx); + controller.set_oz_hook_session(conversation_id, Some(session), ctx); }); }); }); @@ -1112,16 +1149,13 @@ impl AgentDriver { invocation_id: Uuid::new_v4().to_string(), tool_use_id: None, payload: HookPayloadTemplate { - context: HookPayloadContext { - session_id: run_id.clone(), - run_id, - conversation_id: String::new(), - cwd: cwd.to_string_lossy().into_owned(), - model: String::new(), - permission_mode: "supervised".into(), - }, + context: payload_context, event: HookEventFields::SessionStart { - source: SessionStartSource::Startup, + source: if is_resume { + SessionStartSource::Resume + } else { + SessionStartSource::Startup + }, }, }, }) @@ -1129,6 +1163,28 @@ impl AgentDriver { Ok(Some(runtime)) } + async fn current_oz_hook_session( + foreground: &ModelSpawner, + ) -> Result, warpui::ModelDropped> { + foreground + .spawn(|me, ctx| { + let mut session = None; + let conversation_id = me.run_conversation_id; + me.terminal_driver.update(ctx, |driver, ctx| { + driver.with_terminal_view(ctx, |terminal, ctx| { + session = conversation_id.and_then(|conversation_id| { + terminal + .ai_controller() + .as_ref(ctx) + .oz_hook_session(conversation_id) + }); + }); + }); + session + }) + .await + } + async fn finish_oz_hook_runtime( runtime: Option>, status: &SDKConversationOutputStatus, @@ -1137,10 +1193,7 @@ impl AgentDriver { let Some(runtime) = runtime else { return; }; - let Ok((cwd, task_id)) = foreground - .spawn(|me, _| (me.harness_working_dir.clone(), me.task_id)) - .await - else { + let Ok(Some(session)) = Self::current_oz_hook_session(foreground).await else { return; }; let reason = match status { @@ -1149,65 +1202,19 @@ impl AgentDriver { | SDKConversationOutputStatus::Blocked { .. } => SessionEndReason::Failed, SDKConversationOutputStatus::Cancelled { .. } => SessionEndReason::Cancelled, }; - let run_id = task_id - .map(|id| id.to_string()) - .unwrap_or_else(|| Uuid::new_v4().to_string()); - runtime + let _ = runtime .observe(OzHookEvent { invocation_id: Uuid::new_v4().to_string(), tool_use_id: None, payload: HookPayloadTemplate { - context: HookPayloadContext { - session_id: run_id.clone(), - run_id, - conversation_id: String::new(), - cwd: cwd.to_string_lossy().into_owned(), - model: String::new(), - permission_mode: "supervised".into(), - }, + context: session.payload_context, event: HookEventFields::SessionEnd { reason }, }, }) + .with_timeout(Duration::from_secs(3)) .await; runtime.cancel(crate::ai::agent_sdk::hooks::runtime::OzHookCancellationScope::Session); } - async fn observe_oz_prompt( - runtime: Option<&Arc>, - prompt: &AgentRunPrompt, - foreground: &ModelSpawner, - ) { - let (Some(runtime), AgentRunPrompt::Local(prompt)) = (runtime, prompt) else { - return; - }; - let Ok((cwd, task_id)) = foreground - .spawn(|me, _| (me.harness_working_dir.clone(), me.task_id)) - .await - else { - return; - }; - let run_id = task_id - .map(|id| id.to_string()) - .unwrap_or_else(|| Uuid::new_v4().to_string()); - runtime - .observe(OzHookEvent { - invocation_id: Uuid::new_v4().to_string(), - tool_use_id: None, - payload: HookPayloadTemplate { - context: HookPayloadContext { - session_id: run_id.clone(), - run_id, - conversation_id: String::new(), - cwd: cwd.to_string_lossy().into_owned(), - model: String::new(), - permission_mode: "supervised".into(), - }, - event: HookEventFields::user_prompt( - HookRedactor::new([]).redact_text(prompt, MAX_PROMPT_BYTES), - ), - }, - }) - .await; - } async fn observe_oz_stop( runtime: Option<&Arc>, @@ -1217,10 +1224,7 @@ impl AgentDriver { let Some(runtime) = runtime else { return; }; - let Ok((cwd, task_id)) = foreground - .spawn(|me, _| (me.harness_working_dir.clone(), me.task_id)) - .await - else { + let Ok(Some(session)) = Self::current_oz_hook_session(foreground).await else { return; }; let turn_status = match status { @@ -1229,22 +1233,12 @@ impl AgentDriver { SDKConversationOutputStatus::Cancelled { .. } => TurnStatus::Idle, SDKConversationOutputStatus::Blocked { .. } => TurnStatus::Blocked, }; - let run_id = task_id - .map(|id| id.to_string()) - .unwrap_or_else(|| Uuid::new_v4().to_string()); runtime .observe(OzHookEvent { invocation_id: Uuid::new_v4().to_string(), tool_use_id: None, payload: HookPayloadTemplate { - context: HookPayloadContext { - session_id: run_id.clone(), - run_id, - conversation_id: String::new(), - cwd: cwd.to_string_lossy().into_owned(), - model: String::new(), - permission_mode: "supervised".into(), - }, + context: session.payload_context, event: HookEventFields::Stop { turn_status }, }, }) @@ -2354,7 +2348,7 @@ impl AgentDriver { .await?; // For the Oz harness only: set up MCP servers, model overrides, and profile information. - if matches!(&task.harness, HarnessKind::Oz) { + if task.harness.uses_oz_lifecycle_hooks() { let mcp_specs = task.mcp_specs.clone(); let managed_mcp_client = foreground .spawn(|_, ctx| ServerApiProvider::as_ref(ctx).get_managed_mcp_client()) @@ -2644,8 +2638,13 @@ impl AgentDriver { &foreground, ) .await?; - let oz_hook_runtime = Self::initialize_oz_hook_runtime(&foreground).await?; - Self::observe_oz_prompt(oz_hook_runtime.as_ref(), &task.prompt, &foreground).await; + Self::prepare_oz_hook_conversation(&foreground).await?; + let model = task + .model + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| "default".into()); + let oz_hook_runtime = Self::initialize_oz_hook_runtime(model, &foreground).await?; let status_rx = foreground .spawn(move |me, ctx| me.execute_run(task.prompt, ctx)) @@ -2664,8 +2663,14 @@ impl AgentDriver { &foreground, ) .await?; - Self::observe_oz_stop(oz_hook_runtime.as_ref(), &conversation_status, &foreground) + if !matches!(conversation_status, SDKConversationOutputStatus::Success) { + Self::observe_oz_stop( + oz_hook_runtime.as_ref(), + &conversation_status, + &foreground, + ) .await; + } log::info!( "Ambient agent Oz lifecycle: event=run_exit_received idle_on_complete_elapsed_or_not_configured=true next=terminal_teardown_after_flush" @@ -3779,7 +3784,7 @@ impl AgentDriver { run_exit = run_exit.with_wait(wait); } } - let restored_conversation_id = self.restored_conversation_id; + let restored_conversation_id = self.run_conversation_id; // ServerSide prompts enter the agent view and emit // `CloudModeSetupPhaseEnded` to tear down the Cloud Mode Setup V2 chip. diff --git a/app/src/ai/agent_sdk/driver/harness/mod.rs b/app/src/ai/agent_sdk/driver/harness/mod.rs index 1e0c235a118..560ca407841 100644 --- a/app/src/ai/agent_sdk/driver/harness/mod.rs +++ b/app/src/ai/agent_sdk/driver/harness/mod.rs @@ -235,6 +235,12 @@ pub(crate) enum HarnessKind { } impl HarnessKind { + pub(crate) const fn uses_oz_lifecycle_hooks(&self) -> bool { + match self { + Self::Oz => true, + Self::ThirdParty(_) | Self::Unsupported(_) => false, + } + } /// Corresponding [`Harness`] enum value. pub(crate) fn harness(&self) -> Harness { match self { diff --git a/app/src/ai/agent_sdk/driver_tests.rs b/app/src/ai/agent_sdk/driver_tests.rs index 29059178138..3851b1660ec 100644 --- a/app/src/ai/agent_sdk/driver_tests.rs +++ b/app/src/ai/agent_sdk/driver_tests.rs @@ -41,6 +41,7 @@ use crate::ai::agent::{ AIAgentOutputMessage, ArtifactCreatedData, CancellationReason, MessageId, RenderableAIError, UploadArtifactResult, }; +use crate::ai::agent_sdk::driver::harness::harness_kind; use crate::ai::agent_sdk::task_env_vars; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::orchestration_events::{ @@ -55,6 +56,22 @@ use crate::ai::skills::SkillManager; use crate::test_util::assert_eventually; use crate::test_util::terminal::{add_window_with_terminal, initialize_app_for_terminal_view}; +#[test] +fn oz_hooks_runtime_is_not_enabled_for_third_party_harnesses() { + for harness in [ + Harness::Claude, + Harness::Codex, + Harness::Gemini, + Harness::OpenCode, + ] { + assert!( + !harness_kind(harness).unwrap().uses_oz_lifecycle_hooks(), + "{harness} must retain its native hook behavior" + ); + } + assert!(harness_kind(Harness::Oz).unwrap().uses_oz_lifecycle_hooks()); +} + // ── IdleTimeoutSender tests ────────────────────────────────────────────────────── #[test] diff --git a/app/src/ai/agent_sdk/hooks/mod.rs b/app/src/ai/agent_sdk/hooks/mod.rs index f7c22fb1cc5..404689e6860 100644 --- a/app/src/ai/agent_sdk/hooks/mod.rs +++ b/app/src/ai/agent_sdk/hooks/mod.rs @@ -29,6 +29,7 @@ pub(crate) struct OzHookSession { pub(crate) protocol_context: warp_multi_agent_api::OzHookContext, pub(crate) payload_context: payload::HookPayloadContext, pub(crate) redactor: redaction::HookRedactor, + pub(crate) is_driver_owned: bool, } #[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 88ee53c453d..742e92fad9e 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -456,11 +456,12 @@ impl BlocklistAIActionModel { #[cfg(not(target_family = "wasm"))] pub(crate) fn set_oz_hook_session( &mut self, + conversation_id: AIConversationId, session: Option, ctx: &mut ModelContext, ) { self.executor.update(ctx, |executor, _| { - executor.set_oz_hook_session(session); + executor.set_oz_hook_session(conversation_id, session); }); } diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 31c181c7b79..25f62db4a26 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -308,7 +308,7 @@ pub struct BlocklistAIActionExecutor { terminal_model: Arc>, team_context_resolver: TeamContextResolver, #[cfg(not(target_family = "wasm"))] - oz_hook_session: Option, + oz_hook_sessions: std::collections::HashMap, #[cfg(not(target_family = "wasm"))] oz_hook_approved_actions: std::collections::HashSet, } @@ -407,15 +407,23 @@ impl BlocklistAIActionExecutor { ask_user_question_executor, wait_for_events_executor, #[cfg(not(target_family = "wasm"))] - oz_hook_session: None, + oz_hook_sessions: Default::default(), #[cfg(not(target_family = "wasm"))] oz_hook_approved_actions: Default::default(), } } #[cfg(not(target_family = "wasm"))] - pub(crate) fn set_oz_hook_session(&mut self, session: Option) { - self.oz_hook_session = session; + pub(crate) fn set_oz_hook_session( + &mut self, + conversation_id: AIConversationId, + session: Option, + ) { + if let Some(session) = session { + self.oz_hook_sessions.insert(conversation_id, session); + } else { + self.oz_hook_sessions.remove(&conversation_id); + } } pub fn async_executing_action(&self, action_id: &AIAgentActionId) -> Option<&AIAgentAction> { @@ -708,7 +716,9 @@ impl BlocklistAIActionExecutor { } #[cfg(not(target_family = "wasm"))] - if !oz_preflight_complete && let Some(session) = self.oz_hook_session.clone() { + if !oz_preflight_complete + && let Some(session) = self.oz_hook_sessions.get(&conversation_id).cloned() + { let (tool_name, tool_input) = local_action_payload(&action.action, &session.redactor); let tool_use_id = action.id.to_string(); let mut payload_context = session.payload_context.clone(); @@ -1025,7 +1035,7 @@ impl BlocklistAIActionExecutor { ctx: &mut ModelContext, ) { #[cfg(not(target_family = "wasm"))] - if let Some(session) = self.oz_hook_session.clone() { + if let Some(session) = self.oz_hook_sessions.get(&running.conversation_id).cloned() { let action_id = running.action.id.clone(); let tool_use_id = action_id.to_string(); let (tool_name, tool_input) = @@ -1114,7 +1124,7 @@ impl BlocklistAIActionExecutor { } if let Some(running) = self.async_executing_actions.remove(action_id) { #[cfg(not(target_family = "wasm"))] - if let Some(session) = &self.oz_hook_session { + if let Some(session) = self.oz_hook_sessions.get(&running.conversation_id) { session .runtime .cancel(OzHookCancellationScope::Tool(action_id.to_string())); diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 93a005f662c..86054a5608e 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -10,7 +10,7 @@ pub(super) mod shared_session; mod slash_command; use std::collections::{HashMap, HashSet}; #[cfg(not(target_family = "wasm"))] -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; @@ -26,7 +26,7 @@ pub use slash_command::*; use warp_core::assertions::safe_assert; use warp_errors::report_error; use warp_multi_agent_api::{Task, ToolType, message}; -use warpui::r#async::{SpawnedFutureHandle, Timer}; +use warpui::r#async::{FutureExt as _, SpawnedFutureHandle, Timer}; use warpui::{ AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity, WeakViewHandle, }; @@ -58,15 +58,31 @@ use crate::ai::agent_events::AgentMessageEventMetadata; #[cfg(not(target_family = "wasm"))] use crate::ai::agent_sdk::ClaudeHarness; #[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::config::discover_hook_config; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::payload::{HookEventFields, HookPayloadTemplate}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::payload::{ + HookPayloadContext, SessionEndReason, SessionStartSource, TurnStatus, +}; +#[cfg(not(target_family = "wasm"))] use crate::ai::agent_sdk::hooks::protocol::{ event_from_protocol, failed_result, result_for_observation, result_for_pre_tool, }; #[cfg(not(target_family = "wasm"))] -use crate::ai::agent_sdk::hooks::runtime::OzHookCancellationScope; +use crate::ai::agent_sdk::hooks::redaction::HookRedactor; #[cfg(not(target_family = "wasm"))] -use crate::ai::agent_sdk::hooks::runtime::OzPreToolUseEvent; +use crate::ai::agent_sdk::hooks::runtime::{ + OzHookCancellationScope, OzHookEvent, OzHookRuntime, OzHookRuntimeService, OzPreToolUseEvent, +}; #[cfg(not(target_family = "wasm"))] -use crate::ai::agent_sdk::hooks::{HookEventName, OzHookSession, PAYLOAD_SCHEMA_VERSION}; +use crate::ai::agent_sdk::hooks::trust::{ + DenyProjectHookTrust, HookTrustStore, PersistentHookTrustStore, +}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::{ + HookEventName, MAX_PROMPT_BYTES, OzHookSession, PAYLOAD_SCHEMA_VERSION, +}; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, @@ -378,12 +394,14 @@ pub struct BlocklistAIController { )>, >, #[cfg(not(target_family = "wasm"))] - oz_hook_session: Option, + oz_hook_sessions: HashMap, #[cfg(not(target_family = "wasm"))] pending_oz_hook_results: HashMap>, #[cfg(not(target_family = "wasm"))] oz_hook_compatible_streams: HashSet, #[cfg(not(target_family = "wasm"))] + oz_hook_action_streams: HashSet, + #[cfg(not(target_family = "wasm"))] oz_hook_invocations: HashSet<(AIConversationId, String)>, #[cfg(not(target_family = "wasm"))] oz_hook_results_by_invocation: @@ -453,6 +471,186 @@ impl InputQuery { } impl BlocklistAIController { + #[cfg(not(target_family = "wasm"))] + fn end_local_oz_hook_session( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + let Some(session) = self.oz_hook_sessions.get(&conversation_id) else { + return; + }; + if session.is_driver_owned { + return; + } + let session = self + .oz_hook_sessions + .remove(&conversation_id) + .expect("hook session was checked above"); + self.action_model.update(ctx, |model, ctx| { + model.set_oz_hook_session(conversation_id, None, ctx); + }); + let runtime = Arc::clone(&session.runtime); + let event = OzHookEvent { + invocation_id: uuid::Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: session.payload_context, + event: HookEventFields::SessionEnd { + reason: SessionEndReason::Shutdown, + }, + }, + }; + ctx.spawn( + async move { + let _ = runtime + .observe(event) + .with_timeout(Duration::from_secs(3)) + .await; + runtime.cancel(OzHookCancellationScope::Session); + }, + |_, _, _| {}, + ); + } + #[cfg(not(target_family = "wasm"))] + fn defer_oz_stop_if_needed( + &mut self, + stream_id: ResponseStreamId, + finished_event: warp_multi_agent_api::response_event::StreamFinished, + conversation_id: AIConversationId, + did_input_contain_user_query: bool, + response_stream: ModelHandle, + ctx: &mut ModelContext, + ) -> bool { + if !response_stream.as_ref(ctx).is_completion_deferred() { + return false; + } + let has_actions = self.oz_hook_action_streams.contains(&stream_id) + || BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .is_some_and(|conversation| { + conversation + .new_exchange_ids_for_response(&stream_id) + .filter_map(|exchange_id| conversation.exchange_with_id(exchange_id)) + .filter_map(|exchange| exchange.output_status.output()) + .any(|output| output.get().actions().next().is_some()) + }); + let Some(session) = self.oz_hook_sessions.get(&conversation_id) else { + response_stream.update(ctx, |stream, ctx| { + stream.finish_deferred_completion(ctx); + }); + return false; + }; + if has_actions { + response_stream.update(ctx, |stream, ctx| { + stream.finish_deferred_completion(ctx); + }); + return false; + } + let runtime = Arc::clone(&session.runtime); + let event = OzHookEvent { + invocation_id: uuid::Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: session.payload_context.clone(), + event: HookEventFields::Stop { + turn_status: TurnStatus::Completed, + }, + }, + }; + ctx.spawn( + async move { + runtime.observe(event).await; + }, + move |me, _, ctx| { + me.handle_response_stream_finished( + &stream_id, + finished_event, + conversation_id, + did_input_contain_user_query, + ctx, + ); + response_stream.update(ctx, |stream, ctx| { + stream.finish_deferred_completion(ctx); + }); + }, + ); + true + } + #[cfg(not(target_family = "wasm"))] + fn initialize_local_oz_hook_session( + &mut self, + conversation_id: AIConversationId, + model: String, + ctx: &mut ModelContext, + ) -> Option { + if !FeatureFlag::OzLifecycleHooks.is_enabled() + || self.oz_hook_sessions.contains_key(&conversation_id) + { + return None; + } + let cwd = self + .active_session + .as_ref(ctx) + .current_working_directory()? + .clone(); + let trust_store: Box = match PersistentHookTrustStore::load_default() { + Ok(store) => Box::new(store), + Err(error) => { + log::warn!("Failed to load Oz hook trust store: {error}"); + Box::new(DenyProjectHookTrust) + } + }; + let config = discover_hook_config(Path::new(&cwd), trust_store.as_ref()); + for diagnostic in config.diagnostics.iter() { + log::warn!( + "Oz hook configuration diagnostic: kind={:?} path={} hash_present={}", + diagnostic.kind, + diagnostic.path.display(), + diagnostic.definition_hash.is_some() + ); + } + let enabled_events = config + .enabled_events() + .map(|event| event.protocol_value().into()) + .collect(); + let runtime: Arc = Arc::new(OzHookRuntimeService::new(config)); + let payload_context = HookPayloadContext { + session_id: uuid::Uuid::new_v4().to_string(), + run_id: uuid::Uuid::new_v4().to_string(), + conversation_id: conversation_id.to_string(), + cwd, + model, + permission_mode: "supervised".into(), + }; + let source = if BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .is_some_and(|conversation| conversation.exchange_count() > 0) + { + SessionStartSource::Resume + } else { + SessionStartSource::Startup + }; + let session = OzHookSession { + runtime: Arc::clone(&runtime), + protocol_context: warp_multi_agent_api::OzHookContext { + enabled_events, + supported_payload_schema_versions: vec![PAYLOAD_SCHEMA_VERSION.into()], + }, + payload_context: payload_context.clone(), + redactor: HookRedactor::new([]), + is_driver_owned: false, + }; + self.set_oz_hook_session(conversation_id, Some(session), ctx); + Some(OzHookEvent { + invocation_id: uuid::Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: payload_context, + event: HookEventFields::SessionStart { source }, + }, + }) + } /// Returns the bundled-skill catalog origin for this controller's active session. pub fn skill_path_origin(&self, ctx: &AppContext) -> SkillPathOrigin { SessionContext::from_session(self.active_session.as_ref(ctx), ctx).skill_path_origin() @@ -639,6 +837,8 @@ impl BlocklistAIController { ctx, ); } + #[cfg(not(target_family = "wasm"))] + me.end_local_oz_hook_session(*conversation_id, ctx); }); // Subscribe to the orchestration event service to inject events // (e.g. MessagesReceivedFromAgents) into conversations that receive inter-agent messages. @@ -678,12 +878,14 @@ impl BlocklistAIController { pending_passive_follow_ups: HashSet::new(), pending_passive_suggestion_results: HashMap::new(), #[cfg(not(target_family = "wasm"))] - oz_hook_session: None, + oz_hook_sessions: Default::default(), #[cfg(not(target_family = "wasm"))] pending_oz_hook_results: HashMap::new(), #[cfg(not(target_family = "wasm"))] oz_hook_compatible_streams: HashSet::new(), #[cfg(not(target_family = "wasm"))] + oz_hook_action_streams: HashSet::new(), + #[cfg(not(target_family = "wasm"))] oz_hook_invocations: HashSet::new(), #[cfg(not(target_family = "wasm"))] oz_hook_results_by_invocation: HashMap::new(), @@ -2419,15 +2621,28 @@ impl BlocklistAIController { }); } + #[cfg(not(target_family = "wasm"))] + pub(crate) fn oz_hook_session( + &self, + conversation_id: AIConversationId, + ) -> Option { + self.oz_hook_sessions.get(&conversation_id).cloned() + } + #[cfg(not(target_family = "wasm"))] pub(crate) fn set_oz_hook_session( &mut self, + conversation_id: AIConversationId, session: Option, ctx: &mut ModelContext, ) { - self.oz_hook_session = session.clone(); + if let Some(session) = session.clone() { + self.oz_hook_sessions.insert(conversation_id, session); + } else { + self.oz_hook_sessions.remove(&conversation_id); + } self.action_model.update(ctx, |model, ctx| { - model.set_oz_hook_session(session, ctx); + model.set_oz_hook_session(conversation_id, session, ctx); }); } @@ -2453,7 +2668,7 @@ impl BlocklistAIController { } return; } - let Some(session) = &self.oz_hook_session else { + let Some(session) = self.oz_hook_sessions.get(&conversation_id) else { return; }; let event = match event_from_protocol(&action) { @@ -2584,6 +2799,13 @@ impl BlocklistAIController { .expect("Conversation exists- was just created.") } + pub(crate) fn start_oz_hook_conversation( + &self, + ctx: &mut ModelContext, + ) -> AIConversationId { + self.start_new_conversation_for_request(ctx).id() + } + /// Attempts to send a request to the AI model API. Adds context to the input if it /// contains a user query. Returns `Err` if the AI input was not able to be sent due to an /// existing in-flight request. Emits an event containing a receiver for the AI's output. @@ -2729,9 +2951,48 @@ impl BlocklistAIController { request_params.parent_agent_id = parent_agent_id; request_params.agent_name = agent_name; #[cfg(not(target_family = "wasm"))] - if let Some(session) = &self.oz_hook_session { + let startup_hook = request_input + .all_inputs() + .any(AIAgentInput::is_user_query) + .then(|| { + self.initialize_local_oz_hook_session( + conversation_id, + request_params.model.to_string(), + ctx, + ) + }) + .flatten(); + #[cfg(not(target_family = "wasm"))] + if let Some(session) = self.oz_hook_sessions.get(&conversation_id) { request_params.oz_hook_context = Some(session.protocol_context.clone()); } + #[cfg(not(target_family = "wasm"))] + let prompt_hook = self + .oz_hook_sessions + .get(&conversation_id) + .and_then(|session| { + request_input.all_inputs().find_map(|input| { + let AIAgentInput::UserQuery { query, .. } = input else { + return None; + }; + let mut events = startup_hook.clone().into_iter().collect::>(); + events.push(OzHookEvent { + invocation_id: uuid::Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: session.payload_context.clone(), + event: HookEventFields::user_prompt( + session.redactor.redact_text(query, MAX_PROMPT_BYTES), + ), + }, + }); + Some((Arc::clone(&session.runtime), events)) + }) + }); + #[cfg(not(target_family = "wasm"))] + let should_defer_for_prompt_hook = prompt_hook.is_some(); + #[cfg(target_family = "wasm")] + let should_defer_for_prompt_hook = false; let server_conversation_token_for_identifiers = conversation_data.server_conversation_token.clone(); @@ -2745,13 +3006,22 @@ impl BlocklistAIController { client_exchange_id: None, model_id: Some(request_params.model.clone()), }; - ResponseStream::new( - request_params.clone(), - ai_identifiers, - recovery, - team_scope, - ctx, - ) + if should_defer_for_prompt_hook { + ResponseStream::new_deferred( + request_params.clone(), + ai_identifiers, + recovery, + team_scope, + ) + } else { + ResponseStream::new( + request_params.clone(), + ai_identifiers, + recovery, + team_scope, + ctx, + ) + } }); let response_stream_id = response_stream.as_ref(ctx).id().clone(); let response_stream_clone = response_stream.clone(); @@ -2766,6 +3036,20 @@ impl BlocklistAIController { ctx, ); }); + #[cfg(not(target_family = "wasm"))] + if let Some((runtime, events)) = prompt_hook { + let response_stream = response_stream.clone(); + ctx.spawn( + async move { + for event in events { + runtime.observe(event).await; + } + }, + move |_, _, ctx| { + response_stream.update(ctx, |stream, ctx| stream.start_deferred(ctx)); + }, + ); + } for input in request_input.all_inputs() { if let AIAgentInput::UserQuery { @@ -3146,7 +3430,7 @@ impl BlocklistAIController { match event { warp_multi_agent_api::response_event::Type::Init(init_event) => { #[cfg(not(target_family = "wasm"))] - if self.oz_hook_session.is_some() { + if self.oz_hook_sessions.contains_key(&conversation_id) { if init_event .supported_oz_hook_payload_schema_versions .iter() @@ -3189,6 +3473,17 @@ impl BlocklistAIController { warp_multi_agent_api::response_event::Type::Finished( finished_event, ) => { + #[cfg(not(target_family = "wasm"))] + if self.defer_oz_stop_if_needed( + stream_id.clone(), + finished_event.clone(), + conversation_id, + did_input_contain_user_query, + response_stream.clone(), + ctx, + ) { + return; + } self.handle_response_stream_finished( &stream_id, finished_event, @@ -3207,6 +3502,7 @@ impl BlocklistAIController { ), ) = client_action.action.take() { + self.oz_hook_action_streams.insert(stream_id.clone()); if !self.oz_hook_compatible_streams.contains(&stream_id) { report_error!( "Received Oz hook action before successful schema negotiation", @@ -3354,7 +3650,7 @@ impl BlocklistAIController { let history_model = BlocklistAIHistoryModel::handle(ctx); #[cfg(not(target_family = "wasm"))] if cancellation.is_some() { - if let Some(session) = &self.oz_hook_session { + if let Some(session) = self.oz_hook_sessions.get(&conversation_id) { for invocation_key in self .oz_hook_invocations .iter() @@ -3369,6 +3665,7 @@ impl BlocklistAIController { } self.pending_oz_hook_results.remove(&conversation_id); self.oz_hook_compatible_streams.remove(&stream_id); + self.oz_hook_action_streams.remove(&stream_id); } let Some(conversation) = history_model.as_ref(ctx).conversation(&conversation_id) else { @@ -3463,6 +3760,7 @@ impl BlocklistAIController { #[cfg(not(target_family = "wasm"))] { self.oz_hook_compatible_streams.remove(&stream_id); + self.oz_hook_action_streams.remove(&stream_id); self.send_pending_oz_hook_results(conversation_id, ctx); } diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index 6dd9e724434..449e5a03e4b 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -289,6 +289,8 @@ pub struct ResponseStream { /// Whether a retry is parked waiting for a backoff or for connectivity. While set, /// completion of the failed attempt's underlying stream is ignored. deferred_retry_pending: bool, + completion_deferred: bool, + completion_waiting: bool, /// Unique, internal id for the current request. /// @@ -298,6 +300,7 @@ pub struct ResponseStream { /// Note this is unique compared to `id`; this is unique across retry requests while the response /// stream id remains stable. current_request_id: Option, + pending_start: Option<(Uuid, oneshot::Receiver<()>)>, /// Captured once at construction, so retries keep the team the request started on. team_scope: RequestTeamScope, @@ -343,7 +346,10 @@ impl ResponseStream { stream_finished_received: false, error_event_emitted: false, deferred_retry_pending: false, + completion_deferred: false, + completion_waiting: false, current_request_id: Some(Uuid::new_v4()), + pending_start: None, team_scope: RequestTeamScope::from_scope(&TeamlessScopeForTest), } } @@ -375,11 +381,83 @@ impl ResponseStream { stream_finished_received: false, error_event_emitted: false, deferred_retry_pending: false, + completion_deferred: false, + completion_waiting: false, current_request_id: Some(request_id), + pending_start: None, team_scope, } } + pub fn new_deferred( + params: api::RequestParams, + ai_identifiers: AIIdentifiers, + recovery: RecoveryBudget, + team_scope: RequestTeamScope, + ) -> Self { + let (cancellation_tx, cancellation_rx) = oneshot::channel(); + let request_id = Uuid::new_v4(); + Self { + id: ResponseStreamId(Uuid::new_v4().to_string()), + params, + start_time: Local::now(), + time_to_latest_event: TimeDelta::seconds(0), + cancellation_tx: Some(cancellation_tx), + recovery, + retries_sent: 0, + original_error: None, + has_received_client_actions: false, + ai_identifiers, + pending_resume: None, + stream_finished_received: false, + error_event_emitted: false, + deferred_retry_pending: false, + completion_deferred: false, + completion_waiting: false, + current_request_id: None, + pending_start: Some((request_id, cancellation_rx)), + team_scope, + } + } + + pub fn start_deferred(&mut self, ctx: &mut ModelContext) { + let Some((request_id, cancellation_rx)) = self.take_pending_start() else { + return; + }; + Self::spawn_request( + request_id, + self.params.clone(), + self.team_scope, + cancellation_rx, + ctx, + ); + self.current_request_id = Some(request_id); + } + + pub fn finish_deferred_completion(&mut self, ctx: &mut ModelContext) { + if !std::mem::take(&mut self.completion_deferred) { + return; + } + if std::mem::take(&mut self.completion_waiting) + && let Some(request_id) = self.current_request_id + { + self.on_response_stream_complete(request_id, ctx); + } + } + + pub fn is_completion_deferred(&self) -> bool { + self.completion_deferred + } + + fn take_pending_start(&mut self) -> Option<(Uuid, oneshot::Receiver<()>)> { + if self.cancellation_tx.is_none() { + self.pending_start.take(); + None + } else { + self.pending_start.take() + } + } + pub fn id(&self) -> &ResponseStreamId { &self.id } @@ -433,6 +511,8 @@ impl ResponseStream { self.stream_finished_received = false; self.error_event_emitted = false; self.deferred_retry_pending = false; + self.completion_deferred = false; + self.completion_waiting = false; // A retry supersedes any resume this stream had scheduled. Unreachable today (the // eventsource closes on its first error, so a `Resume` decision is never followed by // another error on the same stream), but that depends on a transport detail several @@ -808,6 +888,18 @@ impl ResponseStream { finished_event.reason, Some(warp_multi_agent_api::response_event::stream_finished::Reason::Done(_)) | None ) { + #[cfg(not(target_family = "wasm"))] + { + self.completion_deferred = self + .params + .oz_hook_context + .as_ref() + .is_some_and(|context| { + context + .enabled_events + .contains(&(maa_api::OzHookEvent::Stop as i32)) + }); + } // Emit retry success telemetry if this was a successful completion after retries if self.retries_sent > 0 && let Some(original_error) = &self.original_error { @@ -851,6 +943,10 @@ impl ResponseStream { if self.deferred_retry_pending { return; } + if self.completion_deferred { + self.completion_waiting = true; + return; + } // The server always sends a StreamFinished event before ending the response, // but a transport cut between chunks surfaces as a clean EOF. Synthesize the diff --git a/app/src/ai/blocklist/controller/response_stream_tests.rs b/app/src/ai/blocklist/controller/response_stream_tests.rs index baa558e1573..b9b5a78104d 100644 --- a/app/src/ai/blocklist/controller/response_stream_tests.rs +++ b/app/src/ai/blocklist/controller/response_stream_tests.rs @@ -7,12 +7,18 @@ use super::{FailReason, MAX_RECOVERY_ATTEMPTS, RecoveryAction, RecoveryBudget, r #[cfg(not(target_family = "wasm"))] use super::{ResponseStream, ResponseStreamId}; #[cfg(not(target_family = "wasm"))] +use crate::ai::agent::AIIdentifiers; +#[cfg(not(target_family = "wasm"))] use crate::ai::agent::api::RequestParams; // `agent_sdk` (and so the driver's recovery deadline) is native-only. #[cfg(not(target_family = "wasm"))] use crate::ai::agent_sdk::driver::AUTO_RESUME_TIMEOUT; #[cfg(not(target_family = "wasm"))] use crate::server::retry_strategies::backoff_after_attempts; +#[cfg(not(target_family = "wasm"))] +use crate::server::team_scope::RequestTeamScope; +#[cfg(not(target_family = "wasm"))] +use crate::workspaces::user_workspaces::TeamlessScopeForTest; // Argument order: has_received_client_actions, is_recoverable, recovery, is_online. @@ -225,6 +231,37 @@ fn spending_an_attempt_preserves_resume_eligibility() { ); } +#[cfg(not(target_family = "wasm"))] +fn deferred_stream() -> ResponseStream { + ResponseStream::new_deferred( + RequestParams::new_for_test(), + AIIdentifiers::default(), + RecoveryBudget::fresh(), + RequestTeamScope::from_scope(&TeamlessScopeForTest), + ) +} + +#[cfg(not(target_family = "wasm"))] +#[test] +fn deferred_request_start_is_single_use() { + let mut stream = deferred_stream(); + + assert!(stream.current_request_id.is_none()); + assert!(stream.take_pending_start().is_some()); + assert!(stream.take_pending_start().is_none()); +} + +#[cfg(not(target_family = "wasm"))] +#[test] +fn cancelled_deferred_request_cannot_start() { + let mut stream = deferred_stream(); + stream.cancellation_tx.take(); + + assert!(stream.take_pending_start().is_none()); + assert!(stream.pending_start.is_none()); + assert!(stream.current_request_id.is_none()); +} + #[cfg(not(target_family = "wasm"))] #[test] fn the_recovery_backoff_fits_inside_the_cloud_run_recovery_window() { From 2cafca07b6698107daf93b973cd781f6d57ff18b Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 14:56:39 +0000 Subject: [PATCH 07/11] Fix autonomous action permission classification --- app/src/ai/blocklist/action_model/execute.rs | 37 ++++++++++++------- .../blocklist/action_model/execute_tests.rs | 20 ++++++++++ 2 files changed, 44 insertions(+), 13 deletions(-) diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 25f62db4a26..62cc52dd091 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -272,6 +272,23 @@ impl AsyncExecutingAction { ) } } +fn native_permission_for_action( + is_user_initiated: bool, + can_auto_execute: bool, + is_agent_autonomous: bool, + is_request_command_output: bool, +) -> NativePermission { + if !is_user_initiated && !can_auto_execute && is_agent_autonomous && is_request_command_output { + NativePermission::Deny + } else if !(is_user_initiated + || can_auto_execute + || (is_agent_autonomous && is_request_command_output)) + { + NativePermission::Prompt + } else { + NativePermission::Allow + } +} pub struct BlocklistAIActionExecutor { shell_command_executor: ModelHandle, @@ -675,19 +692,13 @@ impl BlocklistAIActionExecutor { let can_auto_execute = self.should_autoexecute(input, ctx); let is_agent_autonomous = AppExecutionMode::as_ref(ctx).is_autonomous(); - // The agent cannot auto execute and either: - // - the agent is interactive, OR - // - the agent is autonomous and the action was not requesting command output - let needs_confirmation = !(is_user_initiated - || can_auto_execute - || (is_agent_autonomous && action.action.is_request_command_output())); - let native_permission = if !is_user_initiated && !can_auto_execute && is_agent_autonomous { - NativePermission::Deny - } else if needs_confirmation { - NativePermission::Prompt - } else { - NativePermission::Allow - }; + let native_permission = native_permission_for_action( + is_user_initiated, + can_auto_execute, + is_agent_autonomous, + action.action.is_request_command_output(), + ); + let needs_confirmation = native_permission == NativePermission::Prompt; if native_permission == NativePermission::Deny { // It must be the case that the autonomous agent is requesting a denylisted command. if let AIAgentActionType::RequestCommandOutput { command, .. } = &action.action { diff --git a/app/src/ai/blocklist/action_model/execute_tests.rs b/app/src/ai/blocklist/action_model/execute_tests.rs index 5261c355345..ab780954bc6 100644 --- a/app/src/ai/blocklist/action_model/execute_tests.rs +++ b/app/src/ai/blocklist/action_model/execute_tests.rs @@ -1,3 +1,23 @@ +mod native_permissions { + use super::super::native_permission_for_action; + use crate::ai::agent_sdk::hooks::permissions::NativePermission; + + #[test] + fn autonomous_non_command_action_that_cannot_autoexecute_prompts() { + assert_eq!( + native_permission_for_action(false, false, true, false), + NativePermission::Prompt + ); + } + + #[test] + fn autonomous_denylisted_command_is_denied() { + assert_eq!( + native_permission_for_action(false, false, true, true), + NativePermission::Deny + ); + } +} mod binary_detection { use std::io::Write as _; From 6e649815d8457c3d84f760c807be6c48ef218163 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:09:35 +0000 Subject: [PATCH 08/11] Complete local Oz turn lifecycle hooks --- app/src/ai/blocklist/controller.rs | 207 ++++++++++++------ .../blocklist/controller/response_stream.rs | 37 +++- .../controller/response_stream_tests.rs | 33 ++- app/src/ai/blocklist/controller_tests.rs | 9 + 4 files changed, 213 insertions(+), 73 deletions(-) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 86054a5608e..f7fc39c9a19 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -429,6 +429,13 @@ enum WhichTask { task_id: TaskId, }, } +#[cfg(not(target_family = "wasm"))] +fn starts_local_oz_hook_session(input: &AIAgentInput) -> bool { + matches!( + input, + AIAgentInput::UserQuery { .. } | AIAgentInput::ResumeConversation { .. } + ) +} #[derive(Debug, Clone, PartialEq, Eq)] enum LocalClaudeWakeTrigger { @@ -578,6 +585,117 @@ impl BlocklistAIController { true } #[cfg(not(target_family = "wasm"))] + fn defer_failed_oz_stop_if_needed( + &mut self, + stream_id: ResponseStreamId, + error: Arc, + conversation_id: AIConversationId, + response_stream: ModelHandle, + ctx: &mut ModelContext, + ) -> bool { + if !response_stream.as_ref(ctx).is_completion_deferred() { + return false; + } + let Some(session) = self.oz_hook_sessions.get(&conversation_id) else { + response_stream.update(ctx, |stream, ctx| { + stream.finish_deferred_completion(ctx); + }); + return false; + }; + if session.is_driver_owned { + response_stream.update(ctx, |stream, ctx| { + stream.finish_deferred_completion(ctx); + }); + return false; + } + let runtime = Arc::clone(&session.runtime); + let event = OzHookEvent { + invocation_id: uuid::Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: session.payload_context.clone(), + event: HookEventFields::Stop { + turn_status: TurnStatus::Failed, + }, + }, + }; + ctx.spawn( + async move { + runtime.observe(event).await; + }, + move |me, _, ctx| { + me.handle_response_stream_error( + error, + &stream_id, + conversation_id, + &response_stream, + ctx, + ); + response_stream.update(ctx, |stream, ctx| { + stream.finish_deferred_completion(ctx); + }); + }, + ); + true + } + + fn handle_response_stream_error( + &mut self, + error: Arc, + stream_id: &ResponseStreamId, + conversation_id: AIConversationId, + response_stream: &ModelHandle, + ctx: &mut ModelContext, + ) { + if matches!(error.as_ref(), AIApiError::QuotaLimit { .. }) { + // If the error is a quota limit, refresh workspace metadata so AI overages are + // immediately up to date. + TeamUpdateManager::handle(ctx).update(ctx, |team_update_manager, ctx| { + std::mem::drop(team_update_manager.refresh_workspace_metadata(ctx)); + }); + AIRequestUsageModel::handle(ctx).update(ctx, |model, ctx| { + model.enable_buy_credits_banner(ctx); + }); + } + // A resume scheduled for this failure keeps the conversation in the non-terminal + // TransientError status instead of Error. + + let recovery_pending = response_stream + .as_ref(ctx) + .should_resume_conversation_after_stream_finished(); + let mut renderable_error: RenderableAIError = (&error).into(); + if let RenderableAIError::Other { + will_attempt_resume, + waiting_for_network, + .. + } + | RenderableAIError::TransientNetworkError { + will_attempt_resume, + waiting_for_network, + .. + } = &mut renderable_error + { + // Rendering-only hints; state machine consumers key off the TransientError + // conversation status instead. + *will_attempt_resume |= recovery_pending; + if recovery_pending { + let network_status = NetworkStatus::as_ref(ctx); + *waiting_for_network = !network_status.is_online(); + } + } + + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.mark_response_stream_completed_with_error( + renderable_error, + recovery_pending, + stream_id, + conversation_id, + self.terminal_surface_id, + ctx, + ); + }); + } + #[cfg(not(target_family = "wasm"))] fn initialize_local_oz_hook_session( &mut self, conversation_id: AIConversationId, @@ -2953,7 +3071,7 @@ impl BlocklistAIController { #[cfg(not(target_family = "wasm"))] let startup_hook = request_input .all_inputs() - .any(AIAgentInput::is_user_query) + .any(starts_local_oz_hook_session) .then(|| { self.initialize_local_oz_hook_session( conversation_id, @@ -2967,15 +3085,17 @@ impl BlocklistAIController { request_params.oz_hook_context = Some(session.protocol_context.clone()); } #[cfg(not(target_family = "wasm"))] - let prompt_hook = self + let pre_request_hooks = self .oz_hook_sessions .get(&conversation_id) .and_then(|session| { - request_input.all_inputs().find_map(|input| { + let mut events = startup_hook.into_iter().collect::>(); + if let Some(query) = request_input.all_inputs().find_map(|input| { let AIAgentInput::UserQuery { query, .. } = input else { return None; }; - let mut events = startup_hook.clone().into_iter().collect::>(); + Some(query) + }) { events.push(OzHookEvent { invocation_id: uuid::Uuid::new_v4().to_string(), tool_use_id: None, @@ -2986,13 +3106,13 @@ impl BlocklistAIController { ), }, }); - Some((Arc::clone(&session.runtime), events)) - }) + } + (!events.is_empty()).then(|| (Arc::clone(&session.runtime), events)) }); #[cfg(not(target_family = "wasm"))] - let should_defer_for_prompt_hook = prompt_hook.is_some(); + let should_defer_for_pre_request_hooks = pre_request_hooks.is_some(); #[cfg(target_family = "wasm")] - let should_defer_for_prompt_hook = false; + let should_defer_for_pre_request_hooks = false; let server_conversation_token_for_identifiers = conversation_data.server_conversation_token.clone(); @@ -3006,7 +3126,7 @@ impl BlocklistAIController { client_exchange_id: None, model_id: Some(request_params.model.clone()), }; - if should_defer_for_prompt_hook { + if should_defer_for_pre_request_hooks { ResponseStream::new_deferred( request_params.clone(), ai_identifiers, @@ -3037,7 +3157,7 @@ impl BlocklistAIController { ); }); #[cfg(not(target_family = "wasm"))] - if let Some((runtime, events)) = prompt_hook { + if let Some((runtime, events)) = pre_request_hooks { let response_stream = response_stream.clone(); ctx.spawn( async move { @@ -3549,58 +3669,23 @@ impl BlocklistAIController { } } Err(e) => { - if matches!(e.as_ref(), AIApiError::QuotaLimit { .. }) { - // If the error is a quota limit, we want to refresh workspace metadata - // So the current state of AI overages is immediately up to date. - TeamUpdateManager::handle(ctx).update( - ctx, - |team_update_manager, ctx| { - std::mem::drop( - team_update_manager.refresh_workspace_metadata(ctx), - ); - }, - ); - AIRequestUsageModel::handle(ctx).update(ctx, |model, ctx| { - model.enable_buy_credits_banner(ctx); - }); - } - - // A resume scheduled for this failure keeps the conversation in - // the non-terminal TransientError status instead of Error. - let recovery_pending = response_stream - .as_ref(ctx) - .should_resume_conversation_after_stream_finished(); - let mut renderable_error: RenderableAIError = (&e).into(); - if let RenderableAIError::Other { - will_attempt_resume, - waiting_for_network, - .. - } - | RenderableAIError::TransientNetworkError { - will_attempt_resume, - waiting_for_network, - .. - } = &mut renderable_error - { - // Rendering-only hints; state machine consumers key off the - // TransientError conversation status instead. - *will_attempt_resume |= recovery_pending; - if recovery_pending { - let network_status = NetworkStatus::as_ref(ctx); - *waiting_for_network = !network_status.is_online(); - } + #[cfg(not(target_family = "wasm"))] + if self.defer_failed_oz_stop_if_needed( + stream_id.clone(), + Arc::clone(&e), + conversation_id, + response_stream.clone(), + ctx, + ) { + return; } - - history_model.update(ctx, |history_model, ctx| { - history_model.mark_response_stream_completed_with_error( - renderable_error, - recovery_pending, - &stream_id, - conversation_id, - self.terminal_surface_id, - ctx, - ); - }); + self.handle_response_stream_error( + e, + &stream_id, + conversation_id, + response_stream, + ctx, + ); } } } diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index 449e5a03e4b..bdfee639c69 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -448,6 +448,20 @@ impl ResponseStream { pub fn is_completion_deferred(&self) -> bool { self.completion_deferred } + #[cfg(not(target_family = "wasm"))] + fn defer_completion_for_oz_stop(&mut self) { + self.completion_deferred = self.params.oz_hook_context.as_ref().is_some_and(|context| { + context + .enabled_events + .contains(&(maa_api::OzHookEvent::Stop as i32)) + }); + } + #[cfg(not(target_family = "wasm"))] + fn defer_failed_completion_for_oz_stop(&mut self) { + if self.pending_resume.is_none() { + self.defer_completion_for_oz_stop(); + } + } fn take_pending_start(&mut self) -> Option<(Uuid, oneshot::Receiver<()>)> { if self.cancellation_tx.is_none() { @@ -760,6 +774,7 @@ impl ResponseStream { fn surface_grok_refresh_failure(&mut self, request_id: Uuid, ctx: &mut ModelContext) { let error = Arc::new(AIApiError::GrokSubscriptionTokenRefreshFailed); self.error_event_emitted = true; + self.defer_failed_completion_for_oz_stop(); self.report_request_failure( &error, NetworkStatus::as_ref(ctx).is_online(), @@ -847,6 +862,8 @@ impl ResponseStream { NetworkStatus::as_ref(ctx).is_online(), self.recovery.attempts_used(), ); + #[cfg(not(target_family = "wasm"))] + self.defer_failed_completion_for_oz_stop(); ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(Err( error, )))); @@ -889,17 +906,7 @@ impl ResponseStream { Some(warp_multi_agent_api::response_event::stream_finished::Reason::Done(_)) | None ) { #[cfg(not(target_family = "wasm"))] - { - self.completion_deferred = self - .params - .oz_hook_context - .as_ref() - .is_some_and(|context| { - context - .enabled_events - .contains(&(maa_api::OzHookEvent::Stop as i32)) - }); - } + self.defer_completion_for_oz_stop(); // Emit retry success telemetry if this was a successful completion after retries if self.retries_sent > 0 && let Some(original_error) = &self.original_error { @@ -928,6 +935,8 @@ impl ResponseStream { // Don't emit the error event, we're recovering in-request. return; } + #[cfg(not(target_family = "wasm"))] + self.defer_failed_completion_for_oz_stop(); ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(event))); } @@ -962,9 +971,15 @@ impl ResponseStream { ) { return; } + #[cfg(not(target_family = "wasm"))] + self.defer_failed_completion_for_oz_stop(); ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(Err( unexpected_eof, )))); + if self.completion_deferred { + self.completion_waiting = true; + return; + } } ctx.emit(ResponseStreamEvent::AfterStreamFinished { cancellation: None }); diff --git a/app/src/ai/blocklist/controller/response_stream_tests.rs b/app/src/ai/blocklist/controller/response_stream_tests.rs index b9b5a78104d..2c78fff9d69 100644 --- a/app/src/ai/blocklist/controller/response_stream_tests.rs +++ b/app/src/ai/blocklist/controller/response_stream_tests.rs @@ -5,7 +5,7 @@ use std::time::Duration; use super::apply_geap_refresh_to_params; use super::{FailReason, MAX_RECOVERY_ATTEMPTS, RecoveryAction, RecoveryBudget, recovery_action}; #[cfg(not(target_family = "wasm"))] -use super::{ResponseStream, ResponseStreamId}; +use super::{PendingResume, ResponseStream, ResponseStreamId}; #[cfg(not(target_family = "wasm"))] use crate::ai::agent::AIIdentifiers; #[cfg(not(target_family = "wasm"))] @@ -261,6 +261,37 @@ fn cancelled_deferred_request_cannot_start() { assert!(stream.pending_start.is_none()); assert!(stream.current_request_id.is_none()); } +#[cfg(not(target_family = "wasm"))] +#[test] +fn failed_turn_completion_is_deferred_when_oz_stop_is_enabled() { + let mut stream = deferred_stream(); + stream.params.oz_hook_context = Some(warp_multi_agent_api::OzHookContext { + enabled_events: vec![warp_multi_agent_api::OzHookEvent::Stop as i32], + supported_payload_schema_versions: vec![], + }); + + stream.defer_failed_completion_for_oz_stop(); + + assert!(stream.is_completion_deferred()); +} + +#[cfg(not(target_family = "wasm"))] +#[test] +fn failed_turn_completion_is_not_deferred_before_a_scheduled_resume() { + let mut stream = deferred_stream(); + stream.params.oz_hook_context = Some(warp_multi_agent_api::OzHookContext { + enabled_events: vec![warp_multi_agent_api::OzHookEvent::Stop as i32], + supported_payload_schema_versions: vec![], + }); + stream.pending_resume = Some(PendingResume::new_for_test( + RecoveryBudget::fresh().next_attempt(), + Duration::ZERO, + )); + + stream.defer_failed_completion_for_oz_stop(); + + assert!(!stream.is_completion_deferred()); +} #[cfg(not(target_family = "wasm"))] #[test] diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index ded5c0dc056..9ca6ae109f2 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -67,6 +67,15 @@ fn file_attachment(file_name: &str) -> PendingAttachment { }) } +#[cfg(not(target_family = "wasm"))] +#[test] +fn restored_conversation_resume_starts_a_local_oz_hook_session() { + assert!(super::starts_local_oz_hook_session( + &AIAgentInput::ResumeConversation { + context: vec![].into() + } + )); +} #[test] fn passive_suggestions_request_params_omit_ambient_agent_task_id() { App::test((), |mut app| async move { From b189fe711a8fac4f7d5ab3f9022d9f52e7bea642 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:23:56 +0000 Subject: [PATCH 09/11] Fix Oz hook stdin close race --- app/src/ai/agent_sdk/hooks/runtime.rs | 7 +++- app/src/ai/agent_sdk/hooks/runtime_tests.rs | 36 +++++++++++++++++++-- 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/app/src/ai/agent_sdk/hooks/runtime.rs b/app/src/ai/agent_sdk/hooks/runtime.rs index 7af57d2b7a5..3fca551f4bc 100644 --- a/app/src/ai/agent_sdk/hooks/runtime.rs +++ b/app/src/ai/agent_sdk/hooks/runtime.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::ffi::OsString; +use std::io::ErrorKind; use std::path::Path; use std::process::Stdio; use std::sync::Mutex; @@ -457,7 +458,11 @@ async fn run_command( }; let exit_code = status.code(); let outputs = async move { - let stdin_failed = !matches!(stdin_task.await, Ok(Ok(()))); + let stdin_failed = match stdin_task.await { + Ok(Ok(())) => false, + Ok(Err(error)) if error.kind() == ErrorKind::BrokenPipe => false, + Ok(Err(_)) | Err(_) => true, + }; let stdout = join_output(stdout_task, exit_code).await?; let stderr = join_output(stderr_task, exit_code).await?; Ok((stdout, stderr, stdin_failed)) diff --git a/app/src/ai/agent_sdk/hooks/runtime_tests.rs b/app/src/ai/agent_sdk/hooks/runtime_tests.rs index 5dbb0801f25..cb7b46cd747 100644 --- a/app/src/ai/agent_sdk/hooks/runtime_tests.rs +++ b/app/src/ai/agent_sdk/hooks/runtime_tests.rs @@ -47,18 +47,28 @@ fn event(invocation_id: &str, fields: HookEventFields) -> OzHookEvent { } } -fn pre_tool_event(invocation_id: &str) -> OzPreToolUseEvent { +fn pre_tool_event_with_input( + invocation_id: &str, + tool_input: super::super::redaction::RedactedValue, +) -> OzPreToolUseEvent { OzPreToolUseEvent::new(event( invocation_id, HookEventFields::PreToolUse { tool_name: "run_shell_command".into(), tool_use_id: "tool".into(), - tool_input: super::super::redaction::RedactedValue::object([] as [(&str, _); 0]), + tool_input, }, )) .unwrap() } +fn pre_tool_event(invocation_id: &str) -> OzPreToolUseEvent { + pre_tool_event_with_input( + invocation_id, + super::super::redaction::RedactedValue::object([] as [(&str, _); 0]), + ) +} + #[tokio::test] async fn oz_hooks_runtime_runs_matching_handlers_sequentially() { let output = tempfile::NamedTempFile::new().unwrap(); @@ -104,6 +114,28 @@ async fn oz_hooks_runtime_structured_deny_short_circuits_later_handlers() { assert_eq!(fs::read_to_string(marker.path()).unwrap(), ""); } +#[tokio::test] +async fn oz_hooks_runtime_accepts_output_from_hook_that_closes_stdin() { + let deny = r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"deny","permissionDecisionReason":"policy"}}"#; + let runtime = runtime_with_hooks(json!({ + "PreToolUse": [{"hooks": [{ + "type": "command", + "command": format!("exec 0<&-; sleep 0.1; printf '%s' '{deny}'") + }]}] + })); + let event = pre_tool_event_with_input( + "closed-stdin", + super::super::redaction::RedactedValue::String("x".repeat(128 * 1024)), + ); + + let decision = runtime.pre_tool_use(event).await; + + assert!(matches!( + decision, + OzPreToolUseDecision::Deny { ref reason, .. } if reason == "policy" + )); +} + #[tokio::test] async fn oz_hooks_runtime_exit_two_denies_only_pre_tool_use() { let runtime = runtime_with_hooks(json!({ From 029fcf2a9981a0064cdd5a6c1dd81930068f7b97 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:28:56 +0000 Subject: [PATCH 10/11] Complete local Oz terminal lifecycle hooks --- app/src/ai/agent_sdk/driver.rs | 17 +- app/src/ai/agent_sdk/hooks/mod.rs | 43 ++- app/src/ai/agent_sdk/hooks/mod_tests.rs | 60 +++- app/src/ai/agent_sdk/hooks/payload.rs | 4 +- app/src/ai/blocklist/action_model.rs | 63 +++- app/src/ai/blocklist/action_model/execute.rs | 8 + app/src/ai/blocklist/controller.rs | 216 +++++++------ .../blocklist/controller/response_stream.rs | 16 +- app/src/ai/blocklist/controller_tests.rs | 299 +++++++++++++++++- 9 files changed, 613 insertions(+), 113 deletions(-) diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index e910870d0d9..131378855ae 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -1122,16 +1122,16 @@ impl AgentDriver { permission_mode: "supervised".into(), }; let runtime: Arc = Arc::new(OzHookRuntimeService::new(config)); - let session = OzHookSession { - runtime: Arc::clone(&runtime), - protocol_context: warp_multi_agent_api::OzHookContext { + let session = OzHookSession::new( + Arc::clone(&runtime), + warp_multi_agent_api::OzHookContext { enabled_events, supported_payload_schema_versions: vec![PAYLOAD_SCHEMA_VERSION.into()], }, - payload_context: payload_context.clone(), - redactor: HookRedactor::new(secrets.values().flat_map(secret_values)), - is_driver_owned: true, - }; + payload_context.clone(), + HookRedactor::new(secrets.values().flat_map(secret_values)), + true, + ); foreground .spawn(move |me, ctx| { me.terminal_driver.update(ctx, |driver, ctx| { @@ -1227,6 +1227,9 @@ impl AgentDriver { let Ok(Some(session)) = Self::current_oz_hook_session(foreground).await else { return; }; + if !session.claim_stop() { + return; + } let turn_status = match status { SDKConversationOutputStatus::Success => TurnStatus::Completed, SDKConversationOutputStatus::Error { .. } => TurnStatus::Failed, diff --git a/app/src/ai/agent_sdk/hooks/mod.rs b/app/src/ai/agent_sdk/hooks/mod.rs index 404689e6860..660dde467a1 100644 --- a/app/src/ai/agent_sdk/hooks/mod.rs +++ b/app/src/ai/agent_sdk/hooks/mod.rs @@ -1,5 +1,5 @@ use std::fmt; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use serde::{Deserialize, Serialize}; @@ -30,6 +30,47 @@ pub(crate) struct OzHookSession { pub(crate) payload_context: payload::HookPayloadContext, pub(crate) redactor: redaction::HookRedactor, pub(crate) is_driver_owned: bool, + turn_state: Arc>, +} + +#[derive(Default)] +struct OzHookTurnState { + generation: u64, + stopped_generation: Option, +} + +impl OzHookSession { + pub(crate) fn new( + runtime: Arc, + protocol_context: warp_multi_agent_api::OzHookContext, + payload_context: payload::HookPayloadContext, + redactor: redaction::HookRedactor, + is_driver_owned: bool, + ) -> Self { + Self { + runtime, + protocol_context, + payload_context, + redactor, + is_driver_owned, + turn_state: Default::default(), + } + } + + pub(crate) fn begin_turn(&self) { + let mut state = self.turn_state.lock().unwrap(); + state.generation = state.generation.wrapping_add(1); + state.stopped_generation = None; + } + + pub(crate) fn claim_stop(&self) -> bool { + let mut state = self.turn_state.lock().unwrap(); + if state.stopped_generation == Some(state.generation) { + return false; + } + state.stopped_generation = Some(state.generation); + true + } } #[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] diff --git a/app/src/ai/agent_sdk/hooks/mod_tests.rs b/app/src/ai/agent_sdk/hooks/mod_tests.rs index 7d3cbc53109..d6e7af6c41d 100644 --- a/app/src/ai/agent_sdk/hooks/mod_tests.rs +++ b/app/src/ai/agent_sdk/hooks/mod_tests.rs @@ -1,4 +1,48 @@ -use super::HookEventName; +use std::sync::Arc; + +use async_trait::async_trait; + +use super::payload::HookPayloadContext; +use super::redaction::HookRedactor; +use super::runtime::{ + OzHookCancellationScope, OzHookEvent, OzHookObservation, OzHookRuntime, OzPreToolUseDecision, + OzPreToolUseEvent, +}; +use super::{HookEventName, OzHookSession}; + +struct NoopRuntime; + +#[async_trait] +impl OzHookRuntime for NoopRuntime { + async fn observe(&self, _: OzHookEvent) -> OzHookObservation { + OzHookObservation::default() + } + + async fn pre_tool_use(&self, _: OzPreToolUseEvent) -> OzPreToolUseDecision { + OzPreToolUseDecision::Continue { + diagnostics: Vec::new(), + } + } + + fn cancel(&self, _: OzHookCancellationScope) {} +} + +fn session() -> OzHookSession { + OzHookSession::new( + Arc::new(NoopRuntime), + warp_multi_agent_api::OzHookContext::default(), + HookPayloadContext { + session_id: "session".into(), + run_id: "run".into(), + conversation_id: "conversation".into(), + cwd: "/tmp".into(), + model: "model".into(), + permission_mode: "supervised".into(), + }, + HookRedactor::new([]), + false, + ) +} #[test] fn oz_hooks_config_event_names_are_stable() { @@ -15,3 +59,17 @@ fn oz_hooks_config_event_names_are_stable() { ] ); } + +#[test] +fn oz_hook_session_claims_stop_once_per_turn_across_clones() { + let session = session(); + let clone = session.clone(); + + assert!(session.claim_stop()); + assert!(!clone.claim_stop()); + + clone.begin_turn(); + + assert!(session.claim_stop()); + assert!(!clone.claim_stop()); +} diff --git a/app/src/ai/agent_sdk/hooks/payload.rs b/app/src/ai/agent_sdk/hooks/payload.rs index fc12f79bcfa..f64592ea926 100644 --- a/app/src/ai/agent_sdk/hooks/payload.rs +++ b/app/src/ai/agent_sdk/hooks/payload.rs @@ -118,7 +118,7 @@ impl HookEventFields { } } -#[derive(Clone, Copy, Debug, Deserialize, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(rename_all = "lowercase")] pub(crate) enum SessionStartSource { Startup, @@ -170,7 +170,7 @@ impl CompactTrigger { } } -#[derive(Clone, Copy, Debug, Deserialize, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] pub(crate) enum TurnStatus { Idle, diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 742e92fad9e..78da54ea500 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -65,6 +65,10 @@ use crate::ai::agent::{ }; #[cfg(not(target_family = "wasm"))] use crate::ai::agent_sdk::hooks::OzHookSession; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::payload::{HookEventFields, HookPayloadTemplate, TurnStatus}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::runtime::OzHookEvent; use crate::ai::blocklist::action_model::execute::suggest_new_conversation::SuggestNewConversationExecutor; use crate::ai::blocklist::telemetry::send_run_agents_completed_telemetry; use crate::ai::document::ai_document_model::AIDocumentModel; @@ -890,20 +894,71 @@ impl BlocklistAIActionModel { ctx.emit(BlocklistAIActionEvent::ActionBlockedOnUserConfirmation( action.id.clone(), )); + let blocked_action = format!("{:?}", action.action.user_friendly_name()); + #[cfg(not(target_family = "wasm"))] + if let Some(session) = self + .executor + .as_ref(ctx) + .oz_hook_session(conversation_id) + .filter(OzHookSession::claim_stop) + { + let payload_context = session.payload_context.clone(); + let terminal_view_id = self.terminal_view_id; + ctx.spawn( + async move { + session + .runtime + .observe(OzHookEvent { + invocation_id: uuid::Uuid::new_v4().to_string(), + tool_use_id: None, + payload: HookPayloadTemplate { + context: payload_context, + event: HookEventFields::Stop { + turn_status: TurnStatus::Blocked, + }, + }, + }) + .await; + }, + move |_, _, ctx| { + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.update_conversation_status( + terminal_view_id, + conversation_id, + ConversationStatus::Blocked { blocked_action }, + ctx, + ); + }); + }, + ); + return; + } BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { - let blocked_action_user_friendly_str = action.action.user_friendly_name(); history_model.update_conversation_status( self.terminal_view_id, conversation_id, - ConversationStatus::Blocked { - blocked_action: format!("{blocked_action_user_friendly_str:?}"), - }, + ConversationStatus::Blocked { blocked_action }, ctx, ); }); } } + #[cfg(test)] + pub fn block_action_for_test( + &self, + action: &AIAgentAction, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + self.handle_not_executed_action( + action, + NotExecutedReason::NeedsConfirmation, + conversation_id, + ctx, + ); + } + fn action_phase_for_action( &self, action: &AIAgentAction, diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 62cc52dd091..0fad246b4e5 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -443,6 +443,14 @@ impl BlocklistAIActionExecutor { } } + #[cfg(not(target_family = "wasm"))] + pub(crate) fn oz_hook_session( + &self, + conversation_id: AIConversationId, + ) -> Option { + self.oz_hook_sessions.get(&conversation_id).cloned() + } + pub fn async_executing_action(&self, action_id: &AIAgentActionId) -> Option<&AIAgentAction> { self.async_executing_actions .get(action_id) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index f7fc39c9a19..afe4f76b655 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -466,6 +466,25 @@ struct InputQuery { queued_query_id: Option, } +#[cfg(not(target_family = "wasm"))] +fn local_oz_session_start_source<'a>( + inputs: impl Iterator, + has_existing_exchanges: bool, +) -> Option { + let mut has_user_query = false; + for input in inputs { + match input { + AIAgentInput::ResumeConversation { .. } => return Some(SessionStartSource::Resume), + AIAgentInput::UserQuery { .. } => has_user_query = true, + _ => {} + } + } + has_user_query.then_some(if has_existing_exchanges { + SessionStartSource::Resume + } else { + SessionStartSource::Startup + }) +} impl InputQuery { fn query(&self) -> String { match &self.input_query { @@ -532,6 +551,20 @@ impl BlocklistAIController { if !response_stream.as_ref(ctx).is_completion_deferred() { return false; } + let turn_status = match &finished_event.reason { + Some(warp_multi_agent_api::response_event::stream_finished::Reason::Done(_)) | None => { + TurnStatus::Completed + } + Some( + warp_multi_agent_api::response_event::stream_finished::Reason::Other(_) + | warp_multi_agent_api::response_event::stream_finished::Reason::ContextWindowExceeded(_) + | warp_multi_agent_api::response_event::stream_finished::Reason::QuotaLimit(_) + | warp_multi_agent_api::response_event::stream_finished::Reason::LlmUnavailable(_) + | warp_multi_agent_api::response_event::stream_finished::Reason::InvalidApiKey(_) + | warp_multi_agent_api::response_event::stream_finished::Reason::InternalError(_) + | warp_multi_agent_api::response_event::stream_finished::Reason::MaxTokenLimit(_), + ) => TurnStatus::Failed, + }; let has_actions = self.oz_hook_action_streams.contains(&stream_id) || BlocklistAIHistoryModel::as_ref(ctx) .conversation(&conversation_id) @@ -548,7 +581,13 @@ impl BlocklistAIController { }); return false; }; - if has_actions { + if has_actions && matches!(turn_status, TurnStatus::Completed) { + response_stream.update(ctx, |stream, ctx| { + stream.finish_deferred_completion(ctx); + }); + return false; + } + if !session.claim_stop() { response_stream.update(ctx, |stream, ctx| { stream.finish_deferred_completion(ctx); }); @@ -560,9 +599,7 @@ impl BlocklistAIController { tool_use_id: None, payload: HookPayloadTemplate { context: session.payload_context.clone(), - event: HookEventFields::Stop { - turn_status: TurnStatus::Completed, - }, + event: HookEventFields::Stop { turn_status }, }, }; ctx.spawn( @@ -602,7 +639,7 @@ impl BlocklistAIController { }); return false; }; - if session.is_driver_owned { + if !session.claim_stop() { response_stream.update(ctx, |stream, ctx| { stream.finish_deferred_completion(ctx); }); @@ -625,8 +662,8 @@ impl BlocklistAIController { }, move |me, _, ctx| { me.handle_response_stream_error( - error, &stream_id, + error, conversation_id, &response_stream, ctx, @@ -638,68 +675,12 @@ impl BlocklistAIController { ); true } - - fn handle_response_stream_error( - &mut self, - error: Arc, - stream_id: &ResponseStreamId, - conversation_id: AIConversationId, - response_stream: &ModelHandle, - ctx: &mut ModelContext, - ) { - if matches!(error.as_ref(), AIApiError::QuotaLimit { .. }) { - // If the error is a quota limit, refresh workspace metadata so AI overages are - // immediately up to date. - TeamUpdateManager::handle(ctx).update(ctx, |team_update_manager, ctx| { - std::mem::drop(team_update_manager.refresh_workspace_metadata(ctx)); - }); - AIRequestUsageModel::handle(ctx).update(ctx, |model, ctx| { - model.enable_buy_credits_banner(ctx); - }); - } - // A resume scheduled for this failure keeps the conversation in the non-terminal - // TransientError status instead of Error. - - let recovery_pending = response_stream - .as_ref(ctx) - .should_resume_conversation_after_stream_finished(); - let mut renderable_error: RenderableAIError = (&error).into(); - if let RenderableAIError::Other { - will_attempt_resume, - waiting_for_network, - .. - } - | RenderableAIError::TransientNetworkError { - will_attempt_resume, - waiting_for_network, - .. - } = &mut renderable_error - { - // Rendering-only hints; state machine consumers key off the TransientError - // conversation status instead. - *will_attempt_resume |= recovery_pending; - if recovery_pending { - let network_status = NetworkStatus::as_ref(ctx); - *waiting_for_network = !network_status.is_online(); - } - } - - BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { - history_model.mark_response_stream_completed_with_error( - renderable_error, - recovery_pending, - stream_id, - conversation_id, - self.terminal_surface_id, - ctx, - ); - }); - } #[cfg(not(target_family = "wasm"))] fn initialize_local_oz_hook_session( &mut self, conversation_id: AIConversationId, model: String, + source: SessionStartSource, ctx: &mut ModelContext, ) -> Option { if !FeatureFlag::OzLifecycleHooks.is_enabled() @@ -741,24 +722,16 @@ impl BlocklistAIController { model, permission_mode: "supervised".into(), }; - let source = if BlocklistAIHistoryModel::as_ref(ctx) - .conversation(&conversation_id) - .is_some_and(|conversation| conversation.exchange_count() > 0) - { - SessionStartSource::Resume - } else { - SessionStartSource::Startup - }; - let session = OzHookSession { - runtime: Arc::clone(&runtime), - protocol_context: warp_multi_agent_api::OzHookContext { + let session = OzHookSession::new( + Arc::clone(&runtime), + warp_multi_agent_api::OzHookContext { enabled_events, supported_payload_schema_versions: vec![PAYLOAD_SCHEMA_VERSION.into()], }, - payload_context: payload_context.clone(), - redactor: HookRedactor::new([]), - is_driver_owned: false, - }; + payload_context.clone(), + HookRedactor::new([]), + false, + ); self.set_oz_hook_session(conversation_id, Some(session), ctx); Some(OzHookEvent { invocation_id: uuid::Uuid::new_v4().to_string(), @@ -3069,17 +3042,29 @@ impl BlocklistAIController { request_params.parent_agent_id = parent_agent_id; request_params.agent_name = agent_name; #[cfg(not(target_family = "wasm"))] - let startup_hook = request_input - .all_inputs() - .any(starts_local_oz_hook_session) - .then(|| { - self.initialize_local_oz_hook_session( - conversation_id, - request_params.model.to_string(), - ctx, - ) - }) - .flatten(); + let has_existing_exchanges = history_model + .as_ref(ctx) + .conversation(&conversation_id) + .is_some_and(|conversation| conversation.exchange_count() > 0); + #[cfg(not(target_family = "wasm"))] + let session_start_source = + local_oz_session_start_source(request_input.all_inputs(), has_existing_exchanges); + #[cfg(not(target_family = "wasm"))] + let session_start_hook = session_start_source.and_then(|source| { + self.initialize_local_oz_hook_session( + conversation_id, + request_params.model.to_string(), + source, + ctx, + ) + }); + #[cfg(not(target_family = "wasm"))] + if session_start_hook.is_none() + && request_input.all_inputs().any(AIAgentInput::is_user_query) + && let Some(session) = self.oz_hook_sessions.get(&conversation_id) + { + session.begin_turn(); + } #[cfg(not(target_family = "wasm"))] if let Some(session) = self.oz_hook_sessions.get(&conversation_id) { request_params.oz_hook_context = Some(session.protocol_context.clone()); @@ -3089,7 +3074,7 @@ impl BlocklistAIController { .oz_hook_sessions .get(&conversation_id) .and_then(|session| { - let mut events = startup_hook.into_iter().collect::>(); + let mut events = session_start_hook.clone().into_iter().collect::>(); if let Some(query) = request_input.all_inputs().find_map(|input| { let AIAgentInput::UserQuery { query, .. } = input else { return None; @@ -3485,6 +3470,55 @@ impl BlocklistAIController { self.in_flight_response_streams .try_cancel_stream(response_stream_id, reason, ctx) } + fn handle_response_stream_error( + &mut self, + stream_id: &ResponseStreamId, + error: Arc, + conversation_id: AIConversationId, + response_stream: &ModelHandle, + ctx: &mut ModelContext, + ) { + if matches!(error.as_ref(), AIApiError::QuotaLimit { .. }) { + TeamUpdateManager::handle(ctx).update(ctx, |team_update_manager, ctx| { + std::mem::drop(team_update_manager.refresh_workspace_metadata(ctx)); + }); + AIRequestUsageModel::handle(ctx).update(ctx, |model, ctx| { + model.enable_buy_credits_banner(ctx); + }); + } + + let recovery_pending = response_stream + .as_ref(ctx) + .should_resume_conversation_after_stream_finished(); + let mut renderable_error: RenderableAIError = (&error).into(); + if let RenderableAIError::Other { + will_attempt_resume, + waiting_for_network, + .. + } + | RenderableAIError::TransientNetworkError { + will_attempt_resume, + waiting_for_network, + .. + } = &mut renderable_error + { + *will_attempt_resume |= recovery_pending; + if recovery_pending { + *waiting_for_network = !NetworkStatus::as_ref(ctx).is_online(); + } + } + + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { + history_model.mark_response_stream_completed_with_error( + renderable_error, + recovery_pending, + stream_id, + conversation_id, + self.terminal_surface_id, + ctx, + ); + }); + } fn handle_response_stream_event( &mut self, @@ -3680,8 +3714,8 @@ impl BlocklistAIController { return; } self.handle_response_stream_error( - e, &stream_id, + e, conversation_id, response_stream, ctx, diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index bdfee639c69..087216a9a31 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -314,9 +314,15 @@ impl ResponseStream { event: warp_multi_agent_api::ResponseEvent, ctx: &mut ModelContext, ) { - ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(Ok( - event, - )))); + let request_id = self + .current_request_id + .expect("test response stream has a current request"); + self.handle_response_stream_event(request_id, Ok(event), ctx); + } + + #[cfg(test)] + pub fn set_oz_hook_context_for_test(&mut self, context: warp_multi_agent_api::OzHookContext) { + self.params.oz_hook_context = Some(context); } /// Emits the natural-completion `AfterStreamFinished` event (no cancellation) through @@ -900,13 +906,13 @@ impl ResponseStream { } warp_multi_agent_api::response_event::Type::Finished(finished_event) => { self.stream_finished_received = true; + #[cfg(not(target_family = "wasm"))] + self.defer_completion_for_oz_stop(); // Emit retry success telemetry on successful completion if matches!( finished_event.reason, Some(warp_multi_agent_api::response_event::stream_finished::Reason::Done(_)) | None ) { - #[cfg(not(target_family = "wasm"))] - self.defer_completion_for_oz_stop(); // Emit retry success telemetry if this was a successful completion after retries if self.retries_sent > 0 && let Some(original_error) = &self.original_error { diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 9ca6ae109f2..9980fb0ea54 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -6,19 +6,36 @@ use ai::api_keys::{ ApiKeyManager, AwsCredentials, AwsCredentialsState, CustomEndpointParams, GeapCredentials, GeapCredentialsState, }; +#[cfg(not(target_family = "wasm"))] +use async_trait::async_trait; use chrono::Local; +#[cfg(not(target_family = "wasm"))] +use futures::channel::oneshot; use uuid::Uuid; use warp_core::features::FeatureFlag; use warp_multi_agent_api::response_event; use warpui::{App, SingletonEntity, ViewHandle}; +#[cfg(not(target_family = "wasm"))] +use super::local_oz_session_start_source; use super::response_stream::{PendingResume, RecoveryBudget}; use crate::ai::agent::conversation::AIConversationId; use crate::ai::agent::task::TaskId; use crate::ai::agent::{ - AIAgentAttachment, AIAgentContext, AIAgentInput, CancellationReason, ImageContext, - PassiveSuggestionTrigger, UserQueryMode, + AIAgentAction, AIAgentActionId, AIAgentActionType, AIAgentAttachment, AIAgentContext, + AIAgentInput, CancellationReason, ImageContext, PassiveSuggestionTrigger, UserQueryMode, +}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::payload::{HookEventFields, HookPayloadContext, TurnStatus}; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::redaction::HookRedactor; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::runtime::{ + OzHookCancellationScope, OzHookEvent, OzHookObservation, OzHookRuntime, OzPreToolUseDecision, + OzPreToolUseEvent, }; +#[cfg(not(target_family = "wasm"))] +use crate::ai::agent_sdk::hooks::{OzHookSession, PAYLOAD_SCHEMA_VERSION}; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::orchestration_events::{ OrchestrationEventService, PendingEvent, PendingEventDetail, @@ -49,6 +66,81 @@ const GEAP_TEST_SA_EMAIL: &str = "warp-geap@test-project.iam.gserviceaccount.com fn new_ambient_agent_task_id() -> AmbientAgentTaskId { Uuid::new_v4().to_string().parse().unwrap() } +#[cfg(not(target_family = "wasm"))] +struct PausingStopRuntime { + statuses: async_channel::Sender, + release: Mutex>>, +} + +#[cfg(not(target_family = "wasm"))] +#[async_trait] +impl OzHookRuntime for PausingStopRuntime { + async fn observe(&self, event: OzHookEvent) -> OzHookObservation { + if let HookEventFields::Stop { turn_status } = event.payload.event { + self.statuses.send(turn_status).await.unwrap(); + let release = self.release.lock().unwrap().take(); + if let Some(release) = release { + let _ = release.await; + } + } + OzHookObservation::default() + } + + async fn pre_tool_use(&self, _: OzPreToolUseEvent) -> OzPreToolUseDecision { + OzPreToolUseDecision::Continue { + diagnostics: Vec::new(), + } + } + + fn cancel(&self, _: OzHookCancellationScope) {} +} + +#[cfg(not(target_family = "wasm"))] +fn pausing_oz_session( + conversation_id: AIConversationId, +) -> ( + OzHookSession, + async_channel::Receiver, + oneshot::Sender<()>, +) { + let (statuses_tx, statuses_rx) = async_channel::unbounded(); + let (release_tx, release_rx) = oneshot::channel(); + let runtime: Arc = Arc::new(PausingStopRuntime { + statuses: statuses_tx, + release: Mutex::new(Some(release_rx)), + }); + let session = OzHookSession::new( + runtime, + warp_multi_agent_api::OzHookContext { + enabled_events: vec![warp_multi_agent_api::OzHookEvent::Stop.into()], + supported_payload_schema_versions: vec![PAYLOAD_SCHEMA_VERSION.into()], + }, + HookPayloadContext { + session_id: "session".into(), + run_id: "run".into(), + conversation_id: conversation_id.to_string(), + cwd: "/tmp".into(), + model: "model".into(), + permission_mode: "supervised".into(), + }, + HookRedactor::new([]), + false, + ); + (session, statuses_rx, release_tx) +} + +#[cfg(not(target_family = "wasm"))] +#[test] +fn restored_resume_request_starts_an_oz_resume_session() { + let inputs = [AIAgentInput::ResumeConversation { + context: Arc::from([]), + }]; + + assert_eq!( + local_oz_session_start_source(inputs.iter(), false), + Some(crate::ai::agent_sdk::hooks::payload::SessionStartSource::Resume) + ); +} fn image_attachment(file_name: &str) -> PendingAttachment { PendingAttachment::Image(ImageContext { @@ -373,6 +465,209 @@ fn mock_response_stream_updates_history_through_controller() { }); } +#[cfg(not(target_family = "wasm"))] +#[test] +fn failed_response_waits_for_stop_before_publishing_terminal_status() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + let (conversation_id, stream, stop_statuses, release_stop) = + terminal.update(&mut app, |view, ctx| { + let terminal_surface_id = view.id(); + let stream_id = ResponseStreamId::new_for_test(); + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = history.start_new_conversation( + terminal_surface_id, + false, + false, + false, + ctx, + ); + let task_id = history + .conversation(&conversation_id) + .unwrap() + .get_root_task_id() + .clone(); + history + .update_conversation_for_new_request_input( + RequestInput { + conversation_id, + input_messages: HashMap::from([(task_id, vec![])]), + working_directory: None, + model_id: LLMId::from("test-model"), + coding_model_id: LLMId::from("test-coding-model"), + cli_agent_model_id: LLMId::from("test-cli-agent-model"), + computer_use_model_id: LLMId::from("test-computer-use-model"), + shared_session_response_initiator: None, + request_start_ts: Local::now(), + supported_tools_override: None, + }, + stream_id.clone(), + terminal_surface_id, + ctx, + ) + .unwrap(); + conversation_id + }); + let (session, stop_statuses, release_stop) = pausing_oz_session(conversation_id); + let hook_context = session.protocol_context.clone(); + let stream = ctx.add_model(|_| { + let mut stream = ResponseStream::new_for_test(stream_id.clone()); + stream.set_oz_hook_context_for_test(hook_context); + stream + }); + view.ai_controller().update(ctx, |controller, ctx| { + controller.set_oz_hook_session(conversation_id, Some(session), ctx); + controller.register_mock_stream_for_test( + stream_id, + conversation_id, + stream.clone(), + ctx, + ); + }); + (conversation_id, stream, stop_statuses, release_stop) + }); + let (terminal_status_tx, terminal_status_rx) = async_channel::unbounded(); + app.update(|ctx| { + ctx.subscribe_to_model(&BlocklistAIHistoryModel::handle(ctx), move |_, event, _| { + if let BlocklistAIHistoryEvent::UpdatedConversationStatus { + conversation_id: updated_id, + new_status, + .. + } = event + && *updated_id == conversation_id + && new_status.is_error() + { + terminal_status_tx.try_send(new_status.clone()).unwrap(); + } + }); + }); + + stream.update(&mut app, |stream, ctx| { + stream.emit_response_event_for_test( + warp_multi_agent_api::ResponseEvent { + r#type: Some(response_event::Type::Finished( + response_event::StreamFinished { + reason: Some(response_event::stream_finished::Reason::Other( + response_event::stream_finished::Other {}, + )), + conversation_usage_metadata: None, + token_usage: vec![], + should_refresh_model_config: false, + #[allow(deprecated)] + request_cost: None, + request_charges: None, + }, + )), + }, + ctx, + ); + }); + + assert_eq!(stop_statuses.recv().await.unwrap(), TurnStatus::Failed); + BlocklistAIHistoryModel::handle(&app).read(&app, |history, _| { + assert!( + history + .conversation(&conversation_id) + .unwrap() + .status() + .is_in_progress() + ); + }); + + release_stop.send(()).unwrap(); + assert!(terminal_status_rx.recv().await.unwrap().is_error()); + assert!(stop_statuses.try_recv().is_err()); + }); +} + +#[cfg(not(target_family = "wasm"))] +#[test] +fn blocked_action_waits_for_stop_before_publishing_terminal_status() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + let (conversation_id, action_model, stop_statuses, release_stop) = + terminal.update(&mut app, |view, ctx| { + let terminal_surface_id = view.id(); + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = history.start_new_conversation( + terminal_surface_id, + false, + false, + false, + ctx, + ); + history.update_conversation_status( + terminal_surface_id, + conversation_id, + crate::ai::agent::conversation::ConversationStatus::InProgress, + ctx, + ); + conversation_id + }); + let (session, stop_statuses, release_stop) = pausing_oz_session(conversation_id); + let action_model = view.ai_controller().as_ref(ctx).action_model.clone(); + view.ai_controller().update(ctx, |controller, ctx| { + controller.set_oz_hook_session(conversation_id, Some(session), ctx); + }); + (conversation_id, action_model, stop_statuses, release_stop) + }); + let (terminal_status_tx, terminal_status_rx) = async_channel::unbounded(); + app.update(|ctx| { + ctx.subscribe_to_model(&BlocklistAIHistoryModel::handle(ctx), move |_, event, _| { + if let BlocklistAIHistoryEvent::UpdatedConversationStatus { + conversation_id: updated_id, + new_status, + .. + } = event + && *updated_id == conversation_id + && new_status.is_blocked() + { + terminal_status_tx.try_send(new_status.clone()).unwrap(); + } + }); + }); + let action = AIAgentAction { + id: AIAgentActionId::from("blocked-action".to_owned()), + task_id: TaskId::new("task".to_owned()), + action: AIAgentActionType::RequestCommandOutput { + command: "echo blocked".into(), + is_read_only: None, + is_risky: None, + rationale: None, + uses_pager: None, + wait_until_completion: true, + citations: Vec::new(), + }, + requires_result: true, + }; + + action_model.update(&mut app, |model, ctx| { + model.block_action_for_test(&action, conversation_id, ctx); + }); + + assert_eq!(stop_statuses.recv().await.unwrap(), TurnStatus::Blocked); + BlocklistAIHistoryModel::handle(&app).read(&app, |history, _| { + assert!( + history + .conversation(&conversation_id) + .unwrap() + .status() + .is_in_progress() + ); + }); + + release_stop.send(()).unwrap(); + assert!(terminal_status_rx.recv().await.unwrap().is_blocked()); + assert!(stop_statuses.try_recv().is_err()); + }); +} + /// When an agent command exits the shell, the conversation must be finalized as /// `Error` (not `Cancelled`), and a subsequent `ManuallyCancelled` (as fired by /// the pane-close path) must not overwrite that failure. From 1298e992a73156d85cdf39f95c2922277e21e084 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:37:24 +0000 Subject: [PATCH 11/11] Remove obsolete Oz session start helper --- app/src/ai/blocklist/controller.rs | 7 ------- app/src/ai/blocklist/controller_tests.rs | 9 --------- 2 files changed, 16 deletions(-) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index afe4f76b655..4a416768e8f 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -429,13 +429,6 @@ enum WhichTask { task_id: TaskId, }, } -#[cfg(not(target_family = "wasm"))] -fn starts_local_oz_hook_session(input: &AIAgentInput) -> bool { - matches!( - input, - AIAgentInput::UserQuery { .. } | AIAgentInput::ResumeConversation { .. } - ) -} #[derive(Debug, Clone, PartialEq, Eq)] enum LocalClaudeWakeTrigger { diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 9980fb0ea54..00df224bb7c 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -159,15 +159,6 @@ fn file_attachment(file_name: &str) -> PendingAttachment { }) } -#[cfg(not(target_family = "wasm"))] -#[test] -fn restored_conversation_resume_starts_a_local_oz_hook_session() { - assert!(super::starts_local_oz_hook_session( - &AIAgentInput::ResumeConversation { - context: vec![].into() - } - )); -} #[test] fn passive_suggestions_request_params_omit_ambient_agent_task_id() { App::test((), |mut app| async move {