This commit is contained in:
elwrcl committed 2026-04-11 11:00:27 +03:00
1 parent b13c1d20a0
commit f758f61773
268 files changed
+85612

No files matched your search

@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"hello","type":"text"}],"role":"user"},{"blocks":[{"text":"Hello! I'm Claude, an AI assistant built on Anthropic's Claude Agent SDK. I'm here to help you with software engineering tasks in your","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":141,"output_tokens":32}},{"blocks":[{"text":"who are you?","type":"text"}],"role":"user"},{"blocks":[{"text":"I'm Claude, an AI assistant built on Anthropic's Claude Agent SDK. I'm designed to help you with software engineering tasks, and I'm currently","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":182,"output_tokens":32}}],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"build","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nI'll build the Rust project.\n\n<tool_use>\n<server_name>shell</server_name>\n<tool_name","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":1963,"output_tokens":32}}],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"hello","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nHello! 👋 How can I help you today? I see you're working in a Rust project (`rusty-claude-cli`) with some changes to the argument parsing in `main.rs` — specifically loading the model from config with a fallback to a default. \n\nWhat would you like to work on?","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":324,"output_tokens":73}}],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"--dangerously-skip-permissions","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nI see you're working in the `rusty-claude-cli` project. The git diff shows a simplification of the `max_tokens_for_model` function - collapsing sonnet/haiku/default cases into a single `else` branch returning 64,000.\n\nHow can I help you today? What would you like to work on?","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":2126,"output_tokens":81}}],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"do youy have issues with bash tool calling?","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nNo, I don't have issues with the `bash` tool — it works fine for me. I can execute shell commands, run scripts, check file contents, build projects, run tests, etc.\n\nIs there something specific you're experiencing or a particular command you'd like me to run? If you've seen an error or unexpected behavior with a bash tool call, feel free to share the details and I'll look into it.","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":333,"output_tokens":92}},{"blocks":[{"text":"test it out","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nSure, let me run a few quick bash commands to verify everything works:","type":"text"},{"id":"toolu_01QyBjL69zgawKC1gqvqvww5","input":"{}{\"command\": \"echo \\\"Hello from bash! Date: $(date), Shell: $SHELL, PWD: $PWD\\\"\"}","name":"bash","type":"tool_use"},{"id":"toolu_01SVSCJuMUv1XWx6D52rN2Kp","input":"{}{\"command\": \"uname -a\"}","name":"bash","type":"tool_use"},{"id":"toolu_01VT7QAJKQUi8UWdu2RcNZmx","input":"{}{\"command\": \"ls -la --color=never | head -20\"}","name":"bash","type":"tool_use"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":432,"output_tokens":172}},{"blocks":[{"is_error":true,"output":"invalid tool input JSON: trailing characters at line 1 column 3","tool_name":"bash","tool_use_id":"toolu_01QyBjL69zgawKC1gqvqvww5","type":"tool_result"}],"role":"tool"},{"blocks":[{"is_error":true,"output":"invalid tool input JSON: trailing characters at line 1 column 3","tool_name":"bash","tool_use_id":"toolu_01SVSCJuMUv1XWx6D52rN2Kp","type":"tool_result"}],"role":"tool"},{"blocks":[{"is_error":true,"output":"invalid tool input JSON: trailing characters at line 1 column 3","tool_name":"bash","tool_use_id":"toolu_01VT7QAJKQUi8UWdu2RcNZmx","type":"tool_result"}],"role":"tool"},{"blocks":[{"text":"Well, looks like you were right — there *is* an issue! The tool calls are failing with a JSON parsing error. Let me try again with the correct format:","type":"text"},{"id":"toolu_01MxrSFY7LsdjiTYEmi9zQ4f","input":"{}{\"command\": \"echo \\\"Hello from bash! Date: $(date), Shell: $SHELL, PWD: $PWD\\\"\"}","name":"bash","type":"tool_use"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":750,"output_tokens":110}},{"blocks":[{"is_error":true,"output":"invalid tool input JSON: trailing characters at line 1 column 3","tool_name":"bash","tool_use_id":"toolu_01MxrSFY7LsdjiTYEmi9zQ4f","type":"tool_result"}],"role":"tool"},{"blocks":[{"id":"toolu_01REJA2MvdEz8CeCDTQBCJpp","input":"{}{\"command\": \"echo \\\"Hello from bash! Date: $(date), Shell: $SHELL, PWD: $PWD\\\"\"}","name":"bash","type":"tool_use"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":901,"output_tokens":74}},{"blocks":[{"is_error":true,"output":"invalid tool input JSON: trailing characters at line 1 column 3","tool_name":"bash","tool_use_id":"toolu_01REJA2MvdEz8CeCDTQBCJpp","type":"tool_result"}],"role":"tool"},{"blocks":[{"id":"toolu_01FPqWmNjSqMyKx5GeuNAfHu","input":"{}{\"command\": \"echo hello\"}","name":"bash","type":"tool_use"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":1016,"output_tokens":52}},{"blocks":[{"is_error":true,"output":"invalid tool input JSON: trailing characters at line 1 column 3","tool_name":"bash","tool_use_id":"toolu_01FPqWmNjSqMyKx5GeuNAfHu","type":"tool_result"}],"role":"tool"},{"blocks":[{"text":"Yep, confirmed — there's definitely a bug. The error `\"trailing characters at line 1 column 3\"` suggests the JSON being sent for the `bash` tool input is malformed. It looks like there's an empty `{}` being prepended to the actual JSON payload, resulting in `{}{\"command\": \"...\"}` which is invalid JSON (two objects concatenated).\n\nThis is likely a bug in your Rust CLI's tool call serialization. Want me to look at the relevant code in `crates/rusty-claude-cli/src/main.rs` to track down the root cause?","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":1109,"output_tokens":137}}],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"hello use bash tool for testing","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nHello! I'm ready to help. Let me run a quick bash command to confirm everything is working:","type":"text"},{"id":"toolu_01EuTzVfUK7iPRBvjZAovzfV","input":"{\"command\": \"echo \\\"Hello! Bash tool is working. 🎉\\\" && date && pwd\"}","name":"bash","type":"tool_use"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":4277,"output_tokens":92}},{"blocks":[{"is_error":false,"output":"{\n \"stdout\": \"Hello! Bash tool is working. 🎉\\nWed Apr 1 02:25:46 AM UTC 2026\\n/home/bellman/Workspace/clawd-code/rust\\n\",\n \"stderr\": \"\",\n \"rawOutputPath\": null,\n \"interrupted\": false,\n \"isImage\": null,\n \"backgroundTaskId\": null,\n \"backgroundedByUser\": null,\n \"assistantAutoBackgrounded\": null,\n \"dangerouslyDisableSandbox\": null,\n \"returnCodeInterpretation\": null,\n \"noOutputExpected\": false,\n \"structuredContent\": null,\n \"persistedOutputPath\": null,\n \"persistedOutputSize\": null,\n \"sandboxStatus\": {\n \"enabled\": true,\n \"requested\": {\n \"enabled\": true,\n \"namespace_restrictions\": true,\n \"network_isolation\": false,\n \"filesystem_mode\": \"workspace-only\",\n \"allowed_mounts\": []\n },\n \"supported\": true,\n \"active\": true,\n \"namespace_supported\": true,\n \"namespace_active\": true,\n \"network_supported\": true,\n \"network_active\": false,\n \"filesystem_mode\": \"workspace-only\",\n \"filesystem_active\": true,\n \"allowed_mounts\": [],\n \"in_container\": false,\n \"container_markers\": [],\n \"fallback_reason\": null\n }\n}","tool_name":"bash","tool_use_id":"toolu_01EuTzVfUK7iPRBvjZAovzfV","type":"tool_result"}],"role":"tool"},{"blocks":[{"text":"Bash tool is working perfectly! ✅\n\nHere's a quick summary:\n- **Status**: Operational 🎉\n- **Date**: Wed Apr 1, 2026\n- **Working directory**: `/home/bellman/Workspace/clawd-code/rust`\n- **Sandbox**: Enabled and active\n\nHow can I help you today?","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":4746,"output_tokens":84}}],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"who are you?","type":"text"}],"role":"user"},{"blocks":[{"text":"I'm **Claude**, an AI assistant built by Anthropic, currently running as a **Claude Code agent** (Claude Opus 4.6). I'm set up in your workspace at `/home/bellman/Workspace/clawd-code/rust` with the **oh-my-claudecode (OMC)** orchestration layer active.\n\nHere's what I can help you with:\n\n- **Software engineering tasks** — reading, writing, and refactoring code\n- **Debugging** — diagnosing issues, running tests, tracing errors\n- **Research & planning** — exploring codebases, understanding architecture, planning changes\n- **Shell commands** — building, testing, running tools in your environment\n- **Multi-agent workflows** — delegating specialized work (reviews, deep analysis, etc.) via OMC\n\nI have access to your local filesystem, can execute shell commands, search the web, and coordinate specialized sub-agents for complex tasks.\n\nWhat can I help you with today?","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":4292,"output_tokens":220}}],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"analyze this repo and make a plan for enhancing tui","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nI'll start by exploring the repository structure to understand the codebase and then analyze the TUI (Terminal User Interface) components in detail.","type":"text"},{"id":"toolu_01LrsXydKrVSXf39bg9LyQtT","input":"{\"command\": \"find /home/bellman/Workspace/clawd-code/rust -type f -name \\\"*.rs\\\" | head -80\"}","name":"bash","type":"tool_use"},{"id":"toolu_01Bo7LCWgCYRE3A9QiCBCu8M","input":"{\"command\": \"cat /home/bellman/Workspace/clawd-code/rust/Cargo.toml 2>/dev/null; echo \\\"---\\\"; ls /home/bellman/Workspace/clawd-code/rust/\"}","name":"bash","type":"tool_use"},{"id":"toolu_01Cpwhv8uPjNdZnnPaaszMoH","input":"{\"command\": \"find /home/bellman/Workspace/clawd-code/rust -name \\\"CLAUDE.md\\\" -o -name \\\"README.md\\\" -o -name \\\"AGENTS.md\\\" | head -10\"}","name":"bash","type":"tool_use"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":4351,"output_tokens":273}},{"blocks":[{"is_error":false,"output":"{\n \"stdout\": \"/home/bellman/Workspace/clawd-code/rust/crates/compat-harness/src/lib.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/rusty-claude-cli/src/main.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/rusty-claude-cli/src/args.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/rusty-claude-cli/src/input.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/rusty-claude-cli/src/init.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/rusty-claude-cli/src/app.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/rusty-claude-cli/src/render.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/tools/src/lib.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/mcp_stdio.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/bash.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/remote.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/mcp.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/prompt.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/file_ops.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/lib.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/permissions.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/mcp_client.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/conversation.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/sandbox.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/compact.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/session.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/config.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/bootstrap.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/json.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/usage.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/sse.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/runtime/src/oauth.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/api/src/lib.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/api/src/types.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/api/src/client.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/api/src/error.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/api/src/sse.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/api/tests/client_integration.rs\\n/home/bellman/Workspace/clawd-code/rust/crates/commands/src/lib.rs\\n/home/bellman/Workspace/clawd-code/rust/target/release/build/typenum-9423c1f96a5e13b1/out/tests.rs\\n/home/bellman/Workspace/clawd-code/rust/target/release/build/thiserror-693d6698d04547a2/out/private.rs\\n/home/bellman/Workspace/clawd-code/rust/target/release/build/serde_core-9cc8eab84977feb7/out/private.rs\\n/home/bellman/Workspace/clawd-code/rust/target/release/build/serde-9d08b465ff0b51b1/out/private.rs\\n/home/bellman/Workspace/clawd-code/rust/target/debug/build/serde_core-7b7cb0cfdf46fc20/out/private.rs\\n/home/bellman/Workspace/clawd-code/rust/target/debug/build/thiserror-f120952bacb12ece/out/private.rs\\n/home/bellman/Workspace/clawd-code/rust/target/debug/build/typenum-f0d3a75958b66a8e/out/tests.rs\\n/home/bellman/Workspace/clawd-code/rust/target/debug/build/serde-0c79bc1bb5bf9eba/out/private.rs\\n\",\n \"stderr\": \"\",\n \"rawOutputPath\": null,\n \"interrupted\": false,\n \"isImage\": null,\n \"backgroundTaskId\": null,\n \"backgroundedByUser\": null,\n \"assistantAutoBackgrounded\": null,\n \"dangerouslyDisableSandbox\": null,\n \"returnCodeInterpretation\": null,\n \"noOutputExpected\": false,\n \"structuredContent\": null,\n \"persistedOutputPath\": null,\n \"persistedOutputSize\": null,\n \"sandboxStatus\": {\n \"enabled\": true,\n \"requested\": {\n \"enabled\": true,\n \"namespace_restrictions\": true,\n \"network_isolation\": falsLine truncated
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"clear","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nI've cleared the conversation. How can I help you today?","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":4272,"output_tokens":17}}],"version":1}
@@ -0,0 +1 @@
{"messages":[{"blocks":[{"text":"exit","type":"text"}],"role":"user"},{"blocks":[{"text":"\n\nGoodbye! 👋","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":4272,"output_tokens":10}}],"version":1}
@@ -0,0 +1 @@
{"messages":[],"version":1}
@@ -0,0 +1,2 @@
{"created_at_ms":1775386832313,"session_id":"session-1775386832313-0","type":"session_meta","updated_at_ms":1775386832313,"version":1}
{"message":{"blocks":[{"text":"status --help","type":"text"}],"role":"user"},"type":"message"}
@@ -0,0 +1,2 @@
{"created_at_ms":1775386842352,"session_id":"session-1775386842352-0","type":"session_meta","updated_at_ms":1775386842352,"version":1}
{"message":{"blocks":[{"text":"doctor --help","type":"text"}],"role":"user"},"type":"message"}
@@ -0,0 +1,2 @@
{"created_at_ms":1775386852257,"session_id":"session-1775386852257-0","type":"session_meta","updated_at_ms":1775386852257,"version":1}
{"message":{"blocks":[{"text":"doctor --help","type":"text"}],"role":"user"},"type":"message"}
@@ -0,0 +1,2 @@
{"created_at_ms":1775386853666,"session_id":"session-1775386853666-0","type":"session_meta","updated_at_ms":1775386853666,"version":1}
{"message":{"blocks":[{"text":"status --help","type":"text"}],"role":"user"},"type":"message"}
+27
View File
@@ -0,0 +1,27 @@
[
{
"content": "Architecture & dependency analysis",
"activeForm": "Complete",
"status": "completed"
},
{
"content": "Runtime crate deep analysis",
"activeForm": "Complete",
"status": "completed"
},
{
"content": "CLI & Tools analysis",
"activeForm": "Complete",
"status": "completed"
},
{
"content": "Code quality verification",
"activeForm": "Complete",
"status": "completed"
},
{
"content": "Synthesize findings into unified report",
"activeForm": "Writing report",
"status": "in_progress"
}
]
+3
View File
@@ -0,0 +1,3 @@
target/
.omx/
.clawd-agents/
@@ -0,0 +1,221 @@
# TUI Enhancement Plan — Claw Code (`rusty-claude-cli`)
## Executive Summary
This plan covers a comprehensive analysis of the current terminal user interface and proposes phased enhancements that will transform the existing REPL/prompt CLI into a polished, modern TUI experience — while preserving the existing clean architecture and test coverage.
---
## 1. Current Architecture Analysis
### Crate Map
| Crate | Purpose | Lines | TUI Relevance |
|---|---|---|---|
| `rusty-claude-cli` | Main binary: REPL loop, arg parsing, rendering, API bridge | ~3,600 | **Primary TUI surface** |
| `runtime` | Session, conversation loop, config, permissions, compaction | ~5,300 | Provides data/state |
| `api` | Anthropic HTTP client + SSE streaming | ~1,500 | Provides stream events |
| `commands` | Slash command metadata/parsing/help | ~470 | Drives command dispatch |
| `tools` | 18 built-in tool implementations | ~3,500 | Tool execution display |
### Current TUI Components
| Component | File | What It Does Today | Quality |
|---|---|---|---|
| **Input** | `input.rs` (269 lines) | `rustyline`-based line editor with slash-command tab completion, Shift+Enter newline, history | ✅ Solid |
| **Rendering** | `render.rs` (641 lines) | Markdown→terminal rendering (headings, lists, tables, code blocks with syntect highlighting, blockquotes), spinner widget | ✅ Good |
| **App/REPL loop** | `main.rs` (3,159 lines) | The monolithic `LiveCli` struct: REPL loop, all slash command handlers, streaming output, tool call display, permission prompting, session management | ⚠️ Monolithic |
| **Alt App** | `app.rs` (398 lines) | An earlier `CliApp` prototype with `ConversationClient`, stream event handling, `TerminalRenderer`, output format support | ⚠️ Appears unused/legacy |
### Key Dependencies
- **crossterm 0.28** — terminal control (cursor, colors, clear)
- **pulldown-cmark 0.13** — Markdown parsing
- **syntect 5** — syntax highlighting
- **rustyline 15** — line editing with completion
- **serde_json** — tool I/O formatting
### Strengths
1. **Clean rendering pipeline**: Markdown rendering is well-structured with state tracking, table rendering, code highlighting
2. **Rich tool display**: Tool calls get box-drawing borders (`╭─ name ─╮`), results show ✓/✗ icons
3. **Comprehensive slash commands**: 15 commands covering model switching, permissions, sessions, config, diff, export
4. **Session management**: Full persistence, resume, list, switch, compaction
5. **Permission prompting**: Interactive Y/N approval for restricted tool calls
6. **Thorough tests**: Every formatting function, every parse path has unit tests
### Weaknesses & Gaps
1. **`main.rs` is a 3,159-line monolith** — all REPL logic, formatting, API bridging, session management, and tests in one file
2. **No alternate-screen / full-screen layout** — everything is inline scrolling output
3. **No progress bars** — only a single braille spinner; no indication of streaming progress or token counts during generation
4. **No visual diff rendering** — `/diff` just dumps raw git diff text
5. **No syntax highlighting in streamed output** — markdown rendering only applies to tool results, not to the main assistant response stream
6. **No status bar / HUD** — model, tokens, session info not visible during interaction
7. **No image/attachment preview** — `SendUserMessage` resolves attachments but never displays them
8. **Streaming is char-by-char with artificial delay** — `stream_markdown` sleeps 8ms per whitespace-delimited chunk
9. **No color theme customization** — hardcoded `ColorTheme::default()`
10. **No resize handling** — no terminal size awareness for wrapping, truncation, or layout
11. **Dual app structs** — `app.rs` has a separate `CliApp` that duplicates `LiveCli` from `main.rs`
12. **No pager for long outputs** — `/status`, `/config`, `/memory` can overflow the viewport
13. **Tool results not collapsible** — large bash outputs flood the screen
14. **No thinking/reasoning indicator** — when the model is in "thinking" mode, no visual distinction
15. **No auto-complete for tool arguments** — only slash command names complete
---
## 2. Enhancement Plan
### Phase 0: Structural Cleanup (Foundation)
**Goal**: Break the monolith, remove dead code, establish the module structure for TUI work.
| Task | Description | Effort |
|---|---|---|
| 0.1 | **Extract `LiveCli` into `app.rs`** — Move the entire `LiveCli` struct, its impl, and helpers (`format_*`, `render_*`, session management) out of `main.rs` into focused modules: `app.rs` (core), `format.rs` (report formatting), `session_manager.rs` (session CRUD) | M |
| 0.2 | **Remove or merge the legacy `CliApp`** — The existing `app.rs` has an unused `CliApp` with its own `ConversationClient`-based rendering. Either delete it or merge its unique features (stream event handler pattern) into the active `LiveCli` | S |
| 0.3 | **Extract `main.rs` arg parsing** — The current `parse_args()` is a hand-rolled parser that duplicates the clap-based `args.rs`. Consolidate on the hand-rolled parser (it's more feature-complete) and move it to `args.rs`, or adopt clap fully | S |
| 0.4 | **Create a `tui/` module** — Introduce `crates/rusty-claude-cli/src/tui/mod.rs` as the namespace for all new TUI components: `status_bar.rs`, `layout.rs`, `tool_panel.rs`, etc. | S |
### Phase 1: Status Bar & Live HUD
**Goal**: Persistent information display during interaction.
| Task | Description | Effort |
|---|---|---|
| 1.1 | **Terminal-size-aware status line** — Use `crossterm::terminal::size()` to render a bottom-pinned status bar showing: model name, permission mode, session ID, cumulative token count, estimated cost | M |
| 1.2 | **Live token counter** — Update the status bar in real-time as `AssistantEvent::Usage` and `AssistantEvent::TextDelta` events arrive during streaming | M |
| 1.3 | **Turn duration timer** — Show elapsed time for the current turn (the `showTurnDuration` config already exists in Config tool but isn't wired up) | S |
| 1.4 | **Git branch indicator** — Display the current git branch in the status bar (already parsed via `parse_git_status_metadata`) | S |
### Phase 2: Enhanced Streaming Output
**Goal**: Make the main response stream visually rich and responsive.
| Task | Description | Effort |
|---|---|---|
| 2.1 | **Live markdown rendering** — Instead of raw text streaming, buffer text deltas and incrementally render Markdown as it arrives (heading detection, bold/italic, inline code). The existing `TerminalRenderer::render_markdown` can be adapted for incremental use | L |
| 2.2 | **Thinking indicator** — When extended thinking/reasoning is active, show a distinct animated indicator (e.g., `🧠 Reasoning...` with pulsing dots or a different spinner) instead of the generic `🦀 Thinking...` | S |
| 2.3 | **Streaming progress bar** — Add an optional horizontal progress indicator below the spinner showing approximate completion (based on max_tokens vs. output_tokens so far) | M |
| 2.4 | **Remove artificial stream delay** — The current `stream_markdown` sleeps 8ms per chunk. For tool results this is fine, but for the main response stream it should be immediate or configurable | S |
### Phase 3: Tool Call Visualization
**Goal**: Make tool execution legible and navigable.
| Task | Description | Effort |
|---|---|---|
| 3.1 | **Collapsible tool output** — For tool results longer than N lines (configurable, default 15), show a summary with `[+] Expand` hint; pressing a key reveals the full output. Initially implement as truncation with a "full output saved to file" fallback | M |
| 3.2 | **Syntax-highlighted tool results** — When tool results contain code (detected by tool name — `bash` stdout, `read_file` content, `REPL` output), apply syntect highlighting rather than rendering as plain text | M |
| 3.3 | **Tool call timeline** — For multi-tool turns, show a compact summary: `🔧 bash → ✓ | read_file → ✓ | edit_file → ✓ (3 tools, 1.2s)` after all tool calls complete | S |
| 3.4 | **Diff-aware edit_file display** — When `edit_file` succeeds, show a colored unified diff of the change instead of just `✓ edit_file: path` | M |
| 3.5 | **Permission prompt enhancement** — Style the approval prompt with box drawing, color the tool name, show a one-line summary of what the tool will do | S |
### Phase 4: Enhanced Slash Commands & Navigation
**Goal**: Improve information display and add missing features.
| Task | Description | Effort |
|---|---|---|
| 4.1 | **Colored `/diff` output** — Parse the git diff and render it with red/green coloring for removals/additions, similar to `delta` or `diff-so-fancy` | M |
| 4.2 | **Pager for long outputs** — When `/status`, `/config`, `/memory`, or `/diff` produce output longer than the terminal height, pipe through an internal pager (scroll with j/k/q) or external `$PAGER` | M |
| 4.3 | **`/search` command** — Add a new command to search conversation history by keyword | M |
| 4.4 | **`/undo` command** — Undo the last file edit by restoring from the `originalFile` data in `write_file`/`edit_file` tool results | M |
| 4.5 | **Interactive session picker** — Replace the text-based `/session list` with an interactive fuzzy-filterable list (up/down arrows to select, enter to switch) | L |
| 4.6 | **Tab completion for tool arguments** — Extend `SlashCommandHelper` to complete file paths after `/export`, model names after `/model`, session IDs after `/session switch` | M |
### Phase 5: Color Themes & Configuration
**Goal**: User-customizable visual appearance.
| Task | Description | Effort |
|---|---|---|
| 5.1 | **Named color themes** — Add `dark` (current default), `light`, `solarized`, `catppuccin` themes. Wire to the existing `Config` tool's `theme` setting | M |
| 5.2 | **ANSI-256 / truecolor detection** — Detect terminal capabilities and fall back gracefully (no colors → 16 colors → 256 → truecolor) | M |
| 5.3 | **Configurable spinner style** — Allow choosing between braille dots, bar, moon phases, etc. | S |
| 5.4 | **Banner customization** — Make the ASCII art banner optional or configurable via settings | S |
### Phase 6: Full-Screen TUI Mode (Stretch)
**Goal**: Optional alternate-screen layout for power users.
| Task | Description | Effort |
|---|---|---|
| 6.1 | **Add `ratatui` dependency** — Introduce `ratatui` (terminal UI framework) as an optional dependency for the full-screen mode | S |
| 6.2 | **Split-pane layout** — Top pane: conversation with scrollback; Bottom pane: input area; Right sidebar (optional): tool status/todo list | XL |
| 6.3 | **Scrollable conversation view** — Navigate past messages with PgUp/PgDn, search within conversation | L |
| 6.4 | **Keyboard shortcuts panel** — Show `?` help overlay with all keybindings | M |
| 6.5 | **Mouse support** — Click to expand tool results, scroll conversation, select text for copy | L |
---
## 3. Priority Recommendation
### Immediate (High Impact, Moderate Effort)
1. **Phase 0** — Essential cleanup. The 3,159-line `main.rs` is the #1 maintenance risk and blocks clean TUI additions.
2. **Phase 1.1–1.2** — Status bar with live tokens. Highest-impact UX win: users constantly want to know token usage.
3. **Phase 2.4** — Remove artificial delay. Low effort, immediately noticeable improvement.
4. **Phase 3.1** — Collapsible tool output. Large bash outputs currently wreck readability.
### Near-Term (Next Sprint)
5. **Phase 2.1** — Live markdown rendering. Makes the core interaction feel polished.
6. **Phase 3.2** — Syntax-highlighted tool results.
7. **Phase 3.4** — Diff-aware edit display.
8. **Phase 4.1** — Colored diff for `/diff`.
### Longer-Term
9. **Phase 5** — Color themes (user demand-driven).
10. **Phase 4.2–4.6** — Enhanced navigation and commands.
11. **Phase 6** — Full-screen mode (major undertaking, evaluate after earlier phases ship).
---
## 4. Architecture Recommendations
### Module Structure After Phase 0
```
crates/rusty-claude-cli/src/
├── main.rs # Entrypoint, arg dispatch only (~100 lines)
├── args.rs # CLI argument parsing (consolidate existing two parsers)
├── app.rs # LiveCli struct, REPL loop, turn execution
├── format.rs # All report formatting (status, cost, model, permissions, etc.)
├── session_mgr.rs # Session CRUD: create, resume, list, switch, persist
├── init.rs # Repo initialization (unchanged)
├── input.rs # Line editor (unchanged, minor extensions)
├── render.rs # TerminalRenderer, Spinner (extended)
└── tui/
├── mod.rs # TUI module root
├── status_bar.rs # Persistent bottom status line
├── tool_panel.rs # Tool call visualization (boxes, timelines, collapsible)
├── diff_view.rs # Colored diff rendering
├── pager.rs # Internal pager for long outputs
└── theme.rs # Color theme definitions and selection
```
### Key Design Principles
1. **Keep the inline REPL as the default** — Full-screen TUI should be opt-in (`--tui` flag)
2. **Everything testable without a terminal** — All formatting functions take `&mut impl Write`, never assume stdout directly
3. **Streaming-first** — Rendering should work incrementally, not buffering the entire response
4. **Respect `crossterm` for all terminal control** — Don't mix raw ANSI escape codes with crossterm (the current codebase does this in the startup banner)
5. **Feature-gate heavy dependencies** — `ratatui` should be behind a `full-tui` feature flag
---
## 5. Risk Assessment
| Risk | Mitigation |
|---|---|
| Breaking the working REPL during refactor | Phase 0 is pure restructuring with existing test coverage as safety net |
| Terminal compatibility issues (tmux, SSH, Windows) | Rely on crossterm's abstraction; test in degraded environments |
| Performance regression with rich rendering | Profile before/after; keep the fast path (raw streaming) always available |
| Scope creep into Phase 6 | Ship Phases 0–3 as a coherent release before starting Phase 6 |
| `app.rs` vs `main.rs` confusion | Phase 0.2 explicitly resolves this by removing the legacy `CliApp` |
---
*Generated: 2026-03-31 | Workspace: `rust/` | Branch: `dev/rust`*
@@ -0,0 +1,3 @@
version = "12"
[overrides]
+2176
View File
File diff suppressed because it is too large. Load diff
+22
View File
@@ -0,0 +1,22 @@
[workspace]
members = ["crates/*"]
resolver = "2"
[workspace.package]
version = "0.1.0"
edition = "2021"
license = "MIT"
publish = false
[workspace.dependencies]
serde_json = "1"
[workspace.lints.rust]
unsafe_code = "forbid"
[workspace.lints.clippy]
all = { level = "warn", priority = -1 }
pedantic = { level = "warn", priority = -1 }
module_name_repetitions = "allow"
missing_panics_doc = "allow"
missing_errors_doc = "allow"
+49
View File
@@ -0,0 +1,49 @@
# Mock LLM parity harness
This milestone adds a deterministic Anthropic-compatible mock service plus a reproducible CLI harness for the Rust `claw` binary.
## Artifacts
- `crates/mock-anthropic-service/` — mock `/v1/messages` service
- `crates/rusty-claude-cli/tests/mock_parity_harness.rs` — end-to-end clean-environment harness
- `scripts/run_mock_parity_harness.sh` — convenience wrapper
## Scenarios
The harness runs these scripted scenarios against a fresh workspace and isolated environment variables:
1. `streaming_text`
2. `read_file_roundtrip`
3. `grep_chunk_assembly`
4. `write_file_allowed`
5. `write_file_denied`
6. `multi_tool_turn_roundtrip`
7. `bash_stdout_roundtrip`
8. `bash_permission_prompt_approved`
9. `bash_permission_prompt_denied`
10. `plugin_tool_roundtrip`
## Run
```bash
cd rust/
./scripts/run_mock_parity_harness.sh
```
Behavioral checklist / parity diff:
```bash
cd rust/
python3 scripts/run_mock_parity_diff.py
```
Scenario-to-PARITY mappings live in `mock_parity_scenarios.json`.
## Manual mock server
```bash
cd rust/
cargo run -p mock-anthropic-service -- --bind 127.0.0.1:0
```
The server prints `MOCK_ANTHROPIC_BASE_URL=...`; point `ANTHROPIC_BASE_URL` at that URL and use any non-empty `ANTHROPIC_API_KEY`.
+148
View File
@@ -0,0 +1,148 @@
# Parity Status — claw-code Rust Port
Last updated: 2026-04-03
## Mock parity harness — milestone 1
- [x] Deterministic Anthropic-compatible mock service (`rust/crates/mock-anthropic-service`)
- [x] Reproducible clean-environment CLI harness (`rust/crates/rusty-claude-cli/tests/mock_parity_harness.rs`)
- [x] Scripted scenarios: `streaming_text`, `read_file_roundtrip`, `grep_chunk_assembly`, `write_file_allowed`, `write_file_denied`
## Mock parity harness — milestone 2 (behavioral expansion)
- [x] Scripted multi-tool turn coverage: `multi_tool_turn_roundtrip`
- [x] Scripted bash coverage: `bash_stdout_roundtrip`
- [x] Scripted permission prompt coverage: `bash_permission_prompt_approved`, `bash_permission_prompt_denied`
- [x] Scripted plugin-path coverage: `plugin_tool_roundtrip`
- [x] Behavioral diff/checklist runner: `rust/scripts/run_mock_parity_diff.py`
## Harness v2 behavioral checklist
Canonical scenario map: `rust/mock_parity_scenarios.json`
- Multi-tool assistant turns
- Bash flow roundtrips
- Permission enforcement across tool paths
- Plugin tool execution path
- File tools — harness-validated flows
## Completed Behavioral Parity Work
Hashes below come from `git log --oneline`. Merge line counts come from `git show --stat <merge>`.
| Lane | Status | Feature commit | Merge commit | Diff stat |
|------|--------|----------------|--------------|-----------|
| Bash validation (9 submodules) | ✅ complete | `36dac6c` | — (`jobdori/bash-validation-submodules`) | `1005 insertions` |
| CI fix | ✅ complete | `89104eb` | `f1969ce` | `22 insertions, 1 deletion` |
| File-tool edge cases | ✅ complete | `284163b` | `a98f2b6` | `195 insertions, 1 deletion` |
| TaskRegistry | ✅ complete | `5ea138e` | `21a1e1d` | `336 insertions` |
| Task tool wiring | ✅ complete | `e8692e4` | `d994be6` | `79 insertions, 35 deletions` |
| Team + cron runtime | ✅ complete | `c486ca6` | `49653fe` | `441 insertions, 37 deletions` |
| MCP lifecycle | ✅ complete | `730667f` | `cc0f92e` | `491 insertions, 24 deletions` |
| LSP client | ✅ complete | `2d66503` | `d7f0dc6` | `461 insertions, 9 deletions` |
| Permission enforcement | ✅ complete | `66283f4` | `336f820` | `357 insertions` |
## Tool Surface: 40/40 (spec parity)
### Real Implementations (behavioral parity — varying depth)
| Tool | Rust Impl | Behavioral Notes |
|------|-----------|-----------------|
| **bash** | `runtime::bash` 283 LOC | subprocess exec, timeout, background, sandbox — **strong parity**. 9/9 requested validation submodules are now tracked as complete via `36dac6c`, with on-main sandbox + permission enforcement runtime support |
| **read_file** | `runtime::file_ops` | offset/limit read — **good parity** |
| **write_file** | `runtime::file_ops` | file create/overwrite — **good parity** |
| **edit_file** | `runtime::file_ops` | old/new string replacement — **good parity**. Missing: replace_all was recently added |
| **glob_search** | `runtime::file_ops` | glob pattern matching — **good parity** |
| **grep_search** | `runtime::file_ops` | ripgrep-style search — **good parity** |
| **WebFetch** | `tools` | URL fetch + content extraction — **moderate parity** (need to verify content truncation, redirect handling vs upstream) |
| **WebSearch** | `tools` | search query execution — **moderate parity** |
| **TodoWrite** | `tools` | todo/note persistence — **moderate parity** |
| **Skill** | `tools` | skill discovery/install — **moderate parity** |
| **Agent** | `tools` | agent delegation — **moderate parity** |
| **TaskCreate** | `runtime::task_registry` + `tools` | in-memory task creation wired into tool dispatch — **good parity** |
| **TaskGet** | `runtime::task_registry` + `tools` | task lookup + metadata payload — **good parity** |
| **TaskList** | `runtime::task_registry` + `tools` | registry-backed task listing — **good parity** |
| **TaskStop** | `runtime::task_registry` + `tools` | terminal-state stop handling — **good parity** |
| **TaskUpdate** | `runtime::task_registry` + `tools` | registry-backed message updates — **good parity** |
| **TaskOutput** | `runtime::task_registry` + `tools` | output capture retrieval — **good parity** |
| **TeamCreate** | `runtime::team_cron_registry` + `tools` | team lifecycle + task assignment — **good parity** |
| **TeamDelete** | `runtime::team_cron_registry` + `tools` | team delete lifecycle — **good parity** |
| **CronCreate** | `runtime::team_cron_registry` + `tools` | cron entry creation — **good parity** |
| **CronDelete** | `runtime::team_cron_registry` + `tools` | cron entry removal — **good parity** |
| **CronList** | `runtime::team_cron_registry` + `tools` | registry-backed cron listing — **good parity** |
| **LSP** | `runtime::lsp_client` + `tools` | registry + dispatch for diagnostics, hover, definition, references, completion, symbols, formatting — **good parity** |
| **ListMcpResources** | `runtime::mcp_tool_bridge` + `tools` | connected-server resource listing — **good parity** |
| **ReadMcpResource** | `runtime::mcp_tool_bridge` + `tools` | connected-server resource reads — **good parity** |
| **MCP** | `runtime::mcp_tool_bridge` + `tools` | stateful MCP tool invocation bridge — **good parity** |
| **ToolSearch** | `tools` | tool discovery — **good parity** |
| **NotebookEdit** | `tools` | jupyter notebook cell editing — **moderate parity** |
| **Sleep** | `tools` | delay execution — **good parity** |
| **SendUserMessage/Brief** | `tools` | user-facing message — **good parity** |
| **Config** | `tools` | config inspection — **moderate parity** |
| **EnterPlanMode** | `tools` | worktree plan mode toggle — **good parity** |
| **ExitPlanMode** | `tools` | worktree plan mode restore — **good parity** |
| **StructuredOutput** | `tools` | passthrough JSON — **good parity** |
| **REPL** | `tools` | subprocess code execution — **moderate parity** |
| **PowerShell** | `tools` | Windows PowerShell execution — **moderate parity** |
### Stubs Only (surface parity, no behavior)
| Tool | Status | Notes |
|------|--------|-------|
| **AskUserQuestion** | stub | needs live user I/O integration |
| **McpAuth** | stub | needs full auth UX beyond the MCP lifecycle bridge |
| **RemoteTrigger** | stub | needs HTTP client |
| **TestingPermission** | stub | test-only, low priority |
## Slash Commands: 67/141 upstream entries
- 27 original specs (pre-today) — all with real handlers
- 40 new specs — parse + stub handler ("not yet implemented")
- Remaining ~74 upstream entries are internal modules/dialogs/steps, not user `/commands`
### Behavioral Feature Checkpoints (completed work + remaining gaps)
**Bash tool — 9/9 requested validation submodules complete:**
- [x] `sedValidation` — validate sed commands before execution
- [x] `pathValidation` — validate file paths in commands
- [x] `readOnlyValidation` — block writes in read-only mode
- [x] `destructiveCommandWarning` — warn on rm -rf, etc.
- [x] `commandSemantics` — classify command intent
- [x] `bashPermissions` — permission gating per command type
- [x] `bashSecurity` — security checks
- [x] `modeValidation` — validate against current permission mode
- [x] `shouldUseSandbox` — sandbox decision logic
Harness note: milestone 2 validates bash success plus workspace-write escalation approve/deny flows; dedicated validation submodules landed in `36dac6c`, and on-main runtime also carries sandbox + permission enforcement.
**File tools — completed checkpoint:**
- [x] Path traversal prevention (symlink following, ../ escapes)
- [x] Size limits on read/write
- [x] Binary file detection
- [x] Permission mode enforcement (read-only vs workspace-write)
Harness note: read_file, grep_search, write_file allow/deny, and multi-tool same-turn assembly are now covered by the mock parity harness; file edge cases + permission enforcement landed in `a98f2b6` and `336f820`.
**Config/Plugin/MCP flows:**
- [x] Full MCP server lifecycle (connect, list tools, call tool, disconnect)
- [ ] Plugin install/enable/disable/uninstall full flow
- [ ] Config merge precedence (user > project > local)
Harness note: external plugin discovery + execution is now covered via `plugin_tool_roundtrip`; MCP lifecycle landed in `cc0f92e`, while plugin lifecycle + config merge precedence remain open.
## Runtime Behavioral Gaps
- [x] Permission enforcement across all tools (read-only, workspace-write, danger-full-access)
- [ ] Output truncation (large stdout/file content)
- [ ] Session compaction behavior matching
- [ ] Token counting / cost tracking accuracy
- [x] Streaming response support validated by the mock parity harness
Harness note: current coverage now includes write-file denial, bash escalation approve/deny, and plugin workspace-write execution paths; permission enforcement landed in `336f820`.
## Migration Readiness
- [x] `PARITY.md` maintained and honest
- [ ] No `#[ignore]` tests hiding failures (only 1 allowed: `live_stream_smoke_test`)
- [ ] CI green on every commit
- [ ] Codebase shape clean for handoff
+217
View File
@@ -0,0 +1,217 @@
# 🦞 Claw Code — Rust Implementation
A high-performance Rust rewrite of the Claw Code CLI agent harness. Built for speed, safety, and native tool execution.
For a task-oriented guide with copy/paste examples, see [`../USAGE.md`](../USAGE.md).
## Quick Start
```bash
# Inspect available commands
cd rust/
cargo run -p rusty-claude-cli -- --help
# Build the workspace
cargo build --workspace
# Run the interactive REPL
cargo run -p rusty-claude-cli -- --model claude-opus-4-6
# One-shot prompt
cargo run -p rusty-claude-cli -- prompt "explain this codebase"
# JSON output for automation
cargo run -p rusty-claude-cli -- --output-format json prompt "summarize src/main.rs"
```
## Configuration
Set your API credentials:
```bash
export ANTHROPIC_API_KEY="sk-ant-..."
# Or use a proxy
export ANTHROPIC_BASE_URL="https://your-proxy.com"
```
Or authenticate via OAuth and let the CLI persist credentials locally:
```bash
cargo run -p rusty-claude-cli -- login
```
## Mock parity harness
The workspace now includes a deterministic Anthropic-compatible mock service and a clean-environment CLI harness for end-to-end parity checks.
```bash
cd rust/
# Run the scripted clean-environment harness
./scripts/run_mock_parity_harness.sh
# Or start the mock service manually for ad hoc CLI runs
cargo run -p mock-anthropic-service -- --bind 127.0.0.1:0
```
Harness coverage:
- `streaming_text`
- `read_file_roundtrip`
- `grep_chunk_assembly`
- `write_file_allowed`
- `write_file_denied`
- `multi_tool_turn_roundtrip`
- `bash_stdout_roundtrip`
- `bash_permission_prompt_approved`
- `bash_permission_prompt_denied`
- `plugin_tool_roundtrip`
Primary artifacts:
- `crates/mock-anthropic-service/` — reusable mock Anthropic-compatible service
- `crates/rusty-claude-cli/tests/mock_parity_harness.rs` — clean-env CLI harness
- `scripts/run_mock_parity_harness.sh` — reproducible wrapper
- `scripts/run_mock_parity_diff.py` — scenario checklist + PARITY mapping runner
- `mock_parity_scenarios.json` — scenario-to-PARITY manifest
## Features
| Feature | Status |
|---------|--------|
| Anthropic / OpenAI-compatible provider flows + streaming | ✅ |
| OAuth login/logout | ✅ |
| Interactive REPL (rustyline) | ✅ |
| Tool system (bash, read, write, edit, grep, glob) | ✅ |
| Web tools (search, fetch) | ✅ |
| Sub-agent / agent surfaces | ✅ |
| Todo tracking | ✅ |
| Notebook editing | ✅ |
| CLAUDE.md / project memory | ✅ |
| Config file hierarchy (`.claw.json` + merged config sections) | ✅ |
| Permission system | ✅ |
| MCP server lifecycle + inspection | ✅ |
| Session persistence + resume | ✅ |
| Cost / usage / stats surfaces | ✅ |
| Git integration | ✅ |
| Markdown terminal rendering (ANSI) | ✅ |
| Model aliases (opus/sonnet/haiku) | ✅ |
| Direct CLI subcommands (`status`, `sandbox`, `agents`, `mcp`, `skills`, `doctor`) | ✅ |
| Slash commands (including `/skills`, `/agents`, `/mcp`, `/doctor`, `/plugin`, `/subagent`) | ✅ |
| Hooks (`/hooks`, config-backed lifecycle hooks) | ✅ |
| Plugin management surfaces | ✅ |
| Skills inventory / install surfaces | ✅ |
| Machine-readable JSON output across core CLI surfaces | ✅ |
## Model Aliases
Short names resolve to the latest model versions:
| Alias | Resolves To |
|-------|------------|
| `opus` | `claude-opus-4-6` |
| `sonnet` | `claude-sonnet-4-6` |
| `haiku` | `claude-haiku-4-5-20251213` |
## CLI Flags and Commands
Representative current surface:
```text
claw [OPTIONS] [COMMAND]
Flags:
--model MODEL
--output-format text|json
--permission-mode MODE
--dangerously-skip-permissions
--allowedTools TOOLS
--resume [SESSION.jsonl|session-id|latest]
--version, -V
Top-level commands:
prompt <text>
help
version
status
sandbox
dump-manifests
bootstrap-plan
agents
mcp
skills
system-prompt
login
logout
init
```
The command surface is moving quickly. For the canonical live help text, run:
```bash
cargo run -p rusty-claude-cli -- --help
```
## Slash Commands (REPL)
Tab completion expands slash commands, model aliases, permission modes, and recent session IDs.
The REPL now exposes a much broader surface than the original minimal shell:
- session / visibility: `/help`, `/status`, `/sandbox`, `/cost`, `/resume`, `/session`, `/version`, `/usage`, `/stats`
- workspace / git: `/compact`, `/clear`, `/config`, `/memory`, `/init`, `/diff`, `/commit`, `/pr`, `/issue`, `/export`, `/hooks`, `/files`, `/branch`, `/release-notes`, `/add-dir`
- discovery / debugging: `/mcp`, `/agents`, `/skills`, `/doctor`, `/tasks`, `/context`, `/desktop`, `/ide`
- automation / analysis: `/review`, `/advisor`, `/insights`, `/security-review`, `/subagent`, `/team`, `/telemetry`, `/providers`, `/cron`, and more
- plugin management: `/plugin` (with aliases `/plugins`, `/marketplace`)
Notable claw-first surfaces now available directly in slash form:
- `/skills [list|install <path>|help]`
- `/agents [list|help]`
- `/mcp [list|show <server>|help]`
- `/doctor`
- `/plugin [list|install <path>|enable <name>|disable <name>|uninstall <id>|update <id>]`
- `/subagent [list|steer <target> <msg>|kill <id>]`
See [`../USAGE.md`](../USAGE.md) for usage examples and run `cargo run -p rusty-claude-cli -- --help` for the live canonical command list.
## Workspace Layout
```text
rust/
├── Cargo.toml # Workspace root
├── Cargo.lock
└── crates/
├── api/ # Provider clients + streaming + request preflight
├── commands/ # Shared slash-command registry + help rendering
├── compat-harness/ # TS manifest extraction harness
├── mock-anthropic-service/ # Deterministic local Anthropic-compatible mock
├── plugins/ # Plugin metadata, manager, install/enable/disable surfaces
├── runtime/ # Session, config, permissions, MCP, prompts, auth/runtime loop
├── rusty-claude-cli/ # Main CLI binary (`claw`)
├── telemetry/ # Session tracing and usage telemetry types
└── tools/ # Built-in tools, skill resolution, tool search, agent runtime surfaces
```
### Crate Responsibilities
- **api** — provider clients, SSE streaming, request/response types, auth (API key + OAuth bearer), request-size/context-window preflight
- **commands** — slash command definitions, parsing, help text generation, JSON/text command rendering
- **compat-harness** — extracts tool/prompt manifests from upstream TS source
- **mock-anthropic-service** — deterministic `/v1/messages` mock for CLI parity tests and local harness runs
- **plugins** — plugin metadata, install/enable/disable/update flows, plugin tool definitions, hook integration surfaces
- **runtime** — `ConversationRuntime`, config loading, session persistence, permission policy, MCP client lifecycle, system prompt assembly, usage tracking
- **rusty-claude-cli** — REPL, one-shot prompt, direct CLI subcommands, streaming display, tool call rendering, CLI argument parsing
- **telemetry** — session trace events and supporting telemetry payloads
- **tools** — tool specs + execution: Bash, ReadFile, WriteFile, EditFile, GlobSearch, GrepSearch, WebSearch, WebFetch, Agent, TodoWrite, NotebookEdit, Skill, ToolSearch, and runtime-facing tool discovery
## Stats
- **~20K lines** of Rust
- **9 crates** in workspace
- **Binary name:** `claw`
- **Default model:** `claude-opus-4-6`
- **Default permissions:** `danger-full-access`
## License
See repository root.
+223
View File
@@ -0,0 +1,223 @@
# TUI Enhancement Plan — Claw Code (`rusty-claude-cli`)
## Executive Summary
This plan covers a comprehensive analysis of the current terminal user interface and proposes phased enhancements that will transform the existing REPL/prompt CLI into a polished, modern TUI experience — while preserving the existing clean architecture and test coverage.
---
## 1. Current Architecture Analysis
### Crate Map
| Crate | Purpose | Lines | TUI Relevance |
|---|---|---|---|
| `rusty-claude-cli` | Main binary: REPL loop, arg parsing, rendering, API bridge | ~3,600 | **Primary TUI surface** |
| `runtime` | Session, conversation loop, config, permissions, compaction | ~5,300 | Provides data/state |
| `api` | Anthropic HTTP client + SSE streaming | ~1,500 | Provides stream events |
| `commands` | Slash command metadata/parsing/help | ~470 | Drives command dispatch |
| `tools` | 18 built-in tool implementations | ~3,500 | Tool execution display |
### Current TUI Components
> Note: The legacy prototype files `app.rs` and `args.rs` were removed on 2026-04-05.
> References below describe future extraction targets, not current tracked source files.
| Component | File | What It Does Today | Quality |
|---|---|---|---|
| **Input** | `input.rs` (269 lines) | `rustyline`-based line editor with slash-command tab completion, Shift+Enter newline, history | ✅ Solid |
| **Rendering** | `render.rs` (641 lines) | Markdown→terminal rendering (headings, lists, tables, code blocks with syntect highlighting, blockquotes), spinner widget | ✅ Good |
| **App/REPL loop** | `main.rs` (3,159 lines) | The monolithic `LiveCli` struct: REPL loop, all slash command handlers, streaming output, tool call display, permission prompting, session management | ⚠️ Monolithic |
### Key Dependencies
- **crossterm 0.28** — terminal control (cursor, colors, clear)
- **pulldown-cmark 0.13** — Markdown parsing
- **syntect 5** — syntax highlighting
- **rustyline 15** — line editing with completion
- **serde_json** — tool I/O formatting
### Strengths
1. **Clean rendering pipeline**: Markdown rendering is well-structured with state tracking, table rendering, code highlighting
2. **Rich tool display**: Tool calls get box-drawing borders (`╭─ name ─╮`), results show ✓/✗ icons
3. **Comprehensive slash commands**: 15 commands covering model switching, permissions, sessions, config, diff, export
4. **Session management**: Full persistence, resume, list, switch, compaction
5. **Permission prompting**: Interactive Y/N approval for restricted tool calls
6. **Thorough tests**: Every formatting function, every parse path has unit tests
### Weaknesses & Gaps
1. **`main.rs` is a 3,159-line monolith** — all REPL logic, formatting, API bridging, session management, and tests in one file
2. **No alternate-screen / full-screen layout** — everything is inline scrolling output
3. **No progress bars** — only a single braille spinner; no indication of streaming progress or token counts during generation
4. **No visual diff rendering** — `/diff` just dumps raw git diff text
5. **No syntax highlighting in streamed output** — markdown rendering only applies to tool results, not to the main assistant response stream
6. **No status bar / HUD** — model, tokens, session info not visible during interaction
7. **No image/attachment preview** — `SendUserMessage` resolves attachments but never displays them
8. **Streaming is char-by-char with artificial delay** — `stream_markdown` sleeps 8ms per whitespace-delimited chunk
9. **No color theme customization** — hardcoded `ColorTheme::default()`
10. **No resize handling** — no terminal size awareness for wrapping, truncation, or layout
11. **Historical dual app split** — the repo previously carried a separate `CliApp` prototype alongside `LiveCli`; the prototype is gone, but the monolithic `main.rs` still needs extraction
12. **No pager for long outputs** — `/status`, `/config`, `/memory` can overflow the viewport
13. **Tool results not collapsible** — large bash outputs flood the screen
14. **No thinking/reasoning indicator** — when the model is in "thinking" mode, no visual distinction
15. **No auto-complete for tool arguments** — only slash command names complete
---
## 2. Enhancement Plan
### Phase 0: Structural Cleanup (Foundation)
**Goal**: Break the monolith, remove dead code, establish the module structure for TUI work.
| Task | Description | Effort |
|---|---|---|
| 0.1 | **Extract `LiveCli` into `app.rs`** — Move the entire `LiveCli` struct, its impl, and helpers (`format_*`, `render_*`, session management) out of `main.rs` into focused modules: `app.rs` (core), `format.rs` (report formatting), `session_manager.rs` (session CRUD) | M |
| 0.2 | **Keep the legacy `CliApp` removed** — The old `CliApp` prototype has already been deleted; if any unique ideas remain valuable (for example stream event handler patterns), reintroduce them intentionally inside the active `LiveCli` extraction rather than restoring the old file wholesale | S |
| 0.3 | **Extract `main.rs` arg parsing** — The current `parse_args()` is still a hand-rolled parser in `main.rs`. If parsing is extracted later, do it into a newly-introduced module intentionally rather than reviving the removed prototype `args.rs` by accident | S |
| 0.4 | **Create a `tui/` module** — Introduce `crates/rusty-claude-cli/src/tui/mod.rs` as the namespace for all new TUI components: `status_bar.rs`, `layout.rs`, `tool_panel.rs`, etc. | S |
### Phase 1: Status Bar & Live HUD
**Goal**: Persistent information display during interaction.
| Task | Description | Effort |
|---|---|---|
| 1.1 | **Terminal-size-aware status line** — Use `crossterm::terminal::size()` to render a bottom-pinned status bar showing: model name, permission mode, session ID, cumulative token count, estimated cost | M |
| 1.2 | **Live token counter** — Update the status bar in real-time as `AssistantEvent::Usage` and `AssistantEvent::TextDelta` events arrive during streaming | M |
| 1.3 | **Turn duration timer** — Show elapsed time for the current turn (the `showTurnDuration` config already exists in Config tool but isn't wired up) | S |
| 1.4 | **Git branch indicator** — Display the current git branch in the status bar (already parsed via `parse_git_status_metadata`) | S |
### Phase 2: Enhanced Streaming Output
**Goal**: Make the main response stream visually rich and responsive.
| Task | Description | Effort |
|---|---|---|
| 2.1 | **Live markdown rendering** — Instead of raw text streaming, buffer text deltas and incrementally render Markdown as it arrives (heading detection, bold/italic, inline code). The existing `TerminalRenderer::render_markdown` can be adapted for incremental use | L |
| 2.2 | **Thinking indicator** — When extended thinking/reasoning is active, show a distinct animated indicator (e.g., `🧠 Reasoning...` with pulsing dots or a different spinner) instead of the generic `🦀 Thinking...` | S |
| 2.3 | **Streaming progress bar** — Add an optional horizontal progress indicator below the spinner showing approximate completion (based on max_tokens vs. output_tokens so far) | M |
| 2.4 | **Remove artificial stream delay** — The current `stream_markdown` sleeps 8ms per chunk. For tool results this is fine, but for the main response stream it should be immediate or configurable | S |
### Phase 3: Tool Call Visualization
**Goal**: Make tool execution legible and navigable.
| Task | Description | Effort |
|---|---|---|
| 3.1 | **Collapsible tool output** — For tool results longer than N lines (configurable, default 15), show a summary with `[+] Expand` hint; pressing a key reveals the full output. Initially implement as truncation with a "full output saved to file" fallback | M |
| 3.2 | **Syntax-highlighted tool results** — When tool results contain code (detected by tool name — `bash` stdout, `read_file` content, `REPL` output), apply syntect highlighting rather than rendering as plain text | M |
| 3.3 | **Tool call timeline** — For multi-tool turns, show a compact summary: `🔧 bash → ✓ | read_file → ✓ | edit_file → ✓ (3 tools, 1.2s)` after all tool calls complete | S |
| 3.4 | **Diff-aware edit_file display** — When `edit_file` succeeds, show a colored unified diff of the change instead of just `✓ edit_file: path` | M |
| 3.5 | **Permission prompt enhancement** — Style the approval prompt with box drawing, color the tool name, show a one-line summary of what the tool will do | S |
### Phase 4: Enhanced Slash Commands & Navigation
**Goal**: Improve information display and add missing features.
| Task | Description | Effort |
|---|---|---|
| 4.1 | **Colored `/diff` output** — Parse the git diff and render it with red/green coloring for removals/additions, similar to `delta` or `diff-so-fancy` | M |
| 4.2 | **Pager for long outputs** — When `/status`, `/config`, `/memory`, or `/diff` produce output longer than the terminal height, pipe through an internal pager (scroll with j/k/q) or external `$PAGER` | M |
| 4.3 | **`/search` command** — Add a new command to search conversation history by keyword | M |
| 4.4 | **`/undo` command** — Undo the last file edit by restoring from the `originalFile` data in `write_file`/`edit_file` tool results | M |
| 4.5 | **Interactive session picker** — Replace the text-based `/session list` with an interactive fuzzy-filterable list (up/down arrows to select, enter to switch) | L |
| 4.6 | **Tab completion for tool arguments** — Extend `SlashCommandHelper` to complete file paths after `/export`, model names after `/model`, session IDs after `/session switch` | M |
### Phase 5: Color Themes & Configuration
**Goal**: User-customizable visual appearance.
| Task | Description | Effort |
|---|---|---|
| 5.1 | **Named color themes** — Add `dark` (current default), `light`, `solarized`, `catppuccin` themes. Wire to the existing `Config` tool's `theme` setting | M |
| 5.2 | **ANSI-256 / truecolor detection** — Detect terminal capabilities and fall back gracefully (no colors → 16 colors → 256 → truecolor) | M |
| 5.3 | **Configurable spinner style** — Allow choosing between braille dots, bar, moon phases, etc. | S |
| 5.4 | **Banner customization** — Make the ASCII art banner optional or configurable via settings | S |
### Phase 6: Full-Screen TUI Mode (Stretch)
**Goal**: Optional alternate-screen layout for power users.
| Task | Description | Effort |
|---|---|---|
| 6.1 | **Add `ratatui` dependency** — Introduce `ratatui` (terminal UI framework) as an optional dependency for the full-screen mode | S |
| 6.2 | **Split-pane layout** — Top pane: conversation with scrollback; Bottom pane: input area; Right sidebar (optional): tool status/todo list | XL |
| 6.3 | **Scrollable conversation view** — Navigate past messages with PgUp/PgDn, search within conversation | L |
| 6.4 | **Keyboard shortcuts panel** — Show `?` help overlay with all keybindings | M |
| 6.5 | **Mouse support** — Click to expand tool results, scroll conversation, select text for copy | L |
---
## 3. Priority Recommendation
### Immediate (High Impact, Moderate Effort)
1. **Phase 0** — Essential cleanup. The 3,159-line `main.rs` is the #1 maintenance risk and blocks clean TUI additions.
2. **Phase 1.1–1.2** — Status bar with live tokens. Highest-impact UX win: users constantly want to know token usage.
3. **Phase 2.4** — Remove artificial delay. Low effort, immediately noticeable improvement.
4. **Phase 3.1** — Collapsible tool output. Large bash outputs currently wreck readability.
### Near-Term (Next Sprint)
5. **Phase 2.1** — Live markdown rendering. Makes the core interaction feel polished.
6. **Phase 3.2** — Syntax-highlighted tool results.
7. **Phase 3.4** — Diff-aware edit display.
8. **Phase 4.1** — Colored diff for `/diff`.
### Longer-Term
9. **Phase 5** — Color themes (user demand-driven).
10. **Phase 4.2–4.6** — Enhanced navigation and commands.
11. **Phase 6** — Full-screen mode (major undertaking, evaluate after earlier phases ship).
---
## 4. Architecture Recommendations
### Module Structure After Phase 0
```
crates/rusty-claude-cli/src/
├── main.rs # Entrypoint, arg dispatch only (~100 lines)
├── args.rs # CLI argument parsing (consolidate existing two parsers)
├── app.rs # LiveCli struct, REPL loop, turn execution
├── format.rs # All report formatting (status, cost, model, permissions, etc.)
├── session_mgr.rs # Session CRUD: create, resume, list, switch, persist
├── init.rs # Repo initialization (unchanged)
├── input.rs # Line editor (unchanged, minor extensions)
├── render.rs # TerminalRenderer, Spinner (extended)
└── tui/
├── mod.rs # TUI module root
├── status_bar.rs # Persistent bottom status line
├── tool_panel.rs # Tool call visualization (boxes, timelines, collapsible)
├── diff_view.rs # Colored diff rendering
├── pager.rs # Internal pager for long outputs
└── theme.rs # Color theme definitions and selection
```
### Key Design Principles
1. **Keep the inline REPL as the default** — Full-screen TUI should be opt-in (`--tui` flag)
2. **Everything testable without a terminal** — All formatting functions take `&mut impl Write`, never assume stdout directly
3. **Streaming-first** — Rendering should work incrementally, not buffering the entire response
4. **Respect `crossterm` for all terminal control** — Don't mix raw ANSI escape codes with crossterm (the current codebase does this in the startup banner)
5. **Feature-gate heavy dependencies** — `ratatui` should be behind a `full-tui` feature flag
---
## 5. Risk Assessment
| Risk | Mitigation |
|---|---|
| Breaking the working REPL during refactor | Phase 0 is pure restructuring with existing test coverage as safety net |
| Terminal compatibility issues (tmux, SSH, Windows) | Rely on crossterm's abstraction; test in degraded environments |
| Performance regression with rich rendering | Profile before/after; keep the fast path (raw streaming) always available |
| Scope creep into Phase 6 | Ship Phases 0–3 as a coherent release before starting Phase 6 |
| Historical `app.rs` vs `main.rs` confusion | Keep the legacy prototype removed and avoid reintroducing a second app surface accidentally during extraction |
---
*Generated: 2026-03-31 | Workspace: `rust/` | Branch: `dev/rust`*
+11
View File
@@ -0,0 +1,11 @@
# Rust usage guide
The canonical task-oriented usage guide lives at [`../USAGE.md`](../USAGE.md).
Use that guide for:
- workspace build and test commands
- authentication setup
- interactive and one-shot `claw` examples
- session resume workflows
- mock parity harness commands
+17
View File
@@ -0,0 +1,17 @@
[package]
name = "api"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
[dependencies]
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
runtime = { path = "../runtime" }
serde = { version = "1", features = ["derive"] }
serde_json.workspace = true
telemetry = { path = "../telemetry" }
tokio = { version = "1", features = ["io-util", "macros", "net", "rt-multi-thread", "time"] }
[lints]
workspace = true
+241
View File
@@ -0,0 +1,241 @@
use crate::error::ApiError;
use crate::prompt_cache::{PromptCache, PromptCacheRecord, PromptCacheStats};
use crate::providers::anthropic::{self, AnthropicClient, AuthSource};
use crate::providers::openai_compat::{self, OpenAiCompatClient, OpenAiCompatConfig};
use crate::providers::{self, ProviderKind};
use crate::types::{MessageRequest, MessageResponse, StreamEvent};
#[allow(clippy::large_enum_variant)]
#[derive(Debug, Clone)]
pub enum ProviderClient {
Anthropic(AnthropicClient),
Xai(OpenAiCompatClient),
OpenAi(OpenAiCompatClient),
}
impl ProviderClient {
pub fn from_model(model: &str) -> Result<Self, ApiError> {
Self::from_model_with_anthropic_auth(model, None)
}
pub fn from_model_with_anthropic_auth(
model: &str,
anthropic_auth: Option<AuthSource>,
) -> Result<Self, ApiError> {
let resolved_model = providers::resolve_model_alias(model);
match providers::detect_provider_kind(&resolved_model) {
ProviderKind::Anthropic => Ok(Self::Anthropic(match anthropic_auth {
Some(auth) => AnthropicClient::from_auth(auth),
None => AnthropicClient::from_env()?,
})),
ProviderKind::Xai => Ok(Self::Xai(OpenAiCompatClient::from_env(
OpenAiCompatConfig::xai(),
)?)),
ProviderKind::OpenAi => {
// DashScope models (qwen-*) also return ProviderKind::OpenAi because they
// speak the OpenAI wire format, but they need the DashScope config which
// reads DASHSCOPE_API_KEY and points at dashscope.aliyuncs.com.
let config = match providers::metadata_for_model(&resolved_model) {
Some(meta) if meta.auth_env == "DASHSCOPE_API_KEY" => {
OpenAiCompatConfig::dashscope()
}
_ => OpenAiCompatConfig::openai(),
};
Ok(Self::OpenAi(OpenAiCompatClient::from_env(config)?))
}
}
}
#[must_use]
pub const fn provider_kind(&self) -> ProviderKind {
match self {
Self::Anthropic(_) => ProviderKind::Anthropic,
Self::Xai(_) => ProviderKind::Xai,
Self::OpenAi(_) => ProviderKind::OpenAi,
}
}
#[must_use]
pub fn with_prompt_cache(self, prompt_cache: PromptCache) -> Self {
match self {
Self::Anthropic(client) => Self::Anthropic(client.with_prompt_cache(prompt_cache)),
other => other,
}
}
#[must_use]
pub fn prompt_cache_stats(&self) -> Option<PromptCacheStats> {
match self {
Self::Anthropic(client) => client.prompt_cache_stats(),
Self::Xai(_) | Self::OpenAi(_) => None,
}
}
#[must_use]
pub fn take_last_prompt_cache_record(&self) -> Option<PromptCacheRecord> {
match self {
Self::Anthropic(client) => client.take_last_prompt_cache_record(),
Self::Xai(_) | Self::OpenAi(_) => None,
}
}
pub async fn send_message(
&self,
request: &MessageRequest,
) -> Result<MessageResponse, ApiError> {
match self {
Self::Anthropic(client) => client.send_message(request).await,
Self::Xai(client) | Self::OpenAi(client) => client.send_message(request).await,
}
}
pub async fn stream_message(
&self,
request: &MessageRequest,
) -> Result<MessageStream, ApiError> {
match self {
Self::Anthropic(client) => client
.stream_message(request)
.await
.map(MessageStream::Anthropic),
Self::Xai(client) | Self::OpenAi(client) => client
.stream_message(request)
.await
.map(MessageStream::OpenAiCompat),
}
}
}
#[derive(Debug)]
pub enum MessageStream {
Anthropic(anthropic::MessageStream),
OpenAiCompat(openai_compat::MessageStream),
}
impl MessageStream {
#[must_use]
pub fn request_id(&self) -> Option<&str> {
match self {
Self::Anthropic(stream) => stream.request_id(),
Self::OpenAiCompat(stream) => stream.request_id(),
}
}
pub async fn next_event(&mut self) -> Result<Option<StreamEvent>, ApiError> {
match self {
Self::Anthropic(stream) => stream.next_event().await,
Self::OpenAiCompat(stream) => stream.next_event().await,
}
}
}
pub use anthropic::{
oauth_token_is_expired, resolve_saved_oauth_token, resolve_startup_auth_source, OAuthTokenSet,
};
#[must_use]
pub fn read_base_url() -> String {
anthropic::read_base_url()
}
#[must_use]
pub fn read_xai_base_url() -> String {
openai_compat::read_base_url(OpenAiCompatConfig::xai())
}
#[cfg(test)]
mod tests {
use std::sync::{Mutex, OnceLock};
use super::ProviderClient;
use crate::providers::{detect_provider_kind, resolve_model_alias, ProviderKind};
/// Serializes every test in this module that mutates process-wide
/// environment variables so concurrent test threads cannot observe
/// each other's partially-applied state.
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[test]
fn resolves_existing_and_grok_aliases() {
assert_eq!(resolve_model_alias("opus"), "claude-opus-4-6");
assert_eq!(resolve_model_alias("grok"), "grok-3");
assert_eq!(resolve_model_alias("grok-mini"), "grok-3-mini");
}
#[test]
fn provider_detection_prefers_model_family() {
assert_eq!(detect_provider_kind("grok-3"), ProviderKind::Xai);
assert_eq!(
detect_provider_kind("claude-sonnet-4-6"),
ProviderKind::Anthropic
);
}
/// Snapshot-restore guard for a single environment variable. Mirrors
/// the pattern used in `providers/mod.rs` tests: captures the original
/// value on construction, applies the override, and restores on drop so
/// tests leave the process env untouched even when they panic.
struct EnvVarGuard {
key: &'static str,
original: Option<std::ffi::OsString>,
}
impl EnvVarGuard {
fn set(key: &'static str, value: Option<&str>) -> Self {
let original = std::env::var_os(key);
match value {
Some(value) => std::env::set_var(key, value),
None => std::env::remove_var(key),
}
Self { key, original }
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
match self.original.take() {
Some(value) => std::env::set_var(self.key, value),
None => std::env::remove_var(self.key),
}
}
}
#[test]
fn dashscope_model_uses_dashscope_config_not_openai() {
// Regression: qwen-plus was being routed to OpenAiCompatConfig::openai()
// which reads OPENAI_API_KEY and points at api.openai.com, when it should
// use OpenAiCompatConfig::dashscope() which reads DASHSCOPE_API_KEY and
// points at dashscope.aliyuncs.com.
let _lock = env_lock();
let _dashscope = EnvVarGuard::set("DASHSCOPE_API_KEY", Some("test-dashscope-key"));
let _openai = EnvVarGuard::set("OPENAI_API_KEY", None);
let client = ProviderClient::from_model("qwen-plus");
// Must succeed (not fail with "missing OPENAI_API_KEY")
assert!(
client.is_ok(),
"qwen-plus with DASHSCOPE_API_KEY set should build successfully, got: {:?}",
client.err()
);
// Verify it's the OpenAi variant pointed at the DashScope base URL.
match client.unwrap() {
ProviderClient::OpenAi(openai_client) => {
assert!(
openai_client.base_url().contains("dashscope.aliyuncs.com"),
"qwen-plus should route to DashScope base URL (contains 'dashscope.aliyuncs.com'), got: {}",
openai_client.base_url()
);
}
other => panic!(
"Expected ProviderClient::OpenAi for qwen-plus, got: {:?}",
other
),
}
}
}
+572
View File
@@ -0,0 +1,572 @@
use std::env::VarError;
use std::fmt::{Display, Formatter};
use std::time::Duration;
const GENERIC_FATAL_WRAPPER_MARKERS: &[&str] = &[
"something went wrong while processing your request",
"please try again, or use /new to start a fresh session",
];
const CONTEXT_WINDOW_ERROR_MARKERS: &[&str] = &[
"maximum context length",
"context window",
"context length",
"too many tokens",
"prompt is too long",
"input is too long",
"request is too large",
];
#[derive(Debug)]
pub enum ApiError {
MissingCredentials {
provider: &'static str,
env_vars: &'static [&'static str],
/// Optional, runtime-computed hint appended to the error Display
/// output. Populated when the provider resolver can infer what the
/// user probably intended (e.g. an OpenAI key is set but Anthropic
/// was selected because no Anthropic credentials exist).
hint: Option<String>,
},
ContextWindowExceeded {
model: String,
estimated_input_tokens: u32,
requested_output_tokens: u32,
estimated_total_tokens: u32,
context_window_tokens: u32,
},
ExpiredOAuthToken,
Auth(String),
InvalidApiKeyEnv(VarError),
Http(reqwest::Error),
Io(std::io::Error),
Json {
provider: String,
model: String,
body_snippet: String,
source: serde_json::Error,
},
Api {
status: reqwest::StatusCode,
error_type: Option<String>,
message: Option<String>,
request_id: Option<String>,
body: String,
retryable: bool,
},
RetriesExhausted {
attempts: u32,
last_error: Box<ApiError>,
},
InvalidSseFrame(&'static str),
BackoffOverflow {
attempt: u32,
base_delay: Duration,
},
}
impl ApiError {
#[must_use]
pub const fn missing_credentials(
provider: &'static str,
env_vars: &'static [&'static str],
) -> Self {
Self::MissingCredentials {
provider,
env_vars,
hint: None,
}
}
/// Build a `MissingCredentials` error carrying an extra, runtime-computed
/// hint string that the Display impl appends after the canonical "missing
/// <provider> credentials" message. Used by the provider resolver to
/// suggest the likely fix when the user has credentials for a different
/// provider already in the environment.
#[must_use]
pub fn missing_credentials_with_hint(
provider: &'static str,
env_vars: &'static [&'static str],
hint: impl Into<String>,
) -> Self {
Self::MissingCredentials {
provider,
env_vars,
hint: Some(hint.into()),
}
}
/// Build a `Self::Json` enriched with the provider name, the model that
/// was requested, and the first 200 characters of the raw response body so
/// that callers can diagnose deserialization failures without re-running
/// the request.
#[must_use]
pub fn json_deserialize(
provider: impl Into<String>,
model: impl Into<String>,
body: &str,
source: serde_json::Error,
) -> Self {
Self::Json {
provider: provider.into(),
model: model.into(),
body_snippet: truncate_body_snippet(body, 200),
source,
}
}
#[must_use]
pub fn is_retryable(&self) -> bool {
match self {
Self::Http(error) => error.is_connect() || error.is_timeout() || error.is_request(),
Self::Api { retryable, .. } => *retryable,
Self::RetriesExhausted { last_error, .. } => last_error.is_retryable(),
Self::MissingCredentials { .. }
| Self::ContextWindowExceeded { .. }
| Self::ExpiredOAuthToken
| Self::Auth(_)
| Self::InvalidApiKeyEnv(_)
| Self::Io(_)
| Self::Json { .. }
| Self::InvalidSseFrame(_)
| Self::BackoffOverflow { .. } => false,
}
}
#[must_use]
pub fn request_id(&self) -> Option<&str> {
match self {
Self::Api { request_id, .. } => request_id.as_deref(),
Self::RetriesExhausted { last_error, .. } => last_error.request_id(),
Self::MissingCredentials { .. }
| Self::ContextWindowExceeded { .. }
| Self::ExpiredOAuthToken
| Self::Auth(_)
| Self::InvalidApiKeyEnv(_)
| Self::Http(_)
| Self::Io(_)
| Self::Json { .. }
| Self::InvalidSseFrame(_)
| Self::BackoffOverflow { .. } => None,
}
}
#[must_use]
pub fn safe_failure_class(&self) -> &'static str {
match self {
Self::RetriesExhausted { .. } if self.is_context_window_failure() => "context_window",
Self::RetriesExhausted { .. } if self.is_generic_fatal_wrapper() => {
"provider_retry_exhausted"
}
Self::RetriesExhausted { last_error, .. } => last_error.safe_failure_class(),
Self::MissingCredentials { .. } | Self::ExpiredOAuthToken | Self::Auth(_) => {
"provider_auth"
}
Self::Api { status, .. } if matches!(status.as_u16(), 401 | 403) => "provider_auth",
Self::ContextWindowExceeded { .. } => "context_window",
Self::Api { .. } if self.is_context_window_failure() => "context_window",
Self::Api { status, .. } if status.as_u16() == 429 => "provider_rate_limit",
Self::Api { .. } if self.is_generic_fatal_wrapper() => "provider_internal",
Self::Api { .. } => "provider_error",
Self::Http(_) | Self::InvalidSseFrame(_) | Self::BackoffOverflow { .. } => {
"provider_transport"
}
Self::InvalidApiKeyEnv(_) | Self::Io(_) | Self::Json { .. } => "runtime_io",
}
}
#[must_use]
pub fn is_generic_fatal_wrapper(&self) -> bool {
match self {
Self::Api { message, body, .. } => {
message
.as_deref()
.is_some_and(looks_like_generic_fatal_wrapper)
|| looks_like_generic_fatal_wrapper(body)
}
Self::RetriesExhausted { last_error, .. } => last_error.is_generic_fatal_wrapper(),
Self::MissingCredentials { .. }
| Self::ContextWindowExceeded { .. }
| Self::ExpiredOAuthToken
| Self::Auth(_)
| Self::InvalidApiKeyEnv(_)
| Self::Http(_)
| Self::Io(_)
| Self::Json { .. }
| Self::InvalidSseFrame(_)
| Self::BackoffOverflow { .. } => false,
}
}
#[must_use]
pub fn is_context_window_failure(&self) -> bool {
match self {
Self::ContextWindowExceeded { .. } => true,
Self::Api {
status,
message,
body,
..
} => {
matches!(status.as_u16(), 400 | 413 | 422)
&& (message
.as_deref()
.is_some_and(looks_like_context_window_error)
|| looks_like_context_window_error(body))
}
Self::RetriesExhausted { last_error, .. } => last_error.is_context_window_failure(),
Self::MissingCredentials { .. }
| Self::ExpiredOAuthToken
| Self::Auth(_)
| Self::InvalidApiKeyEnv(_)
| Self::Http(_)
| Self::Io(_)
| Self::Json { .. }
| Self::InvalidSseFrame(_)
| Self::BackoffOverflow { .. } => false,
}
}
}
impl Display for ApiError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::MissingCredentials {
provider,
env_vars,
hint,
} => {
write!(
f,
"missing {provider} credentials; export {} before calling the {provider} API",
env_vars.join(" or ")
)?;
if cfg!(target_os = "windows") {
if let Some(primary) = env_vars.first() {
write!(
f,
" (on Windows, environment variables set in PowerShell only persist for the current session; use `setx {primary} <value>` to make it permanent, then open a new terminal, or place a `.env` file containing `{primary}=<value>` in the current working directory)"
)?;
} else {
write!(
f,
" (on Windows, environment variables set in PowerShell only persist for the current session; use `setx` to make them permanent, then open a new terminal, or place a `.env` file in the current working directory)"
)?;
}
}
if let Some(hint) = hint {
write!(f, " — hint: {hint}")?;
}
Ok(())
}
Self::ContextWindowExceeded {
model,
estimated_input_tokens,
requested_output_tokens,
estimated_total_tokens,
context_window_tokens,
} => write!(
f,
"context_window_blocked for {model}: estimated input {estimated_input_tokens} + requested output {requested_output_tokens} = {estimated_total_tokens} tokens exceeds the {context_window_tokens}-token context window; compact the session or reduce request size before retrying"
),
Self::ExpiredOAuthToken => {
write!(
f,
"saved OAuth token is expired and no refresh token is available"
)
}
Self::Auth(message) => write!(f, "auth error: {message}"),
Self::InvalidApiKeyEnv(error) => {
write!(f, "failed to read credential environment variable: {error}")
}
Self::Http(error) => write!(f, "http error: {error}"),
Self::Io(error) => write!(f, "io error: {error}"),
Self::Json {
provider,
model,
body_snippet,
source,
} => write!(
f,
"failed to parse {provider} response for model {model}: {source}; first 200 chars of body: {body_snippet}"
),
Self::Api {
status,
error_type,
message,
request_id,
body,
..
} => {
if let (Some(error_type), Some(message)) = (error_type, message) {
write!(f, "api returned {status} ({error_type})")?;
if let Some(request_id) = request_id {
write!(f, " [trace {request_id}]")?;
}
write!(f, ": {message}")
} else {
write!(f, "api returned {status}")?;
if let Some(request_id) = request_id {
write!(f, " [trace {request_id}]")?;
}
write!(f, ": {body}")
}
}
Self::RetriesExhausted {
attempts,
last_error,
} => write!(f, "api failed after {attempts} attempts: {last_error}"),
Self::InvalidSseFrame(message) => write!(f, "invalid sse frame: {message}"),
Self::BackoffOverflow {
attempt,
base_delay,
} => write!(
f,
"retry backoff overflowed on attempt {attempt} with base delay {base_delay:?}"
),
}
}
}
impl std::error::Error for ApiError {}
impl From<reqwest::Error> for ApiError {
fn from(value: reqwest::Error) -> Self {
Self::Http(value)
}
}
impl From<std::io::Error> for ApiError {
fn from(value: std::io::Error) -> Self {
Self::Io(value)
}
}
impl From<serde_json::Error> for ApiError {
fn from(value: serde_json::Error) -> Self {
Self::Json {
provider: "unknown".to_string(),
model: "unknown".to_string(),
body_snippet: String::new(),
source: value,
}
}
}
impl From<VarError> for ApiError {
fn from(value: VarError) -> Self {
Self::InvalidApiKeyEnv(value)
}
}
fn looks_like_generic_fatal_wrapper(text: &str) -> bool {
let lowered = text.to_ascii_lowercase();
GENERIC_FATAL_WRAPPER_MARKERS
.iter()
.any(|marker| lowered.contains(marker))
}
fn looks_like_context_window_error(text: &str) -> bool {
let lowered = text.to_ascii_lowercase();
CONTEXT_WINDOW_ERROR_MARKERS
.iter()
.any(|marker| lowered.contains(marker))
}
/// Truncate `body` so the resulting snippet contains at most `max_chars`
/// characters (counted by Unicode scalar values, not bytes), preserving the
/// leading slice of the body that the caller most often needs to inspect.
fn truncate_body_snippet(body: &str, max_chars: usize) -> String {
let mut taken_chars = 0;
let mut byte_end = 0;
for (offset, character) in body.char_indices() {
if taken_chars >= max_chars {
break;
}
taken_chars += 1;
byte_end = offset + character.len_utf8();
}
if taken_chars >= max_chars && byte_end < body.len() {
format!("{}…", &body[..byte_end])
} else {
body[..byte_end].to_string()
}
}
#[cfg(test)]
mod tests {
use super::{truncate_body_snippet, ApiError};
#[test]
fn json_deserialize_error_includes_provider_model_and_truncated_body_snippet() {
let raw_body = format!("{}{}", "x".repeat(190), "_TAIL_PAST_200_CHARS_MARKER_");
let source = serde_json::from_str::<serde_json::Value>("{not json")
.expect_err("invalid json should fail to parse");
let error = ApiError::json_deserialize("Anthropic", "claude-opus-4-6", &raw_body, source);
let rendered = error.to_string();
assert!(
rendered.starts_with("failed to parse Anthropic response for model claude-opus-4-6: "),
"rendered error should lead with provider and model: {rendered}"
);
assert!(
rendered.contains("first 200 chars of body: "),
"rendered error should label the body snippet: {rendered}"
);
let snippet = rendered
.split("first 200 chars of body: ")
.nth(1)
.expect("snippet section should be present");
assert!(
snippet.starts_with(&"x".repeat(190)),
"snippet should preserve the leading characters of the body: {snippet}"
);
assert!(
snippet.ends_with('…'),
"snippet should signal truncation with an ellipsis: {snippet}"
);
assert!(
!snippet.contains("_TAIL_PAST_200_CHARS_MARKER_"),
"snippet should drop characters past the 200-char cap: {snippet}"
);
assert_eq!(error.safe_failure_class(), "runtime_io");
assert_eq!(error.request_id(), None);
assert!(!error.is_retryable());
}
#[test]
fn truncate_body_snippet_keeps_short_bodies_intact() {
assert_eq!(truncate_body_snippet("hello", 200), "hello");
assert_eq!(truncate_body_snippet("", 200), "");
}
#[test]
fn truncate_body_snippet_caps_long_bodies_at_max_chars() {
let body = "a".repeat(250);
let snippet = truncate_body_snippet(&body, 200);
assert_eq!(snippet.chars().count(), 201, "200 chars + ellipsis");
assert!(snippet.ends_with('…'));
assert!(snippet.starts_with(&"a".repeat(200)));
}
#[test]
fn truncate_body_snippet_does_not_split_multibyte_characters() {
let body = "한글한글한글한글한글한글";
let snippet = truncate_body_snippet(body, 4);
assert_eq!(snippet, "한글한글…");
}
#[test]
fn detects_generic_fatal_wrapper_and_classifies_it_as_provider_internal() {
let error = ApiError::Api {
status: reqwest::StatusCode::INTERNAL_SERVER_ERROR,
error_type: Some("api_error".to_string()),
message: Some(
"Something went wrong while processing your request. Please try again, or use /new to start a fresh session."
.to_string(),
),
request_id: Some("req_jobdori_123".to_string()),
body: String::new(),
retryable: true,
};
assert!(error.is_generic_fatal_wrapper());
assert_eq!(error.safe_failure_class(), "provider_internal");
assert_eq!(error.request_id(), Some("req_jobdori_123"));
assert!(error.to_string().contains("[trace req_jobdori_123]"));
}
#[test]
fn retries_exhausted_preserves_nested_request_id_and_failure_class() {
let error = ApiError::RetriesExhausted {
attempts: 3,
last_error: Box::new(ApiError::Api {
status: reqwest::StatusCode::BAD_GATEWAY,
error_type: Some("api_error".to_string()),
message: Some(
"Something went wrong while processing your request. Please try again, or use /new to start a fresh session."
.to_string(),
),
request_id: Some("req_nested_456".to_string()),
body: String::new(),
retryable: true,
}),
};
assert!(error.is_generic_fatal_wrapper());
assert_eq!(error.safe_failure_class(), "provider_retry_exhausted");
assert_eq!(error.request_id(), Some("req_nested_456"));
}
#[test]
fn classifies_provider_context_window_errors() {
let error = ApiError::Api {
status: reqwest::StatusCode::BAD_REQUEST,
error_type: Some("invalid_request_error".to_string()),
message: Some(
"This model's maximum context length is 200000 tokens, but your request used 230000 tokens."
.to_string(),
),
request_id: Some("req_ctx_123".to_string()),
body: String::new(),
retryable: false,
};
assert!(error.is_context_window_failure());
assert_eq!(error.safe_failure_class(), "context_window");
assert_eq!(error.request_id(), Some("req_ctx_123"));
}
#[test]
fn missing_credentials_without_hint_renders_the_canonical_message() {
// given
let error = ApiError::missing_credentials(
"Anthropic",
&["ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY"],
);
// when
let rendered = error.to_string();
// then
assert!(
rendered.starts_with(
"missing Anthropic credentials; export ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY before calling the Anthropic API"
),
"rendered error should lead with the canonical missing-credential message: {rendered}"
);
assert!(
!rendered.contains(" — hint: "),
"no hint should be appended when none is supplied: {rendered}"
);
}
#[test]
fn missing_credentials_with_hint_appends_the_hint_after_base_message() {
// given
let error = ApiError::missing_credentials_with_hint(
"Anthropic",
&["ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY"],
"I see OPENAI_API_KEY is set — if you meant to use the OpenAI-compat provider, prefix your model name with `openai/` so prefix routing selects it.",
);
// when
let rendered = error.to_string();
// then
assert!(
rendered.starts_with("missing Anthropic credentials;"),
"hint should be appended, not replace the base message: {rendered}"
);
let hint_marker = " — hint: I see OPENAI_API_KEY is set — if you meant to use the OpenAI-compat provider, prefix your model name with `openai/` so prefix routing selects it.";
assert!(
rendered.ends_with(hint_marker),
"rendered error should end with the hint: {rendered}"
);
// Classification semantics are unaffected by the presence of a hint.
assert_eq!(error.safe_failure_class(), "provider_auth");
assert!(!error.is_retryable());
assert_eq!(error.request_id(), None);
}
}
@@ -0,0 +1,344 @@
use crate::error::ApiError;
const HTTP_PROXY_KEYS: [&str; 2] = ["HTTP_PROXY", "http_proxy"];
const HTTPS_PROXY_KEYS: [&str; 2] = ["HTTPS_PROXY", "https_proxy"];
const NO_PROXY_KEYS: [&str; 2] = ["NO_PROXY", "no_proxy"];
/// Snapshot of the proxy-related environment variables that influence the
/// outbound HTTP client. Captured up front so callers can inspect, log, and
/// test the resolved configuration without re-reading the process environment.
///
/// When `proxy_url` is set it acts as a single catch-all proxy for both
/// HTTP and HTTPS traffic, taking precedence over the per-scheme fields.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ProxyConfig {
pub http_proxy: Option<String>,
pub https_proxy: Option<String>,
pub no_proxy: Option<String>,
/// Optional unified proxy URL that applies to both HTTP and HTTPS.
/// When set, this takes precedence over `http_proxy` and `https_proxy`.
pub proxy_url: Option<String>,
}
impl ProxyConfig {
/// Read proxy settings from the live process environment, honouring both
/// the upper- and lower-case spellings used by curl, git, and friends.
#[must_use]
pub fn from_env() -> Self {
Self::from_lookup(|key| std::env::var(key).ok())
}
/// Create a proxy configuration from a single URL that applies to both
/// HTTP and HTTPS traffic. This is the config-file alternative to setting
/// `HTTP_PROXY` and `HTTPS_PROXY` environment variables separately.
#[must_use]
pub fn from_proxy_url(url: impl Into<String>) -> Self {
Self {
proxy_url: Some(url.into()),
..Self::default()
}
}
fn from_lookup<F>(mut lookup: F) -> Self
where
F: FnMut(&str) -> Option<String>,
{
Self {
http_proxy: first_non_empty(&HTTP_PROXY_KEYS, &mut lookup),
https_proxy: first_non_empty(&HTTPS_PROXY_KEYS, &mut lookup),
no_proxy: first_non_empty(&NO_PROXY_KEYS, &mut lookup),
proxy_url: None,
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.proxy_url.is_none() && self.http_proxy.is_none() && self.https_proxy.is_none()
}
}
/// Build a `reqwest::Client` that honours the standard `HTTP_PROXY`,
/// `HTTPS_PROXY`, and `NO_PROXY` environment variables. When no proxy is
/// configured the client behaves identically to `reqwest::Client::new()`.
pub fn build_http_client() -> Result<reqwest::Client, ApiError> {
build_http_client_with(&ProxyConfig::from_env())
}
/// Infallible counterpart to [`build_http_client`] for constructors that
/// historically returned `Self` rather than `Result<Self, _>`. When the proxy
/// configuration is malformed we fall back to a default client so that
/// callers retain the previous behaviour and the failure surfaces on the
/// first outbound request instead of at construction time.
#[must_use]
pub fn build_http_client_or_default() -> reqwest::Client {
build_http_client().unwrap_or_else(|_| reqwest::Client::new())
}
/// Build a `reqwest::Client` from an explicit [`ProxyConfig`]. Used by tests
/// and by callers that want to override process-level environment lookups.
///
/// When `config.proxy_url` is set it overrides the per-scheme `http_proxy`
/// and `https_proxy` fields and is registered as both an HTTP and HTTPS
/// proxy so a single value can route every outbound request.
pub fn build_http_client_with(config: &ProxyConfig) -> Result<reqwest::Client, ApiError> {
let mut builder = reqwest::Client::builder().no_proxy();
let no_proxy = config
.no_proxy
.as_deref()
.and_then(reqwest::NoProxy::from_string);
let (http_proxy_url, https_proxy_url) = match config.proxy_url.as_deref() {
Some(unified) => (Some(unified), Some(unified)),
None => (config.http_proxy.as_deref(), config.https_proxy.as_deref()),
};
if let Some(url) = https_proxy_url {
let mut proxy = reqwest::Proxy::https(url)?;
if let Some(filter) = no_proxy.clone() {
proxy = proxy.no_proxy(Some(filter));
}
builder = builder.proxy(proxy);
}
if let Some(url) = http_proxy_url {
let mut proxy = reqwest::Proxy::http(url)?;
if let Some(filter) = no_proxy.clone() {
proxy = proxy.no_proxy(Some(filter));
}
builder = builder.proxy(proxy);
}
Ok(builder.build()?)
}
fn first_non_empty<F>(keys: &[&str], lookup: &mut F) -> Option<String>
where
F: FnMut(&str) -> Option<String>,
{
keys.iter()
.find_map(|key| lookup(key).filter(|value| !value.is_empty()))
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use super::{build_http_client_with, ProxyConfig};
fn config_from_map(pairs: &[(&str, &str)]) -> ProxyConfig {
let map: HashMap<String, String> = pairs
.iter()
.map(|(key, value)| ((*key).to_string(), (*value).to_string()))
.collect();
ProxyConfig::from_lookup(|key| map.get(key).cloned())
}
#[test]
fn proxy_config_is_empty_when_no_env_vars_are_set() {
// given
let config = config_from_map(&[]);
// when
let empty = config.is_empty();
// then
assert!(empty);
assert_eq!(config, ProxyConfig::default());
}
#[test]
fn proxy_config_reads_uppercase_http_https_and_no_proxy() {
// given
let pairs = [
("HTTP_PROXY", "http://proxy.internal:3128"),
("HTTPS_PROXY", "http://secure.internal:3129"),
("NO_PROXY", "localhost,127.0.0.1,.corp"),
];
// when
let config = config_from_map(&pairs);
// then
assert_eq!(
config.http_proxy.as_deref(),
Some("http://proxy.internal:3128")
);
assert_eq!(
config.https_proxy.as_deref(),
Some("http://secure.internal:3129")
);
assert_eq!(
config.no_proxy.as_deref(),
Some("localhost,127.0.0.1,.corp")
);
assert!(!config.is_empty());
}
#[test]
fn proxy_config_falls_back_to_lowercase_keys() {
// given
let pairs = [
("http_proxy", "http://lower.internal:3128"),
("https_proxy", "http://lower-secure.internal:3129"),
("no_proxy", ".lower"),
];
// when
let config = config_from_map(&pairs);
// then
assert_eq!(
config.http_proxy.as_deref(),
Some("http://lower.internal:3128")
);
assert_eq!(
config.https_proxy.as_deref(),
Some("http://lower-secure.internal:3129")
);
assert_eq!(config.no_proxy.as_deref(), Some(".lower"));
}
#[test]
fn proxy_config_prefers_uppercase_over_lowercase_when_both_set() {
// given
let pairs = [
("HTTP_PROXY", "http://upper.internal:3128"),
("http_proxy", "http://lower.internal:3128"),
];
// when
let config = config_from_map(&pairs);
// then
assert_eq!(
config.http_proxy.as_deref(),
Some("http://upper.internal:3128")
);
}
#[test]
fn proxy_config_treats_empty_strings_as_unset() {
// given
let pairs = [("HTTP_PROXY", ""), ("http_proxy", "")];
// when
let config = config_from_map(&pairs);
// then
assert!(config.http_proxy.is_none());
}
#[test]
fn build_http_client_succeeds_when_no_proxy_is_configured() {
// given
let config = ProxyConfig::default();
// when
let result = build_http_client_with(&config);
// then
assert!(result.is_ok());
}
#[test]
fn build_http_client_succeeds_with_valid_http_and_https_proxies() {
// given
let config = ProxyConfig {
http_proxy: Some("http://proxy.internal:3128".to_string()),
https_proxy: Some("http://secure.internal:3129".to_string()),
no_proxy: Some("localhost,127.0.0.1".to_string()),
proxy_url: None,
};
// when
let result = build_http_client_with(&config);
// then
assert!(result.is_ok());
}
#[test]
fn build_http_client_returns_http_error_for_invalid_proxy_url() {
// given
let config = ProxyConfig {
http_proxy: None,
https_proxy: Some("not a url".to_string()),
no_proxy: None,
proxy_url: None,
};
// when
let result = build_http_client_with(&config);
// then
let error = result.expect_err("invalid proxy URL must be reported as a build failure");
assert!(
matches!(error, crate::error::ApiError::Http(_)),
"expected ApiError::Http for invalid proxy URL, got: {error:?}"
);
}
#[test]
fn from_proxy_url_sets_unified_field_and_leaves_per_scheme_empty() {
// given / when
let config = ProxyConfig::from_proxy_url("http://unified.internal:3128");
// then
assert_eq!(
config.proxy_url.as_deref(),
Some("http://unified.internal:3128")
);
assert!(config.http_proxy.is_none());
assert!(config.https_proxy.is_none());
assert!(!config.is_empty());
}
#[test]
fn build_http_client_succeeds_with_unified_proxy_url() {
// given
let config = ProxyConfig {
proxy_url: Some("http://unified.internal:3128".to_string()),
no_proxy: Some("localhost".to_string()),
..ProxyConfig::default()
};
// when
let result = build_http_client_with(&config);
// then
assert!(result.is_ok());
}
#[test]
fn proxy_url_takes_precedence_over_per_scheme_fields() {
// given – both per-scheme and unified are set
let config = ProxyConfig {
http_proxy: Some("http://per-scheme.internal:1111".to_string()),
https_proxy: Some("http://per-scheme.internal:2222".to_string()),
no_proxy: None,
proxy_url: Some("http://unified.internal:3128".to_string()),
};
// when – building succeeds (the unified URL is valid)
let result = build_http_client_with(&config);
// then
assert!(result.is_ok());
}
#[test]
fn build_http_client_returns_error_for_invalid_unified_proxy_url() {
// given
let config = ProxyConfig::from_proxy_url("not a url");
// when
let result = build_http_client_with(&config);
// then
assert!(
matches!(result, Err(crate::error::ApiError::Http(_))),
"invalid unified proxy URL should fail: {result:?}"
);
}
}
+39
View File
@@ -0,0 +1,39 @@
mod client;
mod error;
mod http_client;
mod prompt_cache;
mod providers;
mod sse;
mod types;
pub use client::{
oauth_token_is_expired, read_base_url, read_xai_base_url, resolve_saved_oauth_token,
resolve_startup_auth_source, MessageStream, OAuthTokenSet, ProviderClient,
};
pub use error::ApiError;
pub use http_client::{
build_http_client, build_http_client_or_default, build_http_client_with, ProxyConfig,
};
pub use prompt_cache::{
CacheBreakEvent, PromptCache, PromptCacheConfig, PromptCachePaths, PromptCacheRecord,
PromptCacheStats,
};
pub use providers::anthropic::{AnthropicClient, AnthropicClient as ApiClient, AuthSource};
pub use providers::openai_compat::{OpenAiCompatClient, OpenAiCompatConfig};
pub use providers::{
detect_provider_kind, max_tokens_for_model, max_tokens_for_model_with_override,
resolve_model_alias, ProviderKind,
};
pub use sse::{parse_frame, SseParser};
pub use types::{
ContentBlockDelta, ContentBlockDeltaEvent, ContentBlockStartEvent, ContentBlockStopEvent,
InputContentBlock, InputMessage, MessageDelta, MessageDeltaEvent, MessageRequest,
MessageResponse, MessageStartEvent, MessageStopEvent, OutputContentBlock, StreamEvent,
ToolChoice, ToolDefinition, ToolResultContentBlock, Usage,
};
pub use telemetry::{
AnalyticsEvent, AnthropicRequestProfile, ClientIdentity, JsonlTelemetrySink,
MemoryTelemetrySink, SessionTraceRecord, SessionTracer, TelemetryEvent, TelemetrySink,
DEFAULT_ANTHROPIC_VERSION,
};
@@ -0,0 +1,735 @@
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use crate::types::{MessageRequest, MessageResponse, Usage};
const DEFAULT_COMPLETION_TTL_SECS: u64 = 30;
const DEFAULT_PROMPT_TTL_SECS: u64 = 5 * 60;
const DEFAULT_BREAK_MIN_DROP: u32 = 2_000;
const MAX_SANITIZED_LENGTH: usize = 80;
const REQUEST_FINGERPRINT_VERSION: u32 = 1;
const REQUEST_FINGERPRINT_PREFIX: &str = "v1";
const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
#[derive(Debug, Clone)]
pub struct PromptCacheConfig {
pub session_id: String,
pub completion_ttl: Duration,
pub prompt_ttl: Duration,
pub cache_break_min_drop: u32,
}
impl PromptCacheConfig {
#[must_use]
pub fn new(session_id: impl Into<String>) -> Self {
Self {
session_id: session_id.into(),
completion_ttl: Duration::from_secs(DEFAULT_COMPLETION_TTL_SECS),
prompt_ttl: Duration::from_secs(DEFAULT_PROMPT_TTL_SECS),
cache_break_min_drop: DEFAULT_BREAK_MIN_DROP,
}
}
}
impl Default for PromptCacheConfig {
fn default() -> Self {
Self::new("default")
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptCachePaths {
pub root: PathBuf,
pub session_dir: PathBuf,
pub completion_dir: PathBuf,
pub session_state_path: PathBuf,
pub stats_path: PathBuf,
}
impl PromptCachePaths {
#[must_use]
pub fn for_session(session_id: &str) -> Self {
let root = base_cache_root();
let session_dir = root.join(sanitize_path_segment(session_id));
let completion_dir = session_dir.join("completions");
Self {
root,
session_state_path: session_dir.join("session-state.json"),
stats_path: session_dir.join("stats.json"),
session_dir,
completion_dir,
}
}
#[must_use]
pub fn completion_entry_path(&self, request_hash: &str) -> PathBuf {
self.completion_dir.join(format!("{request_hash}.json"))
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptCacheStats {
pub tracked_requests: u64,
pub completion_cache_hits: u64,
pub completion_cache_misses: u64,
pub completion_cache_writes: u64,
pub expected_invalidations: u64,
pub unexpected_cache_breaks: u64,
pub total_cache_creation_input_tokens: u64,
pub total_cache_read_input_tokens: u64,
pub last_cache_creation_input_tokens: Option<u32>,
pub last_cache_read_input_tokens: Option<u32>,
pub last_request_hash: Option<String>,
pub last_completion_cache_key: Option<String>,
pub last_break_reason: Option<String>,
pub last_cache_source: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CacheBreakEvent {
pub unexpected: bool,
pub reason: String,
pub previous_cache_read_input_tokens: u32,
pub current_cache_read_input_tokens: u32,
pub token_drop: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PromptCacheRecord {
pub cache_break: Option<CacheBreakEvent>,
pub stats: PromptCacheStats,
}
#[derive(Debug, Clone)]
pub struct PromptCache {
inner: Arc<Mutex<PromptCacheInner>>,
}
impl PromptCache {
#[must_use]
pub fn new(session_id: impl Into<String>) -> Self {
Self::with_config(PromptCacheConfig::new(session_id))
}
#[must_use]
pub fn with_config(config: PromptCacheConfig) -> Self {
let paths = PromptCachePaths::for_session(&config.session_id);
let stats = read_json::<PromptCacheStats>(&paths.stats_path).unwrap_or_default();
let previous = read_json::<TrackedPromptState>(&paths.session_state_path);
Self {
inner: Arc::new(Mutex::new(PromptCacheInner {
config,
paths,
stats,
previous,
})),
}
}
#[must_use]
pub fn paths(&self) -> PromptCachePaths {
self.lock().paths.clone()
}
#[must_use]
pub fn stats(&self) -> PromptCacheStats {
self.lock().stats.clone()
}
#[must_use]
pub fn lookup_completion(&self, request: &MessageRequest) -> Option<MessageResponse> {
let request_hash = request_hash_hex(request);
let (paths, ttl) = {
let inner = self.lock();
(inner.paths.clone(), inner.config.completion_ttl)
};
let entry_path = paths.completion_entry_path(&request_hash);
let entry = read_json::<CompletionCacheEntry>(&entry_path);
let Some(entry) = entry else {
let mut inner = self.lock();
inner.stats.completion_cache_misses += 1;
inner.stats.last_completion_cache_key = Some(request_hash);
persist_state(&inner);
return None;
};
if entry.fingerprint_version != current_fingerprint_version() {
let mut inner = self.lock();
inner.stats.completion_cache_misses += 1;
inner.stats.last_completion_cache_key = Some(request_hash.clone());
let _ = fs::remove_file(entry_path);
persist_state(&inner);
return None;
}
let expired = now_unix_secs().saturating_sub(entry.cached_at_unix_secs) >= ttl.as_secs();
let mut inner = self.lock();
inner.stats.last_completion_cache_key = Some(request_hash.clone());
if expired {
inner.stats.completion_cache_misses += 1;
let _ = fs::remove_file(entry_path);
persist_state(&inner);
return None;
}
inner.stats.completion_cache_hits += 1;
apply_usage_to_stats(
&mut inner.stats,
&entry.response.usage,
&request_hash,
"completion-cache",
);
inner.previous = Some(TrackedPromptState::from_usage(
request,
&entry.response.usage,
));
persist_state(&inner);
Some(entry.response)
}
#[must_use]
pub fn record_response(
&self,
request: &MessageRequest,
response: &MessageResponse,
) -> PromptCacheRecord {
self.record_usage_internal(request, &response.usage, Some(response))
}
#[must_use]
pub fn record_usage(&self, request: &MessageRequest, usage: &Usage) -> PromptCacheRecord {
self.record_usage_internal(request, usage, None)
}
fn record_usage_internal(
&self,
request: &MessageRequest,
usage: &Usage,
response: Option<&MessageResponse>,
) -> PromptCacheRecord {
let request_hash = request_hash_hex(request);
let mut inner = self.lock();
let previous = inner.previous.clone();
let current = TrackedPromptState::from_usage(request, usage);
let cache_break = detect_cache_break(&inner.config, previous.as_ref(), &current);
inner.stats.tracked_requests += 1;
apply_usage_to_stats(&mut inner.stats, usage, &request_hash, "api-response");
if let Some(event) = &cache_break {
if event.unexpected {
inner.stats.unexpected_cache_breaks += 1;
} else {
inner.stats.expected_invalidations += 1;
}
inner.stats.last_break_reason = Some(event.reason.clone());
}
inner.previous = Some(current);
if let Some(response) = response {
write_completion_entry(&inner.paths, &request_hash, response);
inner.stats.completion_cache_writes += 1;
}
persist_state(&inner);
PromptCacheRecord {
cache_break,
stats: inner.stats.clone(),
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, PromptCacheInner> {
self.inner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
#[derive(Debug)]
struct PromptCacheInner {
config: PromptCacheConfig,
paths: PromptCachePaths,
stats: PromptCacheStats,
previous: Option<TrackedPromptState>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CompletionCacheEntry {
cached_at_unix_secs: u64,
#[serde(default = "current_fingerprint_version")]
fingerprint_version: u32,
response: MessageResponse,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct TrackedPromptState {
observed_at_unix_secs: u64,
#[serde(default = "current_fingerprint_version")]
fingerprint_version: u32,
model_hash: u64,
system_hash: u64,
tools_hash: u64,
messages_hash: u64,
cache_read_input_tokens: u32,
}
impl TrackedPromptState {
fn from_usage(request: &MessageRequest, usage: &Usage) -> Self {
let hashes = RequestFingerprints::from_request(request);
Self {
observed_at_unix_secs: now_unix_secs(),
fingerprint_version: current_fingerprint_version(),
model_hash: hashes.model,
system_hash: hashes.system,
tools_hash: hashes.tools,
messages_hash: hashes.messages,
cache_read_input_tokens: usage.cache_read_input_tokens,
}
}
}
#[derive(Debug, Clone, Copy)]
struct RequestFingerprints {
model: u64,
system: u64,
tools: u64,
messages: u64,
}
impl RequestFingerprints {
fn from_request(request: &MessageRequest) -> Self {
Self {
model: hash_serializable(&request.model),
system: hash_serializable(&request.system),
tools: hash_serializable(&request.tools),
messages: hash_serializable(&request.messages),
}
}
}
fn detect_cache_break(
config: &PromptCacheConfig,
previous: Option<&TrackedPromptState>,
current: &TrackedPromptState,
) -> Option<CacheBreakEvent> {
let previous = previous?;
if previous.fingerprint_version != current.fingerprint_version {
return Some(CacheBreakEvent {
unexpected: false,
reason: format!(
"fingerprint version changed (v{} -> v{})",
previous.fingerprint_version, current.fingerprint_version
),
previous_cache_read_input_tokens: previous.cache_read_input_tokens,
current_cache_read_input_tokens: current.cache_read_input_tokens,
token_drop: previous
.cache_read_input_tokens
.saturating_sub(current.cache_read_input_tokens),
});
}
let token_drop = previous
.cache_read_input_tokens
.saturating_sub(current.cache_read_input_tokens);
if token_drop < config.cache_break_min_drop {
return None;
}
let mut reasons = Vec::new();
if previous.model_hash != current.model_hash {
reasons.push("model changed");
}
if previous.system_hash != current.system_hash {
reasons.push("system prompt changed");
}
if previous.tools_hash != current.tools_hash {
reasons.push("tool definitions changed");
}
if previous.messages_hash != current.messages_hash {
reasons.push("message payload changed");
}
let elapsed = current
.observed_at_unix_secs
.saturating_sub(previous.observed_at_unix_secs);
let (unexpected, reason) = if reasons.is_empty() {
if elapsed > config.prompt_ttl.as_secs() {
(
false,
format!("possible prompt cache TTL expiry after {elapsed}s"),
)
} else {
(
true,
"cache read tokens dropped while prompt fingerprint remained stable".to_string(),
)
}
} else {
(false, reasons.join(", "))
};
Some(CacheBreakEvent {
unexpected,
reason,
previous_cache_read_input_tokens: previous.cache_read_input_tokens,
current_cache_read_input_tokens: current.cache_read_input_tokens,
token_drop,
})
}
fn apply_usage_to_stats(
stats: &mut PromptCacheStats,
usage: &Usage,
request_hash: &str,
source: &str,
) {
stats.total_cache_creation_input_tokens += u64::from(usage.cache_creation_input_tokens);
stats.total_cache_read_input_tokens += u64::from(usage.cache_read_input_tokens);
stats.last_cache_creation_input_tokens = Some(usage.cache_creation_input_tokens);
stats.last_cache_read_input_tokens = Some(usage.cache_read_input_tokens);
stats.last_request_hash = Some(request_hash.to_string());
stats.last_cache_source = Some(source.to_string());
}
fn persist_state(inner: &PromptCacheInner) {
let _ = ensure_cache_dirs(&inner.paths);
let _ = write_json(&inner.paths.stats_path, &inner.stats);
if let Some(previous) = &inner.previous {
let _ = write_json(&inner.paths.session_state_path, previous);
}
}
fn write_completion_entry(
paths: &PromptCachePaths,
request_hash: &str,
response: &MessageResponse,
) {
let _ = ensure_cache_dirs(paths);
let entry = CompletionCacheEntry {
cached_at_unix_secs: now_unix_secs(),
fingerprint_version: current_fingerprint_version(),
response: response.clone(),
};
let _ = write_json(&paths.completion_entry_path(request_hash), &entry);
}
fn ensure_cache_dirs(paths: &PromptCachePaths) -> std::io::Result<()> {
fs::create_dir_all(&paths.completion_dir)
}
fn write_json<T: Serialize>(path: &Path, value: &T) -> std::io::Result<()> {
let json = serde_json::to_vec_pretty(value)
.map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?;
fs::write(path, json)
}
fn read_json<T: for<'de> Deserialize<'de>>(path: &Path) -> Option<T> {
let bytes = fs::read(path).ok()?;
serde_json::from_slice(&bytes).ok()
}
fn request_hash_hex(request: &MessageRequest) -> String {
format!(
"{REQUEST_FINGERPRINT_PREFIX}-{:016x}",
hash_serializable(request)
)
}
fn hash_serializable<T: Serialize>(value: &T) -> u64 {
let json = serde_json::to_vec(value).unwrap_or_default();
stable_hash_bytes(&json)
}
fn sanitize_path_segment(value: &str) -> String {
let sanitized: String = value
.chars()
.map(|ch| if ch.is_ascii_alphanumeric() { ch } else { '-' })
.collect();
if sanitized.len() <= MAX_SANITIZED_LENGTH {
return sanitized;
}
let suffix = format!("-{:x}", hash_string(value));
format!(
"{}{}",
&sanitized[..MAX_SANITIZED_LENGTH.saturating_sub(suffix.len())],
suffix
)
}
fn hash_string(value: &str) -> u64 {
stable_hash_bytes(value.as_bytes())
}
fn base_cache_root() -> PathBuf {
if let Some(config_home) = std::env::var_os("CLAUDE_CONFIG_HOME") {
return PathBuf::from(config_home)
.join("cache")
.join("prompt-cache");
}
if let Some(home) = std::env::var_os("HOME") {
return PathBuf::from(home)
.join(".claude")
.join("cache")
.join("prompt-cache");
}
std::env::temp_dir().join("claude-prompt-cache")
}
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |duration| duration.as_secs())
}
const fn current_fingerprint_version() -> u32 {
REQUEST_FINGERPRINT_VERSION
}
fn stable_hash_bytes(bytes: &[u8]) -> u64 {
let mut hash = FNV_OFFSET_BASIS;
for byte in bytes {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(FNV_PRIME);
}
hash
}
#[cfg(test)]
mod tests {
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use super::{
detect_cache_break, read_json, request_hash_hex, sanitize_path_segment, PromptCache,
PromptCacheConfig, PromptCachePaths, TrackedPromptState, REQUEST_FINGERPRINT_PREFIX,
};
use crate::types::{InputMessage, MessageRequest, MessageResponse, OutputContentBlock, Usage};
fn test_env_lock() -> std::sync::MutexGuard<'static, ()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[test]
fn path_builder_sanitizes_session_identifier() {
let paths = PromptCachePaths::for_session("session:/with spaces");
let session_dir = paths
.session_dir
.file_name()
.and_then(|value| value.to_str())
.expect("session dir name");
assert_eq!(session_dir, "session--with-spaces");
assert!(paths.completion_dir.ends_with("completions"));
assert!(paths.stats_path.ends_with("stats.json"));
assert!(paths.session_state_path.ends_with("session-state.json"));
}
#[test]
fn request_fingerprint_drives_unexpected_break_detection() {
let request = sample_request("same");
let previous = TrackedPromptState::from_usage(
&request,
&Usage {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 6_000,
output_tokens: 0,
},
);
let current = TrackedPromptState::from_usage(
&request,
&Usage {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 1_000,
output_tokens: 0,
},
);
let event = detect_cache_break(&PromptCacheConfig::default(), Some(&previous), &current)
.expect("break should be detected");
assert!(event.unexpected);
assert!(event.reason.contains("stable"));
}
#[test]
fn changed_prompt_marks_break_as_expected() {
let previous_request = sample_request("first");
let current_request = sample_request("second");
let previous = TrackedPromptState::from_usage(
&previous_request,
&Usage {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 6_000,
output_tokens: 0,
},
);
let current = TrackedPromptState::from_usage(
&current_request,
&Usage {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 1_000,
output_tokens: 0,
},
);
let event = detect_cache_break(&PromptCacheConfig::default(), Some(&previous), &current)
.expect("break should be detected");
assert!(!event.unexpected);
assert!(event.reason.contains("message payload changed"));
}
#[test]
fn completion_cache_round_trip_persists_recent_response() {
let _guard = test_env_lock();
let temp_root = std::env::temp_dir().join(format!(
"prompt-cache-test-{}-{}",
std::process::id(),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("time")
.as_nanos()
));
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
let cache = PromptCache::new("unit-test-session");
let request = sample_request("cache me");
let response = sample_response(42, 12, "cached");
assert!(cache.lookup_completion(&request).is_none());
let record = cache.record_response(&request, &response);
assert!(record.cache_break.is_none());
let cached = cache
.lookup_completion(&request)
.expect("cached response should load");
assert_eq!(cached.content, response.content);
let stats = cache.stats();
assert_eq!(stats.completion_cache_hits, 1);
assert_eq!(stats.completion_cache_misses, 1);
assert_eq!(stats.completion_cache_writes, 1);
let persisted = read_json::<super::PromptCacheStats>(&cache.paths().stats_path)
.expect("stats should persist");
assert_eq!(persisted.completion_cache_hits, 1);
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
std::env::remove_var("CLAUDE_CONFIG_HOME");
}
#[test]
fn distinct_requests_do_not_collide_in_completion_cache() {
let _guard = test_env_lock();
let temp_root = std::env::temp_dir().join(format!(
"prompt-cache-distinct-{}-{}",
std::process::id(),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("time")
.as_nanos()
));
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
let cache = PromptCache::new("distinct-request-session");
let first_request = sample_request("first");
let second_request = sample_request("second");
let response = sample_response(42, 12, "cached");
let _ = cache.record_response(&first_request, &response);
assert!(cache.lookup_completion(&second_request).is_none());
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
std::env::remove_var("CLAUDE_CONFIG_HOME");
}
#[test]
fn expired_completion_entries_are_not_reused() {
let _guard = test_env_lock();
let temp_root = std::env::temp_dir().join(format!(
"prompt-cache-expired-{}-{}",
std::process::id(),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("time")
.as_nanos()
));
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
let cache = PromptCache::with_config(PromptCacheConfig {
session_id: "expired-session".to_string(),
completion_ttl: Duration::ZERO,
..PromptCacheConfig::default()
});
let request = sample_request("expire me");
let response = sample_response(7, 3, "stale");
let _ = cache.record_response(&request, &response);
assert!(cache.lookup_completion(&request).is_none());
let stats = cache.stats();
assert_eq!(stats.completion_cache_hits, 0);
assert_eq!(stats.completion_cache_misses, 1);
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
std::env::remove_var("CLAUDE_CONFIG_HOME");
}
#[test]
fn sanitize_path_caps_long_values() {
let long_value = "x".repeat(200);
let sanitized = sanitize_path_segment(&long_value);
assert!(sanitized.len() <= 80);
}
#[test]
fn request_hashes_are_versioned_and_stable() {
let request = sample_request("stable");
let first = request_hash_hex(&request);
let second = request_hash_hex(&request);
assert_eq!(first, second);
assert!(first.starts_with(REQUEST_FINGERPRINT_PREFIX));
}
fn sample_request(text: &str) -> MessageRequest {
MessageRequest {
model: "claude-3-7-sonnet-latest".to_string(),
max_tokens: 64,
messages: vec![InputMessage::user_text(text)],
system: Some("system".to_string()),
tools: None,
tool_choice: None,
stream: false,
..Default::default()
}
}
fn sample_response(
cache_read_input_tokens: u32,
output_tokens: u32,
text: &str,
) -> MessageResponse {
MessageResponse {
id: "msg_test".to_string(),
kind: "message".to_string(),
role: "assistant".to_string(),
content: vec![OutputContentBlock::Text {
text: text.to_string(),
}],
model: "claude-3-7-sonnet-latest".to_string(),
stop_reason: Some("end_turn".to_string()),
stop_sequence: None,
usage: Usage {
input_tokens: 10,
cache_creation_input_tokens: 5,
cache_read_input_tokens,
output_tokens,
},
request_id: Some("req_test".to_string()),
}
}
}
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
+330
View File
@@ -0,0 +1,330 @@
use crate::error::ApiError;
use crate::types::StreamEvent;
#[derive(Debug, Default)]
pub struct SseParser {
buffer: Vec<u8>,
provider: Option<String>,
model: Option<String>,
}
impl SseParser {
#[must_use]
pub fn new() -> Self {
Self::default()
}
/// Attach the provider name and model to this parser so that JSON
/// deserialization failures within streamed frames carry enough context
/// for callers to understand which upstream produced the unparseable
/// payload.
#[must_use]
pub fn with_context(mut self, provider: impl Into<String>, model: impl Into<String>) -> Self {
self.provider = Some(provider.into());
self.model = Some(model.into());
self
}
pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<StreamEvent>, ApiError> {
self.buffer.extend_from_slice(chunk);
let mut events = Vec::new();
while let Some(frame) = self.next_frame() {
if let Some(event) = self.parse_frame_with_context(&frame)? {
events.push(event);
}
}
Ok(events)
}
pub fn finish(&mut self) -> Result<Vec<StreamEvent>, ApiError> {
if self.buffer.is_empty() {
return Ok(Vec::new());
}
let trailing = std::mem::take(&mut self.buffer);
match self.parse_frame_with_context(&String::from_utf8_lossy(&trailing))? {
Some(event) => Ok(vec![event]),
None => Ok(Vec::new()),
}
}
fn parse_frame_with_context(&self, frame: &str) -> Result<Option<StreamEvent>, ApiError> {
let provider = self.provider.as_deref().unwrap_or("unknown");
let model = self.model.as_deref().unwrap_or("unknown");
parse_frame_with_provider(frame, provider, model)
}
fn next_frame(&mut self) -> Option<String> {
let separator = self
.buffer
.windows(2)
.position(|window| window == b"\n\n")
.map(|position| (position, 2))
.or_else(|| {
self.buffer
.windows(4)
.position(|window| window == b"\r\n\r\n")
.map(|position| (position, 4))
})?;
let (position, separator_len) = separator;
let frame = self
.buffer
.drain(..position + separator_len)
.collect::<Vec<_>>();
let frame_len = frame.len().saturating_sub(separator_len);
Some(String::from_utf8_lossy(&frame[..frame_len]).into_owned())
}
}
pub fn parse_frame(frame: &str) -> Result<Option<StreamEvent>, ApiError> {
parse_frame_with_provider(frame, "unknown", "unknown")
}
pub(crate) fn parse_frame_with_provider(
frame: &str,
provider: &str,
model: &str,
) -> Result<Option<StreamEvent>, ApiError> {
let trimmed = frame.trim();
if trimmed.is_empty() {
return Ok(None);
}
let mut data_lines = Vec::new();
let mut event_name: Option<&str> = None;
for line in trimmed.lines() {
if line.starts_with(':') {
continue;
}
if let Some(name) = line.strip_prefix("event:") {
event_name = Some(name.trim());
continue;
}
if let Some(data) = line.strip_prefix("data:") {
data_lines.push(data.trim_start());
}
}
if matches!(event_name, Some("ping")) {
return Ok(None);
}
if data_lines.is_empty() {
return Ok(None);
}
let payload = data_lines.join("\n");
if payload == "[DONE]" {
return Ok(None);
}
serde_json::from_str::<StreamEvent>(&payload)
.map(Some)
.map_err(|error| ApiError::json_deserialize(provider, model, &payload, error))
}
#[cfg(test)]
mod tests {
use super::{parse_frame, SseParser};
use crate::types::{ContentBlockDelta, MessageDelta, OutputContentBlock, StreamEvent, Usage};
#[test]
fn parses_single_frame() {
let frame = concat!(
"event: content_block_start\n",
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"Hi\"}}\n\n"
);
let event = parse_frame(frame).expect("frame should parse");
assert_eq!(
event,
Some(StreamEvent::ContentBlockStart(
crate::types::ContentBlockStartEvent {
index: 0,
content_block: OutputContentBlock::Text {
text: "Hi".to_string(),
},
},
))
);
}
#[test]
fn parses_chunked_stream() {
let mut parser = SseParser::new();
let first = b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hel";
let second = b"lo\"}}\n\n";
assert!(parser
.push(first)
.expect("first chunk should buffer")
.is_empty());
let events = parser.push(second).expect("second chunk should parse");
assert_eq!(
events,
vec![StreamEvent::ContentBlockDelta(
crate::types::ContentBlockDeltaEvent {
index: 0,
delta: ContentBlockDelta::TextDelta {
text: "Hello".to_string(),
},
}
)]
);
}
#[test]
fn ignores_ping_and_done() {
let mut parser = SseParser::new();
let payload = concat!(
": keepalive\n",
"event: ping\n",
"data: {\"type\":\"ping\"}\n\n",
"event: message_delta\n",
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\",\"stop_sequence\":null},\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}\n\n",
"event: message_stop\n",
"data: {\"type\":\"message_stop\"}\n\n",
"data: [DONE]\n\n"
);
let events = parser
.push(payload.as_bytes())
.expect("parser should succeed");
assert_eq!(
events,
vec![
StreamEvent::MessageDelta(crate::types::MessageDeltaEvent {
delta: MessageDelta {
stop_reason: Some("tool_use".to_string()),
stop_sequence: None,
},
usage: Usage {
input_tokens: 1,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
output_tokens: 2,
},
}),
StreamEvent::MessageStop(crate::types::MessageStopEvent {}),
]
);
}
#[test]
fn ignores_data_less_event_frames() {
let frame = "event: ping\n\n";
let event = parse_frame(frame).expect("frame without data should be ignored");
assert_eq!(event, None);
}
#[test]
fn parses_split_json_across_data_lines() {
let frame = concat!(
"event: content_block_delta\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\n",
"data: \"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n"
);
let event = parse_frame(frame).expect("frame should parse");
assert_eq!(
event,
Some(StreamEvent::ContentBlockDelta(
crate::types::ContentBlockDeltaEvent {
index: 0,
delta: ContentBlockDelta::TextDelta {
text: "Hello".to_string(),
},
}
))
);
}
#[test]
fn parses_thinking_content_block_start() {
let frame = concat!(
"event: content_block_start\n",
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\",\"signature\":null}}\n\n"
);
let event = parse_frame(frame).expect("frame should parse");
assert_eq!(
event,
Some(StreamEvent::ContentBlockStart(
crate::types::ContentBlockStartEvent {
index: 0,
content_block: OutputContentBlock::Thinking {
thinking: String::new(),
signature: None,
},
},
))
);
}
#[test]
fn parses_thinking_related_deltas() {
let thinking = concat!(
"event: content_block_delta\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"step 1\"}}\n\n"
);
let signature = concat!(
"event: content_block_delta\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"sig_123\"}}\n\n"
);
let thinking_event = parse_frame(thinking).expect("thinking delta should parse");
let signature_event = parse_frame(signature).expect("signature delta should parse");
assert_eq!(
thinking_event,
Some(StreamEvent::ContentBlockDelta(
crate::types::ContentBlockDeltaEvent {
index: 0,
delta: ContentBlockDelta::ThinkingDelta {
thinking: "step 1".to_string(),
},
}
))
);
assert_eq!(
signature_event,
Some(StreamEvent::ContentBlockDelta(
crate::types::ContentBlockDeltaEvent {
index: 0,
delta: ContentBlockDelta::SignatureDelta {
signature: "sig_123".to_string(),
},
}
))
);
}
#[test]
fn given_message_delta_frame_with_empty_usage_when_parsed_then_usage_defaults_to_zero() {
// given
let frame = concat!(
"event: message_delta\n",
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{}}\n\n"
);
// when
let event = parse_frame(frame).expect("frame should parse");
// then
assert_eq!(
event,
Some(StreamEvent::MessageDelta(crate::types::MessageDeltaEvent {
delta: MessageDelta {
stop_reason: Some("end_turn".to_string()),
stop_sequence: None,
},
usage: Usage::default(),
}))
);
}
}
+310
View File
@@ -0,0 +1,310 @@
use runtime::{pricing_for_model, TokenUsage, UsageCostEstimate};
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct MessageRequest {
pub model: String,
pub max_tokens: u32,
pub messages: Vec<InputMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub system: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<ToolDefinition>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_choice: Option<ToolChoice>,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub stream: bool,
/// OpenAI-compatible tuning parameters. Optional — omitted from payload when None.
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub frequency_penalty: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub presence_penalty: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop: Option<Vec<String>>,
/// Reasoning effort level for OpenAI-compatible reasoning models (e.g. `o4-mini`).
/// Accepted values: `"low"`, `"medium"`, `"high"`. Omitted when `None`.
/// Silently ignored by backends that do not support it.
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_effort: Option<String>,
}
impl MessageRequest {
#[must_use]
pub fn with_streaming(mut self) -> Self {
self.stream = true;
self
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct InputMessage {
pub role: String,
pub content: Vec<InputContentBlock>,
}
impl InputMessage {
#[must_use]
pub fn user_text(text: impl Into<String>) -> Self {
Self {
role: "user".to_string(),
content: vec![InputContentBlock::Text { text: text.into() }],
}
}
#[must_use]
pub fn user_tool_result(
tool_use_id: impl Into<String>,
content: impl Into<String>,
is_error: bool,
) -> Self {
Self {
role: "user".to_string(),
content: vec![InputContentBlock::ToolResult {
tool_use_id: tool_use_id.into(),
content: vec![ToolResultContentBlock::Text {
text: content.into(),
}],
is_error,
}],
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum InputContentBlock {
Text {
text: String,
},
ToolUse {
id: String,
name: String,
input: Value,
},
ToolResult {
tool_use_id: String,
content: Vec<ToolResultContentBlock>,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
is_error: bool,
},
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ToolResultContentBlock {
Text { text: String },
Json { value: Value },
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolDefinition {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
pub input_schema: Value,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ToolChoice {
Auto,
Any,
Tool { name: String },
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MessageResponse {
pub id: String,
#[serde(rename = "type")]
pub kind: String,
pub role: String,
pub content: Vec<OutputContentBlock>,
pub model: String,
#[serde(default)]
pub stop_reason: Option<String>,
#[serde(default)]
pub stop_sequence: Option<String>,
#[serde(default)]
pub usage: Usage,
#[serde(default)]
pub request_id: Option<String>,
}
impl MessageResponse {
#[must_use]
pub fn total_tokens(&self) -> u32 {
self.usage.total_tokens()
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum OutputContentBlock {
Text {
text: String,
},
ToolUse {
id: String,
name: String,
input: Value,
},
Thinking {
#[serde(default)]
thinking: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
signature: Option<String>,
},
RedactedThinking {
data: Value,
},
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
#[serde(default)]
pub input_tokens: u32,
#[serde(default)]
pub cache_creation_input_tokens: u32,
#[serde(default)]
pub cache_read_input_tokens: u32,
#[serde(default)]
pub output_tokens: u32,
}
impl Usage {
#[must_use]
pub const fn total_tokens(&self) -> u32 {
self.input_tokens
+ self.output_tokens
+ self.cache_creation_input_tokens
+ self.cache_read_input_tokens
}
#[must_use]
pub const fn token_usage(&self) -> TokenUsage {
TokenUsage {
input_tokens: self.input_tokens,
output_tokens: self.output_tokens,
cache_creation_input_tokens: self.cache_creation_input_tokens,
cache_read_input_tokens: self.cache_read_input_tokens,
}
}
#[must_use]
pub fn estimated_cost_usd(&self, model: &str) -> UsageCostEstimate {
let usage = self.token_usage();
pricing_for_model(model).map_or_else(
|| usage.estimate_cost_usd(),
|pricing| usage.estimate_cost_usd_with_pricing(pricing),
)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MessageStartEvent {
pub message: MessageResponse,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MessageDeltaEvent {
pub delta: MessageDelta,
#[serde(default)]
pub usage: Usage,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MessageDelta {
#[serde(default)]
pub stop_reason: Option<String>,
#[serde(default)]
pub stop_sequence: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ContentBlockStartEvent {
pub index: u32,
pub content_block: OutputContentBlock,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ContentBlockDeltaEvent {
pub index: u32,
pub delta: ContentBlockDelta,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlockDelta {
TextDelta { text: String },
InputJsonDelta { partial_json: String },
ThinkingDelta { thinking: String },
SignatureDelta { signature: String },
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ContentBlockStopEvent {
pub index: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MessageStopEvent {}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum StreamEvent {
MessageStart(MessageStartEvent),
MessageDelta(MessageDeltaEvent),
ContentBlockStart(ContentBlockStartEvent),
ContentBlockDelta(ContentBlockDeltaEvent),
ContentBlockStop(ContentBlockStopEvent),
MessageStop(MessageStopEvent),
}
#[cfg(test)]
mod tests {
use runtime::format_usd;
use super::{MessageResponse, Usage};
#[test]
fn usage_total_tokens_includes_cache_tokens() {
let usage = Usage {
input_tokens: 10,
cache_creation_input_tokens: 2,
cache_read_input_tokens: 3,
output_tokens: 4,
};
assert_eq!(usage.total_tokens(), 19);
assert_eq!(usage.token_usage().total_tokens(), 19);
}
#[test]
fn message_response_estimates_cost_from_model_usage() {
let response = MessageResponse {
id: "msg_cost".to_string(),
kind: "message".to_string(),
role: "assistant".to_string(),
content: Vec::new(),
model: "claude-sonnet-4-20250514".to_string(),
stop_reason: Some("end_turn".to_string()),
stop_sequence: None,
usage: Usage {
input_tokens: 1_000_000,
cache_creation_input_tokens: 100_000,
cache_read_input_tokens: 200_000,
output_tokens: 500_000,
},
request_id: None,
};
let cost = response.usage.estimated_cost_usd(&response.model);
assert_eq!(format_usd(cost.total_cost_usd()), "$54.6750");
assert_eq!(response.total_tokens(), 1_800_000);
}
}
@@ -0,0 +1,928 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::{Mutex as StdMutex, OnceLock};
use std::time::Duration;
use api::{
AnthropicClient, ApiClient, ApiError, AuthSource, ContentBlockDelta, ContentBlockDeltaEvent,
ContentBlockStartEvent, InputContentBlock, InputMessage, MessageDeltaEvent, MessageRequest,
OutputContentBlock, PromptCache, PromptCacheConfig, ProviderClient, StreamEvent, ToolChoice,
ToolDefinition,
};
use serde_json::json;
use telemetry::{ClientIdentity, MemoryTelemetrySink, SessionTracer, TelemetryEvent};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::Mutex;
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
static LOCK: OnceLock<StdMutex<()>> = OnceLock::new();
LOCK.get_or_init(|| StdMutex::new(()))
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[tokio::test]
async fn send_message_posts_json_and_parses_response() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let body = concat!(
"{",
"\"id\":\"msg_test\",",
"\"type\":\"message\",",
"\"role\":\"assistant\",",
"\"content\":[{\"type\":\"text\",\"text\":\"Hello from Claude\"}],",
"\"model\":\"claude-3-7-sonnet-latest\",",
"\"stop_reason\":\"end_turn\",",
"\"stop_sequence\":null,",
"\"usage\":{\"input_tokens\":12,\"output_tokens\":4},",
"\"request_id\":\"req_body_123\"",
"}"
);
let server = spawn_server(
state.clone(),
vec![http_response("200 OK", "application/json", body)],
)
.await;
let client = ApiClient::new("test-key")
.with_auth_token(Some("proxy-token".to_string()))
.with_base_url(server.base_url());
let response = client
.send_message(&sample_request(false))
.await
.expect("request should succeed");
assert_eq!(response.id, "msg_test");
assert_eq!(response.total_tokens(), 16);
assert_eq!(response.request_id.as_deref(), Some("req_body_123"));
assert_eq!(response.usage.cache_creation_input_tokens, 0);
assert_eq!(response.usage.cache_read_input_tokens, 0);
assert_eq!(
response.content,
vec![OutputContentBlock::Text {
text: "Hello from Claude".to_string(),
}]
);
let captured = state.lock().await;
let request = captured.first().expect("server should capture request");
assert_eq!(request.method, "POST");
assert_eq!(request.path, "/v1/messages");
assert_eq!(
request.headers.get("x-api-key").map(String::as_str),
Some("test-key")
);
assert_eq!(
request.headers.get("authorization").map(String::as_str),
Some("Bearer proxy-token")
);
assert_eq!(
request.headers.get("anthropic-version").map(String::as_str),
Some("2023-06-01")
);
assert_eq!(
request.headers.get("user-agent").map(String::as_str),
Some("claude-code/0.1.0")
);
assert_eq!(
request.headers.get("anthropic-beta").map(String::as_str),
Some("claude-code-20250219,prompt-caching-scope-2026-01-05")
);
let body: serde_json::Value =
serde_json::from_str(&request.body).expect("request body should be json");
assert_eq!(
body.get("model").and_then(serde_json::Value::as_str),
Some("claude-3-7-sonnet-latest")
);
assert!(body.get("stream").is_none());
assert_eq!(body["tools"][0]["name"], json!("get_weather"));
assert_eq!(body["tool_choice"]["type"], json!("auto"));
assert!(
body.get("betas").is_none(),
"betas must travel via the anthropic-beta header, not the request body"
);
}
#[tokio::test]
async fn send_message_blocks_oversized_requests_before_the_http_call() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![http_response("200 OK", "application/json", "{}")],
)
.await;
let client = AnthropicClient::new("test-key").with_base_url(server.base_url());
let error = client
.send_message(&MessageRequest {
model: "claude-sonnet-4-6".to_string(),
max_tokens: 64_000,
messages: vec![InputMessage {
role: "user".to_string(),
content: vec![InputContentBlock::Text {
text: "x".repeat(600_000),
}],
}],
system: Some("Keep the answer short.".to_string()),
tools: None,
tool_choice: None,
stream: false,
..Default::default()
})
.await
.expect_err("oversized request should fail local context-window preflight");
assert!(matches!(error, ApiError::ContextWindowExceeded { .. }));
assert!(
state.lock().await.is_empty(),
"preflight failure should avoid any upstream HTTP request"
);
}
#[tokio::test]
async fn send_message_applies_request_profile_and_records_telemetry() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![http_response_with_headers(
"200 OK",
"application/json",
concat!(
"{",
"\"id\":\"msg_profile\",",
"\"type\":\"message\",",
"\"role\":\"assistant\",",
"\"content\":[{\"type\":\"text\",\"text\":\"ok\"}],",
"\"model\":\"claude-3-7-sonnet-latest\",",
"\"stop_reason\":\"end_turn\",",
"\"stop_sequence\":null,",
"\"usage\":{\"input_tokens\":1,\"cache_creation_input_tokens\":2,\"cache_read_input_tokens\":3,\"output_tokens\":1}",
"}"
),
&[("request-id", "req_profile_123")],
)],
)
.await;
let sink = Arc::new(MemoryTelemetrySink::default());
let client = AnthropicClient::new("test-key")
.with_base_url(server.base_url())
.with_client_identity(ClientIdentity::new("claude-code", "9.9.9").with_runtime("rust-cli"))
.with_beta("tools-2026-04-01")
.with_extra_body_param("metadata", json!({"source": "clawd-code"}))
.with_session_tracer(SessionTracer::new("session-telemetry", sink.clone()));
let response = client
.send_message(&sample_request(false))
.await
.expect("request should succeed");
assert_eq!(response.request_id.as_deref(), Some("req_profile_123"));
let captured = state.lock().await;
let request = captured.first().expect("server should capture request");
assert_eq!(
request.headers.get("anthropic-beta").map(String::as_str),
Some("claude-code-20250219,prompt-caching-scope-2026-01-05,tools-2026-04-01")
);
assert_eq!(
request.headers.get("user-agent").map(String::as_str),
Some("claude-code/9.9.9")
);
let body: serde_json::Value =
serde_json::from_str(&request.body).expect("request body should be json");
assert_eq!(body["metadata"]["source"], json!("clawd-code"));
assert!(
body.get("betas").is_none(),
"betas must travel via the anthropic-beta header, not the request body"
);
let events = sink.events();
assert_eq!(events.len(), 6);
assert!(matches!(
&events[0],
TelemetryEvent::HttpRequestStarted {
session_id,
attempt: 1,
method,
path,
..
} if session_id == "session-telemetry" && method == "POST" && path == "/v1/messages"
));
assert!(matches!(
&events[1],
TelemetryEvent::SessionTrace(trace) if trace.name == "http_request_started"
));
assert!(matches!(
&events[2],
TelemetryEvent::HttpRequestSucceeded {
request_id,
status: 200,
..
} if request_id.as_deref() == Some("req_profile_123")
));
assert!(matches!(
&events[3],
TelemetryEvent::SessionTrace(trace) if trace.name == "http_request_succeeded"
));
assert!(matches!(
&events[4],
TelemetryEvent::Analytics(event)
if event.namespace == "api"
&& event.action == "message_usage"
&& event.properties.get("request_id") == Some(&json!("req_profile_123"))
&& event.properties.get("total_tokens") == Some(&json!(7))
&& event.properties.get("estimated_cost_usd") == Some(&json!("$0.0001"))
));
assert!(matches!(
&events[5],
TelemetryEvent::SessionTrace(trace) if trace.name == "analytics"
));
}
#[tokio::test]
async fn send_message_parses_prompt_cache_token_usage_from_response() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let body = concat!(
"{",
"\"id\":\"msg_cache_tokens\",",
"\"type\":\"message\",",
"\"role\":\"assistant\",",
"\"content\":[{\"type\":\"text\",\"text\":\"Cache tokens\"}],",
"\"model\":\"claude-3-7-sonnet-latest\",",
"\"stop_reason\":\"end_turn\",",
"\"stop_sequence\":null,",
"\"usage\":{\"input_tokens\":12,\"cache_creation_input_tokens\":321,\"cache_read_input_tokens\":654,\"output_tokens\":4}",
"}"
);
let server = spawn_server(
state,
vec![http_response("200 OK", "application/json", body)],
)
.await;
let client = AnthropicClient::new("test-key").with_base_url(server.base_url());
let response = client
.send_message(&sample_request(false))
.await
.expect("request should succeed");
assert_eq!(response.usage.input_tokens, 12);
assert_eq!(response.usage.cache_creation_input_tokens, 321);
assert_eq!(response.usage.cache_read_input_tokens, 654);
assert_eq!(response.usage.output_tokens, 4);
}
#[tokio::test]
async fn given_empty_usage_object_when_send_message_parses_response_then_usage_defaults_to_zero() {
// given
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let body = concat!(
"{",
"\"id\":\"msg_empty_usage\",",
"\"type\":\"message\",",
"\"role\":\"assistant\",",
"\"content\":[{\"type\":\"text\",\"text\":\"Hello from Claude\"}],",
"\"model\":\"claude-3-7-sonnet-latest\",",
"\"stop_reason\":\"end_turn\",",
"\"stop_sequence\":null,",
"\"usage\":{}",
"}"
);
let server = spawn_server(
state,
vec![http_response("200 OK", "application/json", body)],
)
.await;
let client = AnthropicClient::new("test-key").with_base_url(server.base_url());
// when
let response = client
.send_message(&sample_request(false))
.await
.expect("response with empty usage object should still parse");
// then
assert_eq!(response.id, "msg_empty_usage");
assert_eq!(response.total_tokens(), 0);
assert_eq!(response.usage.input_tokens, 0);
assert_eq!(response.usage.cache_creation_input_tokens, 0);
assert_eq!(response.usage.cache_read_input_tokens, 0);
assert_eq!(response.usage.output_tokens, 0);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn stream_message_parses_sse_events_with_tool_use() {
let _guard = env_lock();
let temp_root = std::env::temp_dir().join(format!(
"api-stream-cache-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("time")
.as_nanos()
));
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let sse = concat!(
"event: message_start\n",
"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_stream\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":8,\"cache_creation_input_tokens\":13,\"cache_read_input_tokens\":21,\"output_tokens\":0}}}\n\n",
"event: content_block_start\n",
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"toolu_123\",\"name\":\"get_weather\",\"input\":{}}}\n\n",
"event: content_block_delta\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"city\\\":\\\"Paris\\\"}\"}}\n\n",
"event: content_block_stop\n",
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
"event: message_delta\n",
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\",\"stop_sequence\":null},\"usage\":{\"input_tokens\":8,\"cache_creation_input_tokens\":34,\"cache_read_input_tokens\":55,\"output_tokens\":1}}\n\n",
"event: message_stop\n",
"data: {\"type\":\"message_stop\"}\n\n",
"data: [DONE]\n\n"
);
let server = spawn_server(
state.clone(),
vec![http_response_with_headers(
"200 OK",
"text/event-stream",
sse,
&[("request-id", "req_stream_456")],
)],
)
.await;
let client = ApiClient::new("test-key")
.with_auth_token(Some("proxy-token".to_string()))
.with_base_url(server.base_url())
.with_prompt_cache(PromptCache::new("stream-session"));
let mut stream = client
.stream_message(&sample_request(false))
.await
.expect("stream should start");
assert_eq!(stream.request_id(), Some("req_stream_456"));
let mut events = Vec::new();
while let Some(event) = stream
.next_event()
.await
.expect("stream event should parse")
{
events.push(event);
}
assert_eq!(events.len(), 6);
assert!(matches!(events[0], StreamEvent::MessageStart(_)));
assert!(matches!(
events[1],
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
content_block: OutputContentBlock::ToolUse { .. },
..
})
));
assert!(matches!(
events[2],
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
delta: ContentBlockDelta::InputJsonDelta { .. },
..
})
));
assert!(matches!(events[3], StreamEvent::ContentBlockStop(_)));
assert!(matches!(
events[4],
StreamEvent::MessageDelta(MessageDeltaEvent { .. })
));
assert!(matches!(events[5], StreamEvent::MessageStop(_)));
match &events[1] {
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
content_block: OutputContentBlock::ToolUse { name, input, .. },
..
}) => {
assert_eq!(name, "get_weather");
assert_eq!(input, &json!({}));
}
other => panic!("expected tool_use block, got {other:?}"),
}
let captured = state.lock().await;
let request = captured.first().expect("server should capture request");
assert!(request.body.contains("\"stream\":true"));
let cache_stats = client
.prompt_cache_stats()
.expect("prompt cache stats should exist");
assert_eq!(cache_stats.tracked_requests, 1);
assert_eq!(cache_stats.last_cache_creation_input_tokens, Some(34));
assert_eq!(cache_stats.last_cache_read_input_tokens, Some(55));
assert_eq!(
cache_stats.last_cache_source.as_deref(),
Some("api-response")
);
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
std::env::remove_var("CLAUDE_CONFIG_HOME");
}
#[tokio::test]
async fn retries_retryable_failures_before_succeeding() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![
http_response(
"429 Too Many Requests",
"application/json",
"{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down\"}}",
),
http_response(
"200 OK",
"application/json",
"{\"id\":\"msg_retry\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Recovered\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}",
),
],
)
.await;
let client = ApiClient::new("test-key")
.with_base_url(server.base_url())
.with_retry_policy(2, Duration::from_millis(1), Duration::from_millis(2));
let response = client
.send_message(&sample_request(false))
.await
.expect("retry should eventually succeed");
assert_eq!(response.total_tokens(), 5);
assert_eq!(state.lock().await.len(), 2);
}
#[tokio::test]
async fn provider_client_dispatches_anthropic_requests() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![http_response(
"200 OK",
"application/json",
"{\"id\":\"msg_provider\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Dispatched\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}",
)],
)
.await;
let client = ProviderClient::from_model_with_anthropic_auth(
"claude-sonnet-4-6",
Some(AuthSource::ApiKey("test-key".to_string())),
)
.expect("anthropic provider client should be constructed");
let client = match client {
ProviderClient::Anthropic(client) => {
ProviderClient::Anthropic(client.with_base_url(server.base_url()))
}
other => panic!("expected anthropic provider, got {other:?}"),
};
let response = client
.send_message(&sample_request(false))
.await
.expect("provider-dispatched request should succeed");
assert_eq!(response.total_tokens(), 5);
let captured = state.lock().await;
let request = captured.first().expect("server should capture request");
assert_eq!(request.path, "/v1/messages");
assert_eq!(
request.headers.get("x-api-key").map(String::as_str),
Some("test-key")
);
}
#[tokio::test]
async fn surfaces_retry_exhaustion_for_persistent_retryable_errors() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![
http_response(
"503 Service Unavailable",
"application/json",
"{\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}",
),
http_response(
"503 Service Unavailable",
"application/json",
"{\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"still busy\"}}",
),
],
)
.await;
let client = ApiClient::new("test-key")
.with_base_url(server.base_url())
.with_retry_policy(1, Duration::from_millis(1), Duration::from_millis(2));
let error = client
.send_message(&sample_request(false))
.await
.expect_err("persistent 503 should fail");
match error {
ApiError::RetriesExhausted {
attempts,
last_error,
} => {
assert_eq!(attempts, 2);
assert!(matches!(
*last_error,
ApiError::Api {
status: reqwest::StatusCode::SERVICE_UNAVAILABLE,
retryable: true,
..
}
));
}
other => panic!("expected retries exhausted, got {other:?}"),
}
}
#[tokio::test]
async fn retries_multiple_retryable_failures_with_exponential_backoff_and_jitter() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![
http_response(
"429 Too Many Requests",
"application/json",
"{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down\"}}",
),
http_response(
"500 Internal Server Error",
"application/json",
"{\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"boom\"}}",
),
http_response(
"503 Service Unavailable",
"application/json",
"{\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}",
),
http_response(
"429 Too Many Requests",
"application/json",
"{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down again\"}}",
),
http_response(
"503 Service Unavailable",
"application/json",
"{\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"still busy\"}}",
),
http_response(
"200 OK",
"application/json",
"{\"id\":\"msg_exp_retry\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Recovered after 5\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}",
),
],
)
.await;
let client = ApiClient::new("test-key")
.with_base_url(server.base_url())
.with_retry_policy(8, Duration::from_millis(1), Duration::from_millis(4));
let started_at = std::time::Instant::now();
let response = client
.send_message(&sample_request(false))
.await
.expect("8-retry policy should absorb 5 retryable failures");
let elapsed = started_at.elapsed();
assert_eq!(response.total_tokens(), 5);
assert_eq!(
state.lock().await.len(),
6,
"client should issue 1 original + 5 retry requests before the 200"
);
// Jittered sleeps are bounded by 2 * max_backoff per retry (base + jitter),
// so 5 sleeps fit comfortably below this upper bound with generous slack.
assert!(
elapsed < Duration::from_secs(5),
"retries should complete promptly, took {elapsed:?}"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn send_message_reuses_recent_completion_cache_entries() {
let _guard = env_lock();
let temp_root = std::env::temp_dir().join(format!(
"api-prompt-cache-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("time")
.as_nanos()
));
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![http_response(
"200 OK",
"application/json",
"{\"id\":\"msg_cached\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Cached once\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"cache_creation_input_tokens\":5,\"cache_read_input_tokens\":4000,\"output_tokens\":2}}",
)],
)
.await;
let client = AnthropicClient::new("test-key")
.with_base_url(server.base_url())
.with_prompt_cache(PromptCache::new("integration-session"));
let first = client
.send_message(&sample_request(false))
.await
.expect("first request should succeed");
let second = client
.send_message(&sample_request(false))
.await
.expect("second request should reuse cache");
assert_eq!(first.content, second.content);
assert_eq!(state.lock().await.len(), 1);
let cache_stats = client
.prompt_cache_stats()
.expect("prompt cache stats should exist");
assert_eq!(cache_stats.completion_cache_hits, 1);
assert_eq!(cache_stats.completion_cache_misses, 1);
assert_eq!(cache_stats.completion_cache_writes, 1);
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
std::env::remove_var("CLAUDE_CONFIG_HOME");
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn send_message_tracks_unexpected_prompt_cache_breaks() {
let _guard = env_lock();
let temp_root = std::env::temp_dir().join(format!(
"api-prompt-break-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("time")
.as_nanos()
));
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state,
vec![
http_response(
"200 OK",
"application/json",
"{\"id\":\"msg_one\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"One\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"cache_creation_input_tokens\":5,\"cache_read_input_tokens\":6000,\"output_tokens\":2}}",
),
http_response(
"200 OK",
"application/json",
"{\"id\":\"msg_two\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Two\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":1000,\"output_tokens\":2}}",
),
],
)
.await;
let request = sample_request(false);
let client = AnthropicClient::new("test-key")
.with_base_url(server.base_url())
.with_prompt_cache(PromptCache::with_config(PromptCacheConfig {
session_id: "break-session".to_string(),
completion_ttl: Duration::from_secs(0),
..PromptCacheConfig::default()
}));
client
.send_message(&request)
.await
.expect("first response should succeed");
client
.send_message(&request)
.await
.expect("second response should succeed");
let cache_stats = client
.prompt_cache_stats()
.expect("prompt cache stats should exist");
assert_eq!(cache_stats.unexpected_cache_breaks, 1);
assert_eq!(
cache_stats.last_break_reason.as_deref(),
Some("cache read tokens dropped while prompt fingerprint remained stable")
);
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
std::env::remove_var("CLAUDE_CONFIG_HOME");
}
#[tokio::test]
#[ignore = "requires ANTHROPIC_API_KEY and network access"]
async fn live_stream_smoke_test() {
let client = ApiClient::from_env().expect("ANTHROPIC_API_KEY must be set");
let mut stream = client
.stream_message(&MessageRequest {
model: std::env::var("ANTHROPIC_MODEL")
.unwrap_or_else(|_| "claude-3-7-sonnet-latest".to_string()),
max_tokens: 32,
messages: vec![InputMessage::user_text(
"Reply with exactly: hello from rust",
)],
system: None,
tools: None,
tool_choice: None,
stream: false,
..Default::default()
})
.await
.expect("live stream should start");
while let Some(_event) = stream
.next_event()
.await
.expect("live stream should yield events")
{}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct CapturedRequest {
method: String,
path: String,
headers: HashMap<String, String>,
body: String,
}
struct TestServer {
base_url: String,
join_handle: tokio::task::JoinHandle<()>,
}
impl TestServer {
fn base_url(&self) -> String {
self.base_url.clone()
}
}
impl Drop for TestServer {
fn drop(&mut self) {
self.join_handle.abort();
}
}
async fn spawn_server(
state: Arc<Mutex<Vec<CapturedRequest>>>,
responses: Vec<String>,
) -> TestServer {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let address = listener
.local_addr()
.expect("listener should have local addr");
let join_handle = tokio::spawn(async move {
for response in responses {
let (mut socket, _) = listener.accept().await.expect("server should accept");
let mut buffer = Vec::new();
let mut header_end = None;
loop {
let mut chunk = [0_u8; 1024];
let read = socket
.read(&mut chunk)
.await
.expect("request read should succeed");
if read == 0 {
break;
}
buffer.extend_from_slice(&chunk[..read]);
if let Some(position) = find_header_end(&buffer) {
header_end = Some(position);
break;
}
}
let header_end = header_end.expect("request should include headers");
let (header_bytes, remaining) = buffer.split_at(header_end);
let header_text =
String::from_utf8(header_bytes.to_vec()).expect("headers should be utf8");
let mut lines = header_text.split("\r\n");
let request_line = lines.next().expect("request line should exist");
let mut parts = request_line.split_whitespace();
let method = parts.next().expect("method should exist").to_string();
let path = parts.next().expect("path should exist").to_string();
let mut headers = HashMap::new();
let mut content_length = 0_usize;
for line in lines {
if line.is_empty() {
continue;
}
let (name, value) = line.split_once(':').expect("header should have colon");
let value = value.trim().to_string();
if name.eq_ignore_ascii_case("content-length") {
content_length = value.parse().expect("content length should parse");
}
headers.insert(name.to_ascii_lowercase(), value);
}
let mut body = remaining[4..].to_vec();
while body.len() < content_length {
let mut chunk = vec![0_u8; content_length - body.len()];
let read = socket
.read(&mut chunk)
.await
.expect("body read should succeed");
if read == 0 {
break;
}
body.extend_from_slice(&chunk[..read]);
}
state.lock().await.push(CapturedRequest {
method,
path,
headers,
body: String::from_utf8(body).expect("body should be utf8"),
});
socket
.write_all(response.as_bytes())
.await
.expect("response write should succeed");
}
});
TestServer {
base_url: format!("http://{address}"),
join_handle,
}
}
fn find_header_end(bytes: &[u8]) -> Option<usize> {
bytes.windows(4).position(|window| window == b"\r\n\r\n")
}
fn http_response(status: &str, content_type: &str, body: &str) -> String {
http_response_with_headers(status, content_type, body, &[])
}
fn http_response_with_headers(
status: &str,
content_type: &str,
body: &str,
headers: &[(&str, &str)],
) -> String {
let mut extra_headers = String::new();
for (name, value) in headers {
use std::fmt::Write as _;
write!(&mut extra_headers, "{name}: {value}\r\n").expect("header write should succeed");
}
format!(
"HTTP/1.1 {status}\r\ncontent-type: {content_type}\r\n{extra_headers}content-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
)
}
fn sample_request(stream: bool) -> MessageRequest {
MessageRequest {
model: "claude-3-7-sonnet-latest".to_string(),
max_tokens: 64,
messages: vec![InputMessage {
role: "user".to_string(),
content: vec![
InputContentBlock::Text {
text: "Say hello".to_string(),
},
InputContentBlock::ToolResult {
tool_use_id: "toolu_prev".to_string(),
content: vec![api::ToolResultContentBlock::Json {
value: json!({"forecast": "sunny"}),
}],
is_error: false,
},
],
}],
system: Some("Use tools when needed".to_string()),
tools: Some(vec![ToolDefinition {
name: "get_weather".to_string(),
description: Some("Fetches the weather".to_string()),
input_schema: json!({
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"]
}),
}]),
tool_choice: Some(ToolChoice::Auto),
stream,
..Default::default()
}
}
@@ -0,0 +1,531 @@
use std::collections::HashMap;
use std::ffi::OsString;
use std::sync::Arc;
use std::sync::{Mutex as StdMutex, OnceLock};
use api::{
ApiError, ContentBlockDelta, ContentBlockDeltaEvent, ContentBlockStartEvent,
ContentBlockStopEvent, InputContentBlock, InputMessage, MessageDeltaEvent, MessageRequest,
OpenAiCompatClient, OpenAiCompatConfig, OutputContentBlock, ProviderClient, StreamEvent,
ToolChoice, ToolDefinition,
};
use serde_json::json;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::Mutex;
#[tokio::test]
async fn send_message_uses_openai_compatible_endpoint_and_auth() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let body = concat!(
"{",
"\"id\":\"chatcmpl_test\",",
"\"model\":\"grok-3\",",
"\"choices\":[{",
"\"message\":{\"role\":\"assistant\",\"content\":\"Hello from Grok\",\"tool_calls\":[]},",
"\"finish_reason\":\"stop\"",
"}],",
"\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":5}",
"}"
);
let server = spawn_server(
state.clone(),
vec![http_response("200 OK", "application/json", body)],
)
.await;
let client = OpenAiCompatClient::new("xai-test-key", OpenAiCompatConfig::xai())
.with_base_url(server.base_url());
let response = client
.send_message(&sample_request(false))
.await
.expect("request should succeed");
assert_eq!(response.model, "grok-3");
assert_eq!(response.total_tokens(), 16);
assert_eq!(
response.content,
vec![OutputContentBlock::Text {
text: "Hello from Grok".to_string(),
}]
);
let captured = state.lock().await;
let request = captured.first().expect("server should capture request");
assert_eq!(request.path, "/chat/completions");
assert_eq!(
request.headers.get("authorization").map(String::as_str),
Some("Bearer xai-test-key")
);
let body: serde_json::Value = serde_json::from_str(&request.body).expect("json body");
assert_eq!(body["model"], json!("grok-3"));
assert_eq!(body["messages"][0]["role"], json!("system"));
assert_eq!(body["tools"][0]["type"], json!("function"));
}
#[tokio::test]
async fn send_message_blocks_oversized_xai_requests_before_the_http_call() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![http_response("200 OK", "application/json", "{}")],
)
.await;
let client = OpenAiCompatClient::new("xai-test-key", OpenAiCompatConfig::xai())
.with_base_url(server.base_url());
let error = client
.send_message(&MessageRequest {
model: "grok-3".to_string(),
max_tokens: 64_000,
messages: vec![InputMessage {
role: "user".to_string(),
content: vec![InputContentBlock::Text {
text: "x".repeat(300_000),
}],
}],
system: Some("Keep the answer short.".to_string()),
tools: None,
tool_choice: None,
stream: false,
..Default::default()
})
.await
.expect_err("oversized request should fail local context-window preflight");
assert!(matches!(error, ApiError::ContextWindowExceeded { .. }));
assert!(
state.lock().await.is_empty(),
"preflight failure should avoid any upstream HTTP request"
);
}
#[tokio::test]
async fn send_message_accepts_full_chat_completions_endpoint_override() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let body = concat!(
"{",
"\"id\":\"chatcmpl_full_endpoint\",",
"\"model\":\"grok-3\",",
"\"choices\":[{",
"\"message\":{\"role\":\"assistant\",\"content\":\"Endpoint override works\",\"tool_calls\":[]},",
"\"finish_reason\":\"stop\"",
"}],",
"\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3}",
"}"
);
let server = spawn_server(
state.clone(),
vec![http_response("200 OK", "application/json", body)],
)
.await;
let endpoint_url = format!("{}/chat/completions", server.base_url());
let client = OpenAiCompatClient::new("xai-test-key", OpenAiCompatConfig::xai())
.with_base_url(endpoint_url);
let response = client
.send_message(&sample_request(false))
.await
.expect("request should succeed");
assert_eq!(response.total_tokens(), 10);
let captured = state.lock().await;
let request = captured.first().expect("server should capture request");
assert_eq!(request.path, "/chat/completions");
}
#[tokio::test]
async fn stream_message_normalizes_text_and_multiple_tool_calls() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let sse = concat!(
"data: {\"id\":\"chatcmpl_stream\",\"model\":\"grok-3\",\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n",
"data: {\"id\":\"chatcmpl_stream\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"weather\",\"arguments\":\"{\\\"city\\\":\\\"Paris\\\"}\"}},{\"index\":1,\"id\":\"call_2\",\"function\":{\"name\":\"clock\",\"arguments\":\"{\\\"zone\\\":\\\"UTC\\\"}\"}}]}}]}\n\n",
"data: {\"id\":\"chatcmpl_stream\",\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}]}\n\n",
"data: [DONE]\n\n"
);
let server = spawn_server(
state.clone(),
vec![http_response_with_headers(
"200 OK",
"text/event-stream",
sse,
&[("x-request-id", "req_grok_stream")],
)],
)
.await;
let client = OpenAiCompatClient::new("xai-test-key", OpenAiCompatConfig::xai())
.with_base_url(server.base_url());
let mut stream = client
.stream_message(&sample_request(false))
.await
.expect("stream should start");
assert_eq!(stream.request_id(), Some("req_grok_stream"));
let mut events = Vec::new();
while let Some(event) = stream.next_event().await.expect("event should parse") {
events.push(event);
}
assert!(matches!(events[0], StreamEvent::MessageStart(_)));
assert!(matches!(
events[1],
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
content_block: OutputContentBlock::Text { .. },
..
})
));
assert!(matches!(
events[2],
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
delta: ContentBlockDelta::TextDelta { .. },
..
})
));
assert!(matches!(
events[3],
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
index: 1,
content_block: OutputContentBlock::ToolUse { .. },
})
));
assert!(matches!(
events[4],
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
index: 1,
delta: ContentBlockDelta::InputJsonDelta { .. },
})
));
assert!(matches!(
events[5],
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
index: 2,
content_block: OutputContentBlock::ToolUse { .. },
})
));
assert!(matches!(
events[6],
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
index: 2,
delta: ContentBlockDelta::InputJsonDelta { .. },
})
));
assert!(matches!(
events[7],
StreamEvent::ContentBlockStop(ContentBlockStopEvent { index: 1 })
));
assert!(matches!(
events[8],
StreamEvent::ContentBlockStop(ContentBlockStopEvent { index: 2 })
));
assert!(matches!(
events[9],
StreamEvent::ContentBlockStop(ContentBlockStopEvent { index: 0 })
));
assert!(matches!(events[10], StreamEvent::MessageDelta(_)));
assert!(matches!(events[11], StreamEvent::MessageStop(_)));
let captured = state.lock().await;
let request = captured.first().expect("captured request");
assert_eq!(request.path, "/chat/completions");
assert!(request.body.contains("\"stream\":true"));
}
#[allow(clippy::await_holding_lock)]
#[tokio::test]
async fn openai_streaming_requests_opt_into_usage_chunks() {
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let sse = concat!(
"data: {\"id\":\"chatcmpl_openai_stream\",\"model\":\"gpt-5\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n",
"data: {\"id\":\"chatcmpl_openai_stream\",\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
"data: {\"id\":\"chatcmpl_openai_stream\",\"choices\":[],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":4}}\n\n",
"data: [DONE]\n\n"
);
let server = spawn_server(
state.clone(),
vec![http_response_with_headers(
"200 OK",
"text/event-stream",
sse,
&[("x-request-id", "req_openai_stream")],
)],
)
.await;
let client = OpenAiCompatClient::new("openai-test-key", OpenAiCompatConfig::openai())
.with_base_url(server.base_url());
let mut stream = client
.stream_message(&sample_request(false))
.await
.expect("stream should start");
assert_eq!(stream.request_id(), Some("req_openai_stream"));
let mut events = Vec::new();
while let Some(event) = stream.next_event().await.expect("event should parse") {
events.push(event);
}
assert!(matches!(events[0], StreamEvent::MessageStart(_)));
assert!(matches!(
events[1],
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
content_block: OutputContentBlock::Text { .. },
..
})
));
assert!(matches!(
events[2],
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
delta: ContentBlockDelta::TextDelta { .. },
..
})
));
assert!(matches!(
events[3],
StreamEvent::ContentBlockStop(ContentBlockStopEvent { index: 0 })
));
assert!(matches!(
events[4],
StreamEvent::MessageDelta(MessageDeltaEvent { .. })
));
assert!(matches!(events[5], StreamEvent::MessageStop(_)));
match &events[4] {
StreamEvent::MessageDelta(MessageDeltaEvent { usage, .. }) => {
assert_eq!(usage.input_tokens, 9);
assert_eq!(usage.output_tokens, 4);
}
other => panic!("expected message delta, got {other:?}"),
}
let captured = state.lock().await;
let request = captured.first().expect("captured request");
assert_eq!(request.path, "/chat/completions");
let body: serde_json::Value = serde_json::from_str(&request.body).expect("json body");
assert_eq!(body["stream"], json!(true));
assert_eq!(body["stream_options"], json!({"include_usage": true}));
}
#[allow(clippy::await_holding_lock)]
#[tokio::test]
async fn provider_client_dispatches_xai_requests_from_env() {
let _lock = env_lock();
let _api_key = ScopedEnvVar::set("XAI_API_KEY", "xai-test-key");
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let server = spawn_server(
state.clone(),
vec![http_response(
"200 OK",
"application/json",
"{\"id\":\"chatcmpl_provider\",\"model\":\"grok-3\",\"choices\":[{\"message\":{\"role\":\"assistant\",\"content\":\"Through provider client\",\"tool_calls\":[]},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":4}}",
)],
)
.await;
let _base_url = ScopedEnvVar::set("XAI_BASE_URL", server.base_url());
let client =
ProviderClient::from_model("grok").expect("xAI provider client should be constructed");
assert!(matches!(client, ProviderClient::Xai(_)));
let response = client
.send_message(&sample_request(false))
.await
.expect("provider-dispatched request should succeed");
assert_eq!(response.total_tokens(), 13);
let captured = state.lock().await;
let request = captured.first().expect("captured request");
assert_eq!(request.path, "/chat/completions");
assert_eq!(
request.headers.get("authorization").map(String::as_str),
Some("Bearer xai-test-key")
);
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct CapturedRequest {
path: String,
headers: HashMap<String, String>,
body: String,
}
struct TestServer {
base_url: String,
join_handle: tokio::task::JoinHandle<()>,
}
impl TestServer {
fn base_url(&self) -> String {
self.base_url.clone()
}
}
impl Drop for TestServer {
fn drop(&mut self) {
self.join_handle.abort();
}
}
async fn spawn_server(
state: Arc<Mutex<Vec<CapturedRequest>>>,
responses: Vec<String>,
) -> TestServer {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let address = listener.local_addr().expect("listener addr");
let join_handle = tokio::spawn(async move {
for response in responses {
let (mut socket, _) = listener.accept().await.expect("accept");
let mut buffer = Vec::new();
let mut header_end = None;
loop {
let mut chunk = [0_u8; 1024];
let read = socket.read(&mut chunk).await.expect("read request");
if read == 0 {
break;
}
buffer.extend_from_slice(&chunk[..read]);
if let Some(position) = find_header_end(&buffer) {
header_end = Some(position);
break;
}
}
let header_end = header_end.expect("headers should exist");
let (header_bytes, remaining) = buffer.split_at(header_end);
let header_text = String::from_utf8(header_bytes.to_vec()).expect("utf8 headers");
let mut lines = header_text.split("\r\n");
let request_line = lines.next().expect("request line");
let path = request_line
.split_whitespace()
.nth(1)
.expect("path")
.to_string();
let mut headers = HashMap::new();
let mut content_length = 0_usize;
for line in lines {
if line.is_empty() {
continue;
}
let (name, value) = line.split_once(':').expect("header");
let value = value.trim().to_string();
if name.eq_ignore_ascii_case("content-length") {
content_length = value.parse().expect("content length");
}
headers.insert(name.to_ascii_lowercase(), value);
}
let mut body = remaining[4..].to_vec();
while body.len() < content_length {
let mut chunk = vec![0_u8; content_length - body.len()];
let read = socket.read(&mut chunk).await.expect("read body");
if read == 0 {
break;
}
body.extend_from_slice(&chunk[..read]);
}
state.lock().await.push(CapturedRequest {
path,
headers,
body: String::from_utf8(body).expect("utf8 body"),
});
socket
.write_all(response.as_bytes())
.await
.expect("write response");
}
});
TestServer {
base_url: format!("http://{address}"),
join_handle,
}
}
fn find_header_end(bytes: &[u8]) -> Option<usize> {
bytes.windows(4).position(|window| window == b"\r\n\r\n")
}
fn http_response(status: &str, content_type: &str, body: &str) -> String {
http_response_with_headers(status, content_type, body, &[])
}
fn http_response_with_headers(
status: &str,
content_type: &str,
body: &str,
headers: &[(&str, &str)],
) -> String {
let mut extra_headers = String::new();
for (name, value) in headers {
use std::fmt::Write as _;
write!(&mut extra_headers, "{name}: {value}\r\n").expect("header write");
}
format!(
"HTTP/1.1 {status}\r\ncontent-type: {content_type}\r\n{extra_headers}content-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
)
}
fn sample_request(stream: bool) -> MessageRequest {
MessageRequest {
model: "grok-3".to_string(),
max_tokens: 64,
messages: vec![InputMessage {
role: "user".to_string(),
content: vec![InputContentBlock::Text {
text: "Say hello".to_string(),
}],
}],
system: Some("Use tools when needed".to_string()),
tools: Some(vec![ToolDefinition {
name: "weather".to_string(),
description: Some("Fetches weather".to_string()),
input_schema: json!({
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"]
}),
}]),
tool_choice: Some(ToolChoice::Auto),
stream,
..Default::default()
}
}
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
static LOCK: OnceLock<StdMutex<()>> = OnceLock::new();
LOCK.get_or_init(|| StdMutex::new(()))
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
struct ScopedEnvVar {
key: &'static str,
previous: Option<OsString>,
}
impl ScopedEnvVar {
fn set(key: &'static str, value: impl AsRef<std::ffi::OsStr>) -> Self {
let previous = std::env::var_os(key);
std::env::set_var(key, value);
Self { key, previous }
}
}
impl Drop for ScopedEnvVar {
fn drop(&mut self) {
match &self.previous {
Some(value) => std::env::set_var(self.key, value),
None => std::env::remove_var(self.key),
}
}
}
@@ -0,0 +1,88 @@
use std::ffi::OsString;
use std::sync::{Mutex, OnceLock};
use api::{read_xai_base_url, ApiError, AuthSource, ProviderClient, ProviderKind};
#[test]
fn provider_client_routes_grok_aliases_through_xai() {
let _lock = env_lock();
let _xai_api_key = EnvVarGuard::set("XAI_API_KEY", Some("xai-test-key"));
let client = ProviderClient::from_model("grok-mini").expect("grok alias should resolve");
assert_eq!(client.provider_kind(), ProviderKind::Xai);
}
#[test]
fn provider_client_reports_missing_xai_credentials_for_grok_models() {
let _lock = env_lock();
let _xai_api_key = EnvVarGuard::set("XAI_API_KEY", None);
let error = ProviderClient::from_model("grok-3")
.expect_err("grok requests without XAI_API_KEY should fail fast");
match error {
ApiError::MissingCredentials {
provider, env_vars, ..
} => {
assert_eq!(provider, "xAI");
assert_eq!(env_vars, &["XAI_API_KEY"]);
}
other => panic!("expected missing xAI credentials, got {other:?}"),
}
}
#[test]
fn provider_client_uses_explicit_anthropic_auth_without_env_lookup() {
let _lock = env_lock();
let _anthropic_api_key = EnvVarGuard::set("ANTHROPIC_API_KEY", None);
let _anthropic_auth_token = EnvVarGuard::set("ANTHROPIC_AUTH_TOKEN", None);
let client = ProviderClient::from_model_with_anthropic_auth(
"claude-sonnet-4-6",
Some(AuthSource::ApiKey("anthropic-test-key".to_string())),
)
.expect("explicit anthropic auth should avoid env lookup");
assert_eq!(client.provider_kind(), ProviderKind::Anthropic);
}
#[test]
fn read_xai_base_url_prefers_env_override() {
let _lock = env_lock();
let _xai_base_url = EnvVarGuard::set("XAI_BASE_URL", Some("https://example.xai.test/v1"));
assert_eq!(read_xai_base_url(), "https://example.xai.test/v1");
}
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
struct EnvVarGuard {
key: &'static str,
original: Option<OsString>,
}
impl EnvVarGuard {
fn set(key: &'static str, value: Option<&str>) -> Self {
let original = std::env::var_os(key);
match value {
Some(value) => std::env::set_var(key, value),
None => std::env::remove_var(key),
}
Self { key, original }
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
match &self.original {
Some(value) => std::env::set_var(self.key, value),
None => std::env::remove_var(self.key),
}
}
}
@@ -0,0 +1,173 @@
use std::ffi::OsString;
use std::sync::{Mutex, OnceLock};
use api::{build_http_client_with, ProxyConfig};
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
struct EnvVarGuard {
key: &'static str,
original: Option<OsString>,
}
impl EnvVarGuard {
fn set(key: &'static str, value: Option<&str>) -> Self {
let original = std::env::var_os(key);
match value {
Some(value) => std::env::set_var(key, value),
None => std::env::remove_var(key),
}
Self { key, original }
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
match &self.original {
Some(value) => std::env::set_var(self.key, value),
None => std::env::remove_var(self.key),
}
}
}
#[test]
fn proxy_config_from_env_reads_uppercase_proxy_vars() {
// given
let _lock = env_lock();
let _http = EnvVarGuard::set("HTTP_PROXY", Some("http://proxy.corp:3128"));
let _https = EnvVarGuard::set("HTTPS_PROXY", Some("http://secure.corp:3129"));
let _no = EnvVarGuard::set("NO_PROXY", Some("localhost,127.0.0.1"));
let _http_lower = EnvVarGuard::set("http_proxy", None);
let _https_lower = EnvVarGuard::set("https_proxy", None);
let _no_lower = EnvVarGuard::set("no_proxy", None);
// when
let config = ProxyConfig::from_env();
// then
assert_eq!(config.http_proxy.as_deref(), Some("http://proxy.corp:3128"));
assert_eq!(
config.https_proxy.as_deref(),
Some("http://secure.corp:3129")
);
assert_eq!(config.no_proxy.as_deref(), Some("localhost,127.0.0.1"));
assert!(config.proxy_url.is_none());
assert!(!config.is_empty());
}
#[test]
fn proxy_config_from_env_reads_lowercase_proxy_vars() {
// given
let _lock = env_lock();
let _http = EnvVarGuard::set("HTTP_PROXY", None);
let _https = EnvVarGuard::set("HTTPS_PROXY", None);
let _no = EnvVarGuard::set("NO_PROXY", None);
let _http_lower = EnvVarGuard::set("http_proxy", Some("http://lower.corp:3128"));
let _https_lower = EnvVarGuard::set("https_proxy", Some("http://lower-secure.corp:3129"));
let _no_lower = EnvVarGuard::set("no_proxy", Some(".internal"));
// when
let config = ProxyConfig::from_env();
// then
assert_eq!(config.http_proxy.as_deref(), Some("http://lower.corp:3128"));
assert_eq!(
config.https_proxy.as_deref(),
Some("http://lower-secure.corp:3129")
);
assert_eq!(config.no_proxy.as_deref(), Some(".internal"));
assert!(!config.is_empty());
}
#[test]
fn proxy_config_from_env_is_empty_when_no_vars_set() {
// given
let _lock = env_lock();
let _http = EnvVarGuard::set("HTTP_PROXY", None);
let _https = EnvVarGuard::set("HTTPS_PROXY", None);
let _no = EnvVarGuard::set("NO_PROXY", None);
let _http_lower = EnvVarGuard::set("http_proxy", None);
let _https_lower = EnvVarGuard::set("https_proxy", None);
let _no_lower = EnvVarGuard::set("no_proxy", None);
// when
let config = ProxyConfig::from_env();
// then
assert!(config.is_empty());
assert!(config.http_proxy.is_none());
assert!(config.https_proxy.is_none());
assert!(config.no_proxy.is_none());
}
#[test]
fn proxy_config_from_env_treats_empty_values_as_unset() {
// given
let _lock = env_lock();
let _http = EnvVarGuard::set("HTTP_PROXY", Some(""));
let _https = EnvVarGuard::set("HTTPS_PROXY", Some(""));
let _http_lower = EnvVarGuard::set("http_proxy", Some(""));
let _https_lower = EnvVarGuard::set("https_proxy", Some(""));
let _no = EnvVarGuard::set("NO_PROXY", Some(""));
let _no_lower = EnvVarGuard::set("no_proxy", Some(""));
// when
let config = ProxyConfig::from_env();
// then
assert!(config.is_empty());
}
#[test]
fn build_client_with_env_proxy_config_succeeds() {
// given
let _lock = env_lock();
let _http = EnvVarGuard::set("HTTP_PROXY", Some("http://proxy.corp:3128"));
let _https = EnvVarGuard::set("HTTPS_PROXY", Some("http://secure.corp:3129"));
let _no = EnvVarGuard::set("NO_PROXY", Some("localhost"));
let _http_lower = EnvVarGuard::set("http_proxy", None);
let _https_lower = EnvVarGuard::set("https_proxy", None);
let _no_lower = EnvVarGuard::set("no_proxy", None);
let config = ProxyConfig::from_env();
// when
let result = build_http_client_with(&config);
// then
assert!(result.is_ok());
}
#[test]
fn build_client_with_proxy_url_config_succeeds() {
// given
let config = ProxyConfig::from_proxy_url("http://unified.corp:3128");
// when
let result = build_http_client_with(&config);
// then
assert!(result.is_ok());
}
#[test]
fn proxy_config_from_env_prefers_uppercase_over_lowercase() {
// given
let _lock = env_lock();
let _http_upper = EnvVarGuard::set("HTTP_PROXY", Some("http://upper.corp:3128"));
let _http_lower = EnvVarGuard::set("http_proxy", Some("http://lower.corp:3128"));
let _https = EnvVarGuard::set("HTTPS_PROXY", None);
let _https_lower = EnvVarGuard::set("https_proxy", None);
let _no = EnvVarGuard::set("NO_PROXY", None);
let _no_lower = EnvVarGuard::set("no_proxy", None);
// when
let config = ProxyConfig::from_env();
// then
assert_eq!(config.http_proxy.as_deref(), Some("http://upper.corp:3128"));
}
+14
View File
@@ -0,0 +1,14 @@
[package]
name = "commands"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
[lints]
workspace = true
[dependencies]
plugins = { path = "../plugins" }
runtime = { path = "../runtime" }
serde_json.workspace = true
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,14 @@
[package]
name = "compat-harness"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
[dependencies]
commands = { path = "../commands" }
tools = { path = "../tools" }
runtime = { path = "../runtime" }
[lints]
workspace = true
@@ -0,0 +1,357 @@
use std::fs;
use std::path::{Path, PathBuf};
use commands::{CommandManifestEntry, CommandRegistry, CommandSource};
use runtime::{BootstrapPhase, BootstrapPlan};
use tools::{ToolManifestEntry, ToolRegistry, ToolSource};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamPaths {
repo_root: PathBuf,
}
impl UpstreamPaths {
#[must_use]
pub fn from_repo_root(repo_root: impl Into<PathBuf>) -> Self {
Self {
repo_root: repo_root.into(),
}
}
#[must_use]
pub fn from_workspace_dir(workspace_dir: impl AsRef<Path>) -> Self {
let workspace_dir = workspace_dir
.as_ref()
.canonicalize()
.unwrap_or_else(|_| workspace_dir.as_ref().to_path_buf());
let primary_repo_root = workspace_dir
.parent()
.map_or_else(|| PathBuf::from(".."), Path::to_path_buf);
let repo_root = resolve_upstream_repo_root(&primary_repo_root);
Self { repo_root }
}
#[must_use]
pub fn commands_path(&self) -> PathBuf {
self.repo_root.join("src/commands.ts")
}
#[must_use]
pub fn tools_path(&self) -> PathBuf {
self.repo_root.join("src/tools.ts")
}
#[must_use]
pub fn cli_path(&self) -> PathBuf {
self.repo_root.join("src/entrypoints/cli.tsx")
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExtractedManifest {
pub commands: CommandRegistry,
pub tools: ToolRegistry,
pub bootstrap: BootstrapPlan,
}
fn resolve_upstream_repo_root(primary_repo_root: &Path) -> PathBuf {
let candidates = upstream_repo_candidates(primary_repo_root);
candidates
.into_iter()
.find(|candidate| candidate.join("src/commands.ts").is_file())
.unwrap_or_else(|| primary_repo_root.to_path_buf())
}
fn upstream_repo_candidates(primary_repo_root: &Path) -> Vec<PathBuf> {
let mut candidates = vec![primary_repo_root.to_path_buf()];
if let Some(explicit) = std::env::var_os("CLAUDE_CODE_UPSTREAM") {
candidates.push(PathBuf::from(explicit));
}
for ancestor in primary_repo_root.ancestors().take(4) {
candidates.push(ancestor.join("claw-code"));
candidates.push(ancestor.join("clawd-code"));
}
candidates.push(primary_repo_root.join("reference-source").join("claw-code"));
candidates.push(primary_repo_root.join("vendor").join("claw-code"));
let mut deduped = Vec::new();
for candidate in candidates {
if !deduped.iter().any(|seen: &PathBuf| seen == &candidate) {
deduped.push(candidate);
}
}
deduped
}
pub fn extract_manifest(paths: &UpstreamPaths) -> std::io::Result<ExtractedManifest> {
let commands_source = fs::read_to_string(paths.commands_path())?;
let tools_source = fs::read_to_string(paths.tools_path())?;
let cli_source = fs::read_to_string(paths.cli_path())?;
Ok(ExtractedManifest {
commands: extract_commands(&commands_source),
tools: extract_tools(&tools_source),
bootstrap: extract_bootstrap_plan(&cli_source),
})
}
#[must_use]
pub fn extract_commands(source: &str) -> CommandRegistry {
let mut entries = Vec::new();
let mut in_internal_block = false;
for raw_line in source.lines() {
let line = raw_line.trim();
if line.starts_with("export const INTERNAL_ONLY_COMMANDS = [") {
in_internal_block = true;
continue;
}
if in_internal_block {
if line.starts_with(']') {
in_internal_block = false;
continue;
}
if let Some(name) = first_identifier(line) {
entries.push(CommandManifestEntry {
name,
source: CommandSource::InternalOnly,
});
}
continue;
}
if line.starts_with("import ") {
for imported in imported_symbols(line) {
entries.push(CommandManifestEntry {
name: imported,
source: CommandSource::Builtin,
});
}
}
if line.contains("feature('") && line.contains("./commands/") {
if let Some(name) = first_assignment_identifier(line) {
entries.push(CommandManifestEntry {
name,
source: CommandSource::FeatureGated,
});
}
}
}
dedupe_commands(entries)
}
#[must_use]
pub fn extract_tools(source: &str) -> ToolRegistry {
let mut entries = Vec::new();
for raw_line in source.lines() {
let line = raw_line.trim();
if line.starts_with("import ") && line.contains("./tools/") {
for imported in imported_symbols(line) {
if imported.ends_with("Tool") {
entries.push(ToolManifestEntry {
name: imported,
source: ToolSource::Base,
});
}
}
}
if line.contains("feature('") && line.contains("Tool") {
if let Some(name) = first_assignment_identifier(line) {
if name.ends_with("Tool") || name.ends_with("Tools") {
entries.push(ToolManifestEntry {
name,
source: ToolSource::Conditional,
});
}
}
}
}
dedupe_tools(entries)
}
#[must_use]
pub fn extract_bootstrap_plan(source: &str) -> BootstrapPlan {
let mut phases = vec![BootstrapPhase::CliEntry];
if source.contains("--version") {
phases.push(BootstrapPhase::FastPathVersion);
}
if source.contains("startupProfiler") {
phases.push(BootstrapPhase::StartupProfiler);
}
if source.contains("--dump-system-prompt") {
phases.push(BootstrapPhase::SystemPromptFastPath);
}
if source.contains("--claude-in-chrome-mcp") {
phases.push(BootstrapPhase::ChromeMcpFastPath);
}
if source.contains("--daemon-worker") {
phases.push(BootstrapPhase::DaemonWorkerFastPath);
}
if source.contains("remote-control") {
phases.push(BootstrapPhase::BridgeFastPath);
}
if source.contains("args[0] === 'daemon'") {
phases.push(BootstrapPhase::DaemonFastPath);
}
if source.contains("args[0] === 'ps'") || source.contains("args.includes('--bg')") {
phases.push(BootstrapPhase::BackgroundSessionFastPath);
}
if source.contains("args[0] === 'new' || args[0] === 'list' || args[0] === 'reply'") {
phases.push(BootstrapPhase::TemplateFastPath);
}
if source.contains("environment-runner") {
phases.push(BootstrapPhase::EnvironmentRunnerFastPath);
}
phases.push(BootstrapPhase::MainRuntime);
BootstrapPlan::from_phases(phases)
}
fn imported_symbols(line: &str) -> Vec<String> {
let Some(after_import) = line.strip_prefix("import ") else {
return Vec::new();
};
let before_from = after_import
.split(" from ")
.next()
.unwrap_or_default()
.trim();
if before_from.starts_with('{') {
return before_from
.trim_matches(|c| c == '{' || c == '}')
.split(',')
.filter_map(|part| {
let trimmed = part.trim();
if trimmed.is_empty() {
return None;
}
Some(trimmed.split_whitespace().next()?.to_string())
})
.collect();
}
let first = before_from.split(',').next().unwrap_or_default().trim();
if first.is_empty() {
Vec::new()
} else {
vec![first.to_string()]
}
}
fn first_assignment_identifier(line: &str) -> Option<String> {
let trimmed = line.trim_start();
let candidate = trimmed.split('=').next()?.trim();
first_identifier(candidate)
}
fn first_identifier(line: &str) -> Option<String> {
let mut out = String::new();
for ch in line.chars() {
if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' {
out.push(ch);
} else if !out.is_empty() {
break;
}
}
(!out.is_empty()).then_some(out)
}
fn dedupe_commands(entries: Vec<CommandManifestEntry>) -> CommandRegistry {
let mut deduped = Vec::new();
for entry in entries {
let exists = deduped.iter().any(|seen: &CommandManifestEntry| {
seen.name == entry.name && seen.source == entry.source
});
if !exists {
deduped.push(entry);
}
}
CommandRegistry::new(deduped)
}
fn dedupe_tools(entries: Vec<ToolManifestEntry>) -> ToolRegistry {
let mut deduped = Vec::new();
for entry in entries {
let exists = deduped
.iter()
.any(|seen: &ToolManifestEntry| seen.name == entry.name && seen.source == entry.source);
if !exists {
deduped.push(entry);
}
}
ToolRegistry::new(deduped)
}
#[cfg(test)]
mod tests {
use super::*;
fn fixture_paths() -> UpstreamPaths {
let workspace_dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("../..");
UpstreamPaths::from_workspace_dir(workspace_dir)
}
fn has_upstream_fixture(paths: &UpstreamPaths) -> bool {
paths.commands_path().is_file()
&& paths.tools_path().is_file()
&& paths.cli_path().is_file()
}
#[test]
fn extracts_non_empty_manifests_from_upstream_repo() {
let paths = fixture_paths();
if !has_upstream_fixture(&paths) {
return;
}
let manifest = extract_manifest(&paths).expect("manifest should load");
assert!(!manifest.commands.entries().is_empty());
assert!(!manifest.tools.entries().is_empty());
assert!(!manifest.bootstrap.phases().is_empty());
}
#[test]
fn detects_known_upstream_command_symbols() {
let paths = fixture_paths();
if !paths.commands_path().is_file() {
return;
}
let commands =
extract_commands(&fs::read_to_string(paths.commands_path()).expect("commands.ts"));
let names: Vec<_> = commands
.entries()
.iter()
.map(|entry| entry.name.as_str())
.collect();
assert!(names.contains(&"addDir"));
assert!(names.contains(&"review"));
assert!(!names.contains(&"INTERNAL_ONLY_COMMANDS"));
}
#[test]
fn detects_known_upstream_tool_symbols() {
let paths = fixture_paths();
if !paths.tools_path().is_file() {
return;
}
let tools = extract_tools(&fs::read_to_string(paths.tools_path()).expect("tools.ts"));
let names: Vec<_> = tools
.entries()
.iter()
.map(|entry| entry.name.as_str())
.collect();
assert!(names.contains(&"AgentTool"));
assert!(names.contains(&"BashTool"));
}
}
@@ -0,0 +1,18 @@
[package]
name = "mock-anthropic-service"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
[[bin]]
name = "mock-anthropic-service"
path = "src/main.rs"
[dependencies]
api = { path = "../api" }
serde_json.workspace = true
tokio = { version = "1", features = ["io-util", "macros", "net", "rt-multi-thread", "signal", "sync"] }
[lints]
workspace = true
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,34 @@
use std::env;
use mock_anthropic_service::MockAnthropicService;
#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut bind_addr = String::from("127.0.0.1:0");
let mut args = env::args().skip(1);
while let Some(arg) = args.next() {
match arg.as_str() {
"--bind" => {
bind_addr = args
.next()
.ok_or_else(|| "missing value for --bind".to_string())?;
}
flag if flag.starts_with("--bind=") => {
bind_addr = flag[7..].to_string();
}
"--help" | "-h" => {
println!("Usage: mock-anthropic-service [--bind HOST:PORT]");
return Ok(());
}
other => {
return Err(format!("unsupported argument: {other}").into());
}
}
}
let server = MockAnthropicService::spawn_on(&bind_addr).await?;
println!("MOCK_ANTHROPIC_BASE_URL={}", server.base_url());
tokio::signal::ctrl_c().await?;
drop(server);
Ok(())
}
+13
View File
@@ -0,0 +1,13 @@
[package]
name = "plugins"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
[dependencies]
serde = { version = "1", features = ["derive"] }
serde_json.workspace = true
[lints]
workspace = true
@@ -0,0 +1,10 @@
{
"name": "example-bundled",
"version": "0.1.0",
"description": "Example bundled plugin scaffold for the Rust plugin system",
"defaultEnabled": false,
"hooks": {
"PreToolUse": ["./hooks/pre.sh"],
"PostToolUse": ["./hooks/post.sh"]
}
}
@@ -0,0 +1,2 @@
#!/bin/sh
printf '%s\n' 'example bundled post hook'
@@ -0,0 +1,2 @@
#!/bin/sh
printf '%s\n' 'example bundled pre hook'
@@ -0,0 +1,10 @@
{
"name": "sample-hooks",
"version": "0.1.0",
"description": "Bundled sample plugin scaffold for hook integration tests.",
"defaultEnabled": false,
"hooks": {
"PreToolUse": ["./hooks/pre.sh"],
"PostToolUse": ["./hooks/post.sh"]
}
}
@@ -0,0 +1,2 @@
#!/bin/sh
printf 'sample bundled post hook'
@@ -0,0 +1,2 @@
#!/bin/sh
printf 'sample bundled pre hook'
+564
View File
@@ -0,0 +1,564 @@
use std::ffi::OsStr;
use std::path::Path;
use std::process::Command;
use serde_json::json;
use crate::{PluginError, PluginHooks, PluginRegistry};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HookEvent {
PreToolUse,
PostToolUse,
PostToolUseFailure,
}
impl HookEvent {
fn as_str(self) -> &'static str {
match self {
Self::PreToolUse => "PreToolUse",
Self::PostToolUse => "PostToolUse",
Self::PostToolUseFailure => "PostToolUseFailure",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HookRunResult {
denied: bool,
failed: bool,
messages: Vec<String>,
}
impl HookRunResult {
#[must_use]
pub fn allow(messages: Vec<String>) -> Self {
Self {
denied: false,
failed: false,
messages,
}
}
#[must_use]
pub fn is_denied(&self) -> bool {
self.denied
}
#[must_use]
pub fn is_failed(&self) -> bool {
self.failed
}
#[must_use]
pub fn messages(&self) -> &[String] {
&self.messages
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct HookRunner {
hooks: PluginHooks,
}
impl HookRunner {
#[must_use]
pub fn new(hooks: PluginHooks) -> Self {
Self { hooks }
}
pub fn from_registry(plugin_registry: &PluginRegistry) -> Result<Self, PluginError> {
Ok(Self::new(plugin_registry.aggregated_hooks()?))
}
#[must_use]
pub fn run_pre_tool_use(&self, tool_name: &str, tool_input: &str) -> HookRunResult {
Self::run_commands(
HookEvent::PreToolUse,
&self.hooks.pre_tool_use,
tool_name,
tool_input,
None,
false,
)
}
#[must_use]
pub fn run_post_tool_use(
&self,
tool_name: &str,
tool_input: &str,
tool_output: &str,
is_error: bool,
) -> HookRunResult {
Self::run_commands(
HookEvent::PostToolUse,
&self.hooks.post_tool_use,
tool_name,
tool_input,
Some(tool_output),
is_error,
)
}
#[must_use]
pub fn run_post_tool_use_failure(
&self,
tool_name: &str,
tool_input: &str,
tool_error: &str,
) -> HookRunResult {
Self::run_commands(
HookEvent::PostToolUseFailure,
&self.hooks.post_tool_use_failure,
tool_name,
tool_input,
Some(tool_error),
true,
)
}
fn run_commands(
event: HookEvent,
commands: &[String],
tool_name: &str,
tool_input: &str,
tool_output: Option<&str>,
is_error: bool,
) -> HookRunResult {
if commands.is_empty() {
return HookRunResult::allow(Vec::new());
}
let payload = hook_payload(event, tool_name, tool_input, tool_output, is_error).to_string();
let mut messages = Vec::new();
for command in commands {
match Self::run_command(
command,
event,
tool_name,
tool_input,
tool_output,
is_error,
&payload,
) {
HookCommandOutcome::Allow { message } => {
if let Some(message) = message {
messages.push(message);
}
}
HookCommandOutcome::Deny { message } => {
messages.push(message.unwrap_or_else(|| {
format!("{} hook denied tool `{tool_name}`", event.as_str())
}));
return HookRunResult {
denied: true,
failed: false,
messages,
};
}
HookCommandOutcome::Failed { message } => {
messages.push(message);
return HookRunResult {
denied: false,
failed: true,
messages,
};
}
}
}
HookRunResult::allow(messages)
}
#[allow(clippy::too_many_arguments)]
fn run_command(
command: &str,
event: HookEvent,
tool_name: &str,
tool_input: &str,
tool_output: Option<&str>,
is_error: bool,
payload: &str,
) -> HookCommandOutcome {
let mut child = shell_command(command);
child.stdin(std::process::Stdio::piped());
child.stdout(std::process::Stdio::piped());
child.stderr(std::process::Stdio::piped());
child.env("HOOK_EVENT", event.as_str());
child.env("HOOK_TOOL_NAME", tool_name);
child.env("HOOK_TOOL_INPUT", tool_input);
child.env("HOOK_TOOL_IS_ERROR", if is_error { "1" } else { "0" });
if let Some(tool_output) = tool_output {
child.env("HOOK_TOOL_OUTPUT", tool_output);
}
match child.output_with_stdin(payload.as_bytes()) {
Ok(output) => {
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
let message = (!stdout.is_empty()).then_some(stdout);
match output.status.code() {
Some(0) => HookCommandOutcome::Allow { message },
Some(2) => HookCommandOutcome::Deny { message },
Some(code) => HookCommandOutcome::Failed {
message: format_hook_warning(
command,
code,
message.as_deref(),
stderr.as_str(),
),
},
None => HookCommandOutcome::Failed {
message: format!(
"{} hook `{command}` terminated by signal while handling `{tool_name}`",
event.as_str()
),
},
}
}
Err(error) => HookCommandOutcome::Failed {
message: format!(
"{} hook `{command}` failed to start for `{tool_name}`: {error}",
event.as_str()
),
},
}
}
}
enum HookCommandOutcome {
Allow { message: Option<String> },
Deny { message: Option<String> },
Failed { message: String },
}
fn hook_payload(
event: HookEvent,
tool_name: &str,
tool_input: &str,
tool_output: Option<&str>,
is_error: bool,
) -> serde_json::Value {
match event {
HookEvent::PostToolUseFailure => json!({
"hook_event_name": event.as_str(),
"tool_name": tool_name,
"tool_input": parse_tool_input(tool_input),
"tool_input_json": tool_input,
"tool_error": tool_output,
"tool_result_is_error": true,
}),
_ => json!({
"hook_event_name": event.as_str(),
"tool_name": tool_name,
"tool_input": parse_tool_input(tool_input),
"tool_input_json": tool_input,
"tool_output": tool_output,
"tool_result_is_error": is_error,
}),
}
}
fn parse_tool_input(tool_input: &str) -> serde_json::Value {
serde_json::from_str(tool_input).unwrap_or_else(|_| json!({ "raw": tool_input }))
}
fn format_hook_warning(command: &str, code: i32, stdout: Option<&str>, stderr: &str) -> String {
let mut message = format!("Hook `{command}` exited with status {code}");
if let Some(stdout) = stdout.filter(|stdout| !stdout.is_empty()) {
message.push_str(": ");
message.push_str(stdout);
} else if !stderr.is_empty() {
message.push_str(": ");
message.push_str(stderr);
}
message
}
fn shell_command(command: &str) -> CommandWithStdin {
#[cfg(windows)]
let command_builder = {
let mut command_builder = Command::new("cmd");
command_builder.arg("/C").arg(command);
CommandWithStdin::new(command_builder)
};
#[cfg(not(windows))]
let command_builder = if Path::new(command).exists() {
let mut command_builder = Command::new("sh");
command_builder.arg(command);
CommandWithStdin::new(command_builder)
} else {
let mut command_builder = Command::new("sh");
command_builder.arg("-lc").arg(command);
CommandWithStdin::new(command_builder)
};
command_builder
}
struct CommandWithStdin {
command: Command,
}
impl CommandWithStdin {
fn new(command: Command) -> Self {
Self { command }
}
fn stdin(&mut self, cfg: std::process::Stdio) -> &mut Self {
self.command.stdin(cfg);
self
}
fn stdout(&mut self, cfg: std::process::Stdio) -> &mut Self {
self.command.stdout(cfg);
self
}
fn stderr(&mut self, cfg: std::process::Stdio) -> &mut Self {
self.command.stderr(cfg);
self
}
fn env<K, V>(&mut self, key: K, value: V) -> &mut Self
where
K: AsRef<OsStr>,
V: AsRef<OsStr>,
{
self.command.env(key, value);
self
}
fn output_with_stdin(&mut self, stdin: &[u8]) -> std::io::Result<std::process::Output> {
let mut child = self.command.spawn()?;
if let Some(mut child_stdin) = child.stdin.take() {
use std::io::Write as _;
// Tolerate BrokenPipe: a hook script that runs to completion
// (or exits early without reading stdin) closes its stdin
// before the parent finishes writing the JSON payload, and
// the kernel raises EPIPE on the parent's write_all. That is
// not a hook failure — the child still exited cleanly and we
// still need to wait_with_output() to capture stdout/stderr
// and the real exit code. Other write errors (e.g. EIO,
// permission, OOM) still propagate.
//
// This was the root cause of the Linux CI flake on
// hooks::tests::collects_and_runs_hooks_from_enabled_plugins
// (ROADMAP #25, runs 24120271422 / 24120538408 / 24121392171
// / 24121776826): the test hook scripts run in microseconds
// and the parent's stdin write races against child exit.
// macOS pipes happen to buffer the small payload before the
// child exits; Linux pipes do not, so the race shows up
// deterministically on ubuntu runners.
match child_stdin.write_all(stdin) {
Ok(()) => {}
Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => {}
Err(error) => return Err(error),
}
}
child.wait_with_output()
}
}
#[cfg(test)]
mod tests {
use super::{HookRunResult, HookRunner};
use crate::{PluginManager, PluginManagerConfig};
use std::fs;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};
fn temp_dir(label: &str) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("time should be after epoch")
.as_nanos();
std::env::temp_dir().join(format!("plugins-hook-runner-{label}-{nanos}"))
}
fn make_executable(path: &Path) {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let perms = fs::Permissions::from_mode(0o755);
fs::set_permissions(path, perms)
.unwrap_or_else(|e| panic!("chmod +x {}: {e}", path.display()));
}
#[cfg(not(unix))]
let _ = path;
}
fn write_hook_plugin(
root: &Path,
name: &str,
pre_message: &str,
post_message: &str,
failure_message: &str,
) {
fs::create_dir_all(root.join(".claude-plugin")).expect("manifest dir");
fs::create_dir_all(root.join("hooks")).expect("hooks dir");
let pre_path = root.join("hooks").join("pre.sh");
fs::write(
&pre_path,
format!("#!/bin/sh\nprintf '%s\\n' '{pre_message}'\n"),
)
.expect("write pre hook");
make_executable(&pre_path);
let post_path = root.join("hooks").join("post.sh");
fs::write(
&post_path,
format!("#!/bin/sh\nprintf '%s\\n' '{post_message}'\n"),
)
.expect("write post hook");
make_executable(&post_path);
let failure_path = root.join("hooks").join("failure.sh");
fs::write(
&failure_path,
format!("#!/bin/sh\nprintf '%s\\n' '{failure_message}'\n"),
)
.expect("write failure hook");
make_executable(&failure_path);
fs::write(
root.join(".claude-plugin").join("plugin.json"),
format!(
"{{\n \"name\": \"{name}\",\n \"version\": \"1.0.0\",\n \"description\": \"hook plugin\",\n \"hooks\": {{\n \"PreToolUse\": [\"./hooks/pre.sh\"],\n \"PostToolUse\": [\"./hooks/post.sh\"],\n \"PostToolUseFailure\": [\"./hooks/failure.sh\"]\n }}\n}}"
),
)
.expect("write plugin manifest");
}
#[test]
fn collects_and_runs_hooks_from_enabled_plugins() {
// given
let config_home = temp_dir("config");
let first_source_root = temp_dir("source-a");
let second_source_root = temp_dir("source-b");
write_hook_plugin(
&first_source_root,
"first",
"plugin pre one",
"plugin post one",
"plugin failure one",
);
write_hook_plugin(
&second_source_root,
"second",
"plugin pre two",
"plugin post two",
"plugin failure two",
);
let mut manager = PluginManager::new(PluginManagerConfig::new(&config_home));
manager
.install(first_source_root.to_str().expect("utf8 path"))
.expect("first plugin install should succeed");
manager
.install(second_source_root.to_str().expect("utf8 path"))
.expect("second plugin install should succeed");
let registry = manager.plugin_registry().expect("registry should build");
// when
let runner = HookRunner::from_registry(&registry).expect("plugin hooks should load");
// then
assert_eq!(
runner.run_pre_tool_use("Read", r#"{"path":"README.md"}"#),
HookRunResult::allow(vec![
"plugin pre one".to_string(),
"plugin pre two".to_string(),
])
);
assert_eq!(
runner.run_post_tool_use("Read", r#"{"path":"README.md"}"#, "ok", false),
HookRunResult::allow(vec![
"plugin post one".to_string(),
"plugin post two".to_string(),
])
);
assert_eq!(
runner.run_post_tool_use_failure("Read", r#"{"path":"README.md"}"#, "tool failed",),
HookRunResult::allow(vec![
"plugin failure one".to_string(),
"plugin failure two".to_string(),
])
);
let _ = fs::remove_dir_all(config_home);
let _ = fs::remove_dir_all(first_source_root);
let _ = fs::remove_dir_all(second_source_root);
}
#[test]
fn pre_tool_use_denies_when_plugin_hook_exits_two() {
// given
let runner = HookRunner::new(crate::PluginHooks {
pre_tool_use: vec!["printf 'blocked by plugin'; exit 2".to_string()],
post_tool_use: Vec::new(),
post_tool_use_failure: Vec::new(),
});
// when
let result = runner.run_pre_tool_use("Bash", r#"{"command":"pwd"}"#);
// then
assert!(result.is_denied());
assert_eq!(result.messages(), &["blocked by plugin".to_string()]);
}
#[test]
fn propagates_plugin_hook_failures() {
// given
let runner = HookRunner::new(crate::PluginHooks {
pre_tool_use: vec![
"printf 'broken plugin hook'; exit 1".to_string(),
"printf 'later plugin hook'".to_string(),
],
post_tool_use: Vec::new(),
post_tool_use_failure: Vec::new(),
});
// when
let result = runner.run_pre_tool_use("Bash", r#"{"command":"pwd"}"#);
// then
assert!(result.is_failed());
assert!(result
.messages()
.iter()
.any(|message| message.contains("broken plugin hook")));
assert!(!result
.messages()
.iter()
.any(|message| message == "later plugin hook"));
}
#[test]
#[cfg(unix)]
fn generated_hook_scripts_are_executable() {
use std::os::unix::fs::PermissionsExt;
// given
let root = temp_dir("exec-guard");
write_hook_plugin(&root, "exec-check", "pre", "post", "fail");
// then
for script in ["pre.sh", "post.sh", "failure.sh"] {
let path = root.join("hooks").join(script);
let mode = fs::metadata(&path)
.unwrap_or_else(|e| panic!("{script} metadata: {e}"))
.permissions()
.mode();
assert!(
mode & 0o111 != 0,
"{script} must have at least one execute bit set, got mode {mode:#o}"
);
}
}
}
File diff suppressed because it is too large. Load diff
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "runtime"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
[dependencies]
sha2 = "0.10"
glob = "0.3"
plugins = { path = "../plugins" }
regex = "1"
serde = { version = "1", features = ["derive"] }
serde_json.workspace = true
telemetry = { path = "../telemetry" }
tokio = { version = "1", features = ["io-std", "io-util", "macros", "process", "rt", "rt-multi-thread", "time"] }
walkdir = "2"
[lints]
workspace = true
+336
View File
@@ -0,0 +1,336 @@
use std::env;
use std::io;
use std::process::{Command, Stdio};
use std::time::Duration;
use serde::{Deserialize, Serialize};
use tokio::process::Command as TokioCommand;
use tokio::runtime::Builder;
use tokio::time::timeout;
use crate::sandbox::{
build_linux_sandbox_command, resolve_sandbox_status_for_request, FilesystemIsolationMode,
SandboxConfig, SandboxStatus,
};
use crate::ConfigLoader;
/// Input schema for the built-in bash execution tool.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct BashCommandInput {
pub command: String,
pub timeout: Option<u64>,
pub description: Option<String>,
#[serde(rename = "run_in_background")]
pub run_in_background: Option<bool>,
#[serde(rename = "dangerouslyDisableSandbox")]
pub dangerously_disable_sandbox: Option<bool>,
#[serde(rename = "namespaceRestrictions")]
pub namespace_restrictions: Option<bool>,
#[serde(rename = "isolateNetwork")]
pub isolate_network: Option<bool>,
#[serde(rename = "filesystemMode")]
pub filesystem_mode: Option<FilesystemIsolationMode>,
#[serde(rename = "allowedMounts")]
pub allowed_mounts: Option<Vec<String>>,
}
/// Output returned from a bash tool invocation.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct BashCommandOutput {
pub stdout: String,
pub stderr: String,
#[serde(rename = "rawOutputPath")]
pub raw_output_path: Option<String>,
pub interrupted: bool,
#[serde(rename = "isImage")]
pub is_image: Option<bool>,
#[serde(rename = "backgroundTaskId")]
pub background_task_id: Option<String>,
#[serde(rename = "backgroundedByUser")]
pub backgrounded_by_user: Option<bool>,
#[serde(rename = "assistantAutoBackgrounded")]
pub assistant_auto_backgrounded: Option<bool>,
#[serde(rename = "dangerouslyDisableSandbox")]
pub dangerously_disable_sandbox: Option<bool>,
#[serde(rename = "returnCodeInterpretation")]
pub return_code_interpretation: Option<String>,
#[serde(rename = "noOutputExpected")]
pub no_output_expected: Option<bool>,
#[serde(rename = "structuredContent")]
pub structured_content: Option<Vec<serde_json::Value>>,
#[serde(rename = "persistedOutputPath")]
pub persisted_output_path: Option<String>,
#[serde(rename = "persistedOutputSize")]
pub persisted_output_size: Option<u64>,
#[serde(rename = "sandboxStatus")]
pub sandbox_status: Option<SandboxStatus>,
}
/// Executes a shell command with the requested sandbox settings.
pub fn execute_bash(input: BashCommandInput) -> io::Result<BashCommandOutput> {
let cwd = env::current_dir()?;
let sandbox_status = sandbox_status_for_input(&input, &cwd);
if input.run_in_background.unwrap_or(false) {
let mut child = prepare_command(&input.command, &cwd, &sandbox_status, false);
let child = child
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()?;
return Ok(BashCommandOutput {
stdout: String::new(),
stderr: String::new(),
raw_output_path: None,
interrupted: false,
is_image: None,
background_task_id: Some(child.id().to_string()),
backgrounded_by_user: Some(false),
assistant_auto_backgrounded: Some(false),
dangerously_disable_sandbox: input.dangerously_disable_sandbox,
return_code_interpretation: None,
no_output_expected: Some(true),
structured_content: None,
persisted_output_path: None,
persisted_output_size: None,
sandbox_status: Some(sandbox_status),
});
}
let runtime = Builder::new_current_thread().enable_all().build()?;
runtime.block_on(execute_bash_async(input, sandbox_status, cwd))
}
async fn execute_bash_async(
input: BashCommandInput,
sandbox_status: SandboxStatus,
cwd: std::path::PathBuf,
) -> io::Result<BashCommandOutput> {
let mut command = prepare_tokio_command(&input.command, &cwd, &sandbox_status, true);
let output_result = if let Some(timeout_ms) = input.timeout {
match timeout(Duration::from_millis(timeout_ms), command.output()).await {
Ok(result) => (result?, false),
Err(_) => {
return Ok(BashCommandOutput {
stdout: String::new(),
stderr: format!("Command exceeded timeout of {timeout_ms} ms"),
raw_output_path: None,
interrupted: true,
is_image: None,
background_task_id: None,
backgrounded_by_user: None,
assistant_auto_backgrounded: None,
dangerously_disable_sandbox: input.dangerously_disable_sandbox,
return_code_interpretation: Some(String::from("timeout")),
no_output_expected: Some(true),
structured_content: None,
persisted_output_path: None,
persisted_output_size: None,
sandbox_status: Some(sandbox_status),
});
}
}
} else {
(command.output().await?, false)
};
let (output, interrupted) = output_result;
let stdout = truncate_output(&String::from_utf8_lossy(&output.stdout));
let stderr = truncate_output(&String::from_utf8_lossy(&output.stderr));
let no_output_expected = Some(stdout.trim().is_empty() && stderr.trim().is_empty());
let return_code_interpretation = output.status.code().and_then(|code| {
if code == 0 {
None
} else {
Some(format!("exit_code:{code}"))
}
});
Ok(BashCommandOutput {
stdout,
stderr,
raw_output_path: None,
interrupted,
is_image: None,
background_task_id: None,
backgrounded_by_user: None,
assistant_auto_backgrounded: None,
dangerously_disable_sandbox: input.dangerously_disable_sandbox,
return_code_interpretation,
no_output_expected,
structured_content: None,
persisted_output_path: None,
persisted_output_size: None,
sandbox_status: Some(sandbox_status),
})
}
fn sandbox_status_for_input(input: &BashCommandInput, cwd: &std::path::Path) -> SandboxStatus {
let config = ConfigLoader::default_for(cwd).load().map_or_else(
|_| SandboxConfig::default(),
|runtime_config| runtime_config.sandbox().clone(),
);
let request = config.resolve_request(
input.dangerously_disable_sandbox.map(|disabled| !disabled),
input.namespace_restrictions,
input.isolate_network,
input.filesystem_mode,
input.allowed_mounts.clone(),
);
resolve_sandbox_status_for_request(&request, cwd)
}
fn prepare_command(
command: &str,
cwd: &std::path::Path,
sandbox_status: &SandboxStatus,
create_dirs: bool,
) -> Command {
if create_dirs {
prepare_sandbox_dirs(cwd);
}
if let Some(launcher) = build_linux_sandbox_command(command, cwd, sandbox_status) {
let mut prepared = Command::new(launcher.program);
prepared.args(launcher.args);
prepared.current_dir(cwd);
prepared.envs(launcher.env);
return prepared;
}
let mut prepared = Command::new("sh");
prepared.arg("-lc").arg(command).current_dir(cwd);
if sandbox_status.filesystem_active {
prepared.env("HOME", cwd.join(".sandbox-home"));
prepared.env("TMPDIR", cwd.join(".sandbox-tmp"));
}
prepared
}
fn prepare_tokio_command(
command: &str,
cwd: &std::path::Path,
sandbox_status: &SandboxStatus,
create_dirs: bool,
) -> TokioCommand {
if create_dirs {
prepare_sandbox_dirs(cwd);
}
if let Some(launcher) = build_linux_sandbox_command(command, cwd, sandbox_status) {
let mut prepared = TokioCommand::new(launcher.program);
prepared.args(launcher.args);
prepared.current_dir(cwd);
prepared.envs(launcher.env);
return prepared;
}
let mut prepared = TokioCommand::new("sh");
prepared.arg("-lc").arg(command).current_dir(cwd);
if sandbox_status.filesystem_active {
prepared.env("HOME", cwd.join(".sandbox-home"));
prepared.env("TMPDIR", cwd.join(".sandbox-tmp"));
}
prepared
}
fn prepare_sandbox_dirs(cwd: &std::path::Path) {
let _ = std::fs::create_dir_all(cwd.join(".sandbox-home"));
let _ = std::fs::create_dir_all(cwd.join(".sandbox-tmp"));
}
#[cfg(test)]
mod tests {
use super::{execute_bash, BashCommandInput};
use crate::sandbox::FilesystemIsolationMode;
#[test]
fn executes_simple_command() {
let output = execute_bash(BashCommandInput {
command: String::from("printf 'hello'"),
timeout: Some(1_000),
description: None,
run_in_background: Some(false),
dangerously_disable_sandbox: Some(false),
namespace_restrictions: Some(false),
isolate_network: Some(false),
filesystem_mode: Some(FilesystemIsolationMode::WorkspaceOnly),
allowed_mounts: None,
})
.expect("bash command should execute");
assert_eq!(output.stdout, "hello");
assert!(!output.interrupted);
assert!(output.sandbox_status.is_some());
}
#[test]
fn disables_sandbox_when_requested() {
let output = execute_bash(BashCommandInput {
command: String::from("printf 'hello'"),
timeout: Some(1_000),
description: None,
run_in_background: Some(false),
dangerously_disable_sandbox: Some(true),
namespace_restrictions: None,
isolate_network: None,
filesystem_mode: None,
allowed_mounts: None,
})
.expect("bash command should execute");
assert!(!output.sandbox_status.expect("sandbox status").enabled);
}
}
/// Maximum output bytes before truncation (16 KiB, matching upstream).
const MAX_OUTPUT_BYTES: usize = 16_384;
/// Truncate output to `MAX_OUTPUT_BYTES`, appending a marker when trimmed.
fn truncate_output(s: &str) -> String {
if s.len() <= MAX_OUTPUT_BYTES {
return s.to_string();
}
// Find the last valid UTF-8 boundary at or before MAX_OUTPUT_BYTES
let mut end = MAX_OUTPUT_BYTES;
while end > 0 && !s.is_char_boundary(end) {
end -= 1;
}
let mut truncated = s[..end].to_string();
truncated.push_str("\n\n[output truncated — exceeded 16384 bytes]");
truncated
}
#[cfg(test)]
mod truncation_tests {
use super::*;
#[test]
fn short_output_unchanged() {
let s = "hello world";
assert_eq!(truncate_output(s), s);
}
#[test]
fn long_output_truncated() {
let s = "x".repeat(20_000);
let result = truncate_output(&s);
assert!(result.len() < 20_000);
assert!(result.ends_with("[output truncated — exceeded 16384 bytes]"));
}
#[test]
fn exact_boundary_unchanged() {
let s = "a".repeat(MAX_OUTPUT_BYTES);
assert_eq!(truncate_output(&s), s);
}
#[test]
fn one_over_boundary_truncated() {
let s = "a".repeat(MAX_OUTPUT_BYTES + 1);
let result = truncate_output(&s);
assert!(result.contains("[output truncated"));
}
}
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,111 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BootstrapPhase {
CliEntry,
FastPathVersion,
StartupProfiler,
SystemPromptFastPath,
ChromeMcpFastPath,
DaemonWorkerFastPath,
BridgeFastPath,
DaemonFastPath,
BackgroundSessionFastPath,
TemplateFastPath,
EnvironmentRunnerFastPath,
MainRuntime,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BootstrapPlan {
phases: Vec<BootstrapPhase>,
}
impl BootstrapPlan {
#[must_use]
pub fn claude_code_default() -> Self {
Self::from_phases(vec![
BootstrapPhase::CliEntry,
BootstrapPhase::FastPathVersion,
BootstrapPhase::StartupProfiler,
BootstrapPhase::SystemPromptFastPath,
BootstrapPhase::ChromeMcpFastPath,
BootstrapPhase::DaemonWorkerFastPath,
BootstrapPhase::BridgeFastPath,
BootstrapPhase::DaemonFastPath,
BootstrapPhase::BackgroundSessionFastPath,
BootstrapPhase::TemplateFastPath,
BootstrapPhase::EnvironmentRunnerFastPath,
BootstrapPhase::MainRuntime,
])
}
#[must_use]
pub fn from_phases(phases: Vec<BootstrapPhase>) -> Self {
let mut deduped = Vec::new();
for phase in phases {
if !deduped.contains(&phase) {
deduped.push(phase);
}
}
Self { phases: deduped }
}
#[must_use]
pub fn phases(&self) -> &[BootstrapPhase] {
&self.phases
}
}
#[cfg(test)]
mod tests {
use super::{BootstrapPhase, BootstrapPlan};
#[test]
fn from_phases_deduplicates_while_preserving_order() {
// given
let phases = vec![
BootstrapPhase::CliEntry,
BootstrapPhase::FastPathVersion,
BootstrapPhase::CliEntry,
BootstrapPhase::MainRuntime,
BootstrapPhase::FastPathVersion,
];
// when
let plan = BootstrapPlan::from_phases(phases);
// then
assert_eq!(
plan.phases(),
&[
BootstrapPhase::CliEntry,
BootstrapPhase::FastPathVersion,
BootstrapPhase::MainRuntime,
]
);
}
#[test]
fn claude_code_default_covers_each_phase_once() {
// given
let expected = [
BootstrapPhase::CliEntry,
BootstrapPhase::FastPathVersion,
BootstrapPhase::StartupProfiler,
BootstrapPhase::SystemPromptFastPath,
BootstrapPhase::ChromeMcpFastPath,
BootstrapPhase::DaemonWorkerFastPath,
BootstrapPhase::BridgeFastPath,
BootstrapPhase::DaemonFastPath,
BootstrapPhase::BackgroundSessionFastPath,
BootstrapPhase::TemplateFastPath,
BootstrapPhase::EnvironmentRunnerFastPath,
BootstrapPhase::MainRuntime,
];
// when
let plan = BootstrapPlan::claude_code_default();
// then
assert_eq!(plan.phases(), &expected);
}
}
@@ -0,0 +1,144 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BranchLockIntent {
#[serde(rename = "laneId")]
pub lane_id: String,
pub branch: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub worktree: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub modules: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BranchLockCollision {
pub branch: String,
pub module: String,
#[serde(rename = "laneIds")]
pub lane_ids: Vec<String>,
}
#[must_use]
pub fn detect_branch_lock_collisions(intents: &[BranchLockIntent]) -> Vec<BranchLockCollision> {
let mut collisions = Vec::new();
for (index, left) in intents.iter().enumerate() {
for right in &intents[index + 1..] {
if left.branch != right.branch {
continue;
}
for module in overlapping_modules(&left.modules, &right.modules) {
collisions.push(BranchLockCollision {
branch: left.branch.clone(),
module,
lane_ids: vec![left.lane_id.clone(), right.lane_id.clone()],
});
}
}
}
collisions.sort_by(|a, b| {
a.branch
.cmp(&b.branch)
.then(a.module.cmp(&b.module))
.then(a.lane_ids.cmp(&b.lane_ids))
});
collisions.dedup();
collisions
}
fn overlapping_modules(left: &[String], right: &[String]) -> Vec<String> {
let mut overlaps = Vec::new();
for left_module in left {
for right_module in right {
if modules_overlap(left_module, right_module) {
overlaps.push(shared_scope(left_module, right_module));
}
}
}
overlaps.sort();
overlaps.dedup();
overlaps
}
fn modules_overlap(left: &str, right: &str) -> bool {
left == right
|| left.starts_with(&format!("{right}/"))
|| right.starts_with(&format!("{left}/"))
}
fn shared_scope(left: &str, right: &str) -> String {
if left.starts_with(&format!("{right}/")) || left == right {
right.to_string()
} else {
left.to_string()
}
}
#[cfg(test)]
mod tests {
use super::{detect_branch_lock_collisions, BranchLockIntent};
#[test]
fn detects_same_branch_same_module_collisions() {
let collisions = detect_branch_lock_collisions(&[
BranchLockIntent {
lane_id: "lane-a".to_string(),
branch: "feature/lock".to_string(),
worktree: Some("wt-a".to_string()),
modules: vec!["runtime/mcp".to_string()],
},
BranchLockIntent {
lane_id: "lane-b".to_string(),
branch: "feature/lock".to_string(),
worktree: Some("wt-b".to_string()),
modules: vec!["runtime/mcp".to_string()],
},
]);
assert_eq!(collisions.len(), 1);
assert_eq!(collisions[0].branch, "feature/lock");
assert_eq!(collisions[0].module, "runtime/mcp");
}
#[test]
fn detects_nested_module_scope_collisions() {
let collisions = detect_branch_lock_collisions(&[
BranchLockIntent {
lane_id: "lane-a".to_string(),
branch: "feature/lock".to_string(),
worktree: None,
modules: vec!["runtime".to_string()],
},
BranchLockIntent {
lane_id: "lane-b".to_string(),
branch: "feature/lock".to_string(),
worktree: None,
modules: vec!["runtime/mcp".to_string()],
},
]);
assert_eq!(collisions[0].module, "runtime");
}
#[test]
fn ignores_different_branches() {
let collisions = detect_branch_lock_collisions(&[
BranchLockIntent {
lane_id: "lane-a".to_string(),
branch: "feature/a".to_string(),
worktree: None,
modules: vec!["runtime/mcp".to_string()],
},
BranchLockIntent {
lane_id: "lane-b".to_string(),
branch: "feature/b".to_string(),
worktree: None,
modules: vec!["runtime/mcp".to_string()],
},
]);
assert!(collisions.is_empty());
}
}
@@ -0,0 +1,827 @@
use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
const COMPACT_CONTINUATION_PREAMBLE: &str =
"This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.\n\n";
const COMPACT_RECENT_MESSAGES_NOTE: &str = "Recent messages are preserved verbatim.";
const COMPACT_DIRECT_RESUME_INSTRUCTION: &str = "Continue the conversation from where it left off without asking the user any further questions. Resume directly — do not acknowledge the summary, do not recap what was happening, and do not preface with continuation text.";
/// Thresholds controlling when and how a session is compacted.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CompactionConfig {
pub preserve_recent_messages: usize,
pub max_estimated_tokens: usize,
}
impl Default for CompactionConfig {
fn default() -> Self {
Self {
preserve_recent_messages: 4,
max_estimated_tokens: 10_000,
}
}
}
/// Result of compacting a session into a summary plus preserved tail messages.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CompactionResult {
pub summary: String,
pub formatted_summary: String,
pub compacted_session: Session,
pub removed_message_count: usize,
}
/// Roughly estimates the token footprint of the current session transcript.
#[must_use]
pub fn estimate_session_tokens(session: &Session) -> usize {
session.messages.iter().map(estimate_message_tokens).sum()
}
/// Returns `true` when the session exceeds the configured compaction budget.
#[must_use]
pub fn should_compact(session: &Session, config: CompactionConfig) -> bool {
let start = compacted_summary_prefix_len(session);
let compactable = &session.messages[start..];
compactable.len() > config.preserve_recent_messages
&& compactable
.iter()
.map(estimate_message_tokens)
.sum::<usize>()
>= config.max_estimated_tokens
}
/// Normalizes a compaction summary into user-facing continuation text.
#[must_use]
pub fn format_compact_summary(summary: &str) -> String {
let without_analysis = strip_tag_block(summary, "analysis");
let formatted = if let Some(content) = extract_tag_block(&without_analysis, "summary") {
without_analysis.replace(
&format!("<summary>{content}</summary>"),
&format!("Summary:\n{}", content.trim()),
)
} else {
without_analysis
};
collapse_blank_lines(&formatted).trim().to_string()
}
/// Builds the synthetic system message used after session compaction.
#[must_use]
pub fn get_compact_continuation_message(
summary: &str,
suppress_follow_up_questions: bool,
recent_messages_preserved: bool,
) -> String {
let mut base = format!(
"{COMPACT_CONTINUATION_PREAMBLE}{}",
format_compact_summary(summary)
);
if recent_messages_preserved {
base.push_str("\n\n");
base.push_str(COMPACT_RECENT_MESSAGES_NOTE);
}
if suppress_follow_up_questions {
base.push('\n');
base.push_str(COMPACT_DIRECT_RESUME_INSTRUCTION);
}
base
}
/// Compacts a session by summarizing older messages and preserving the recent tail.
#[must_use]
pub fn compact_session(session: &Session, config: CompactionConfig) -> CompactionResult {
if !should_compact(session, config) {
return CompactionResult {
summary: String::new(),
formatted_summary: String::new(),
compacted_session: session.clone(),
removed_message_count: 0,
};
}
let existing_summary = session
.messages
.first()
.and_then(extract_existing_compacted_summary);
let compacted_prefix_len = usize::from(existing_summary.is_some());
let raw_keep_from = session
.messages
.len()
.saturating_sub(config.preserve_recent_messages);
// Ensure we do not split a tool-use / tool-result pair at the compaction
// boundary. If the first preserved message is a user message whose first
// block is a ToolResult, the assistant message with the matching ToolUse
// was slated for removal — that produces an orphaned tool role message on
// the OpenAI-compat path (400: tool message must follow assistant with
// tool_calls). Walk the boundary back until we start at a safe point.
let keep_from = {
let mut k = raw_keep_from;
// If the first preserved message is a tool-result turn, ensure its
// paired assistant tool-use turn is preserved too. Without this fix,
// the OpenAI-compat adapter sends an orphaned 'tool' role message
// with no preceding assistant 'tool_calls', which providers reject
// with a 400. We walk back only if the immediately preceding message
// is NOT an assistant message that contains a ToolUse block (i.e. the
// pair is actually broken at the boundary).
loop {
if k == 0 || k <= compacted_prefix_len {
break;
}
let first_preserved = &session.messages[k];
let starts_with_tool_result = first_preserved
.blocks
.first()
.map(|b| matches!(b, ContentBlock::ToolResult { .. }))
.unwrap_or(false);
if !starts_with_tool_result {
break;
}
// Check the message just before the current boundary.
let preceding = &session.messages[k - 1];
let preceding_has_tool_use = preceding
.blocks
.iter()
.any(|b| matches!(b, ContentBlock::ToolUse { .. }));
if preceding_has_tool_use {
// Pair is intact — walk back one more to include the assistant turn.
k = k.saturating_sub(1);
break;
}
// Preceding message has no ToolUse but we have a ToolResult —
// this is already an orphaned pair; walk back to try to fix it.
k = k.saturating_sub(1);
}
k
};
let removed = &session.messages[compacted_prefix_len..keep_from];
let preserved = session.messages[keep_from..].to_vec();
let summary =
merge_compact_summaries(existing_summary.as_deref(), &summarize_messages(removed));
let formatted_summary = format_compact_summary(&summary);
let continuation = get_compact_continuation_message(&summary, true, !preserved.is_empty());
let mut compacted_messages = vec![ConversationMessage {
role: MessageRole::System,
blocks: vec![ContentBlock::Text { text: continuation }],
usage: None,
}];
compacted_messages.extend(preserved);
let mut compacted_session = session.clone();
compacted_session.messages = compacted_messages;
compacted_session.record_compaction(summary.clone(), removed.len());
CompactionResult {
summary,
formatted_summary,
compacted_session,
removed_message_count: removed.len(),
}
}
fn compacted_summary_prefix_len(session: &Session) -> usize {
usize::from(
session
.messages
.first()
.and_then(extract_existing_compacted_summary)
.is_some(),
)
}
fn summarize_messages(messages: &[ConversationMessage]) -> String {
let user_messages = messages
.iter()
.filter(|message| message.role == MessageRole::User)
.count();
let assistant_messages = messages
.iter()
.filter(|message| message.role == MessageRole::Assistant)
.count();
let tool_messages = messages
.iter()
.filter(|message| message.role == MessageRole::Tool)
.count();
let mut tool_names = messages
.iter()
.flat_map(|message| message.blocks.iter())
.filter_map(|block| match block {
ContentBlock::ToolUse { name, .. } => Some(name.as_str()),
ContentBlock::ToolResult { tool_name, .. } => Some(tool_name.as_str()),
ContentBlock::Text { .. } => None,
})
.collect::<Vec<_>>();
tool_names.sort_unstable();
tool_names.dedup();
let mut lines = vec![
"<summary>".to_string(),
"Conversation summary:".to_string(),
format!(
"- Scope: {} earlier messages compacted (user={}, assistant={}, tool={}).",
messages.len(),
user_messages,
assistant_messages,
tool_messages
),
];
if !tool_names.is_empty() {
lines.push(format!("- Tools mentioned: {}.", tool_names.join(", ")));
}
let recent_user_requests = collect_recent_role_summaries(messages, MessageRole::User, 3);
if !recent_user_requests.is_empty() {
lines.push("- Recent user requests:".to_string());
lines.extend(
recent_user_requests
.into_iter()
.map(|request| format!(" - {request}")),
);
}
let pending_work = infer_pending_work(messages);
if !pending_work.is_empty() {
lines.push("- Pending work:".to_string());
lines.extend(pending_work.into_iter().map(|item| format!(" - {item}")));
}
let key_files = collect_key_files(messages);
if !key_files.is_empty() {
lines.push(format!("- Key files referenced: {}.", key_files.join(", ")));
}
if let Some(current_work) = infer_current_work(messages) {
lines.push(format!("- Current work: {current_work}"));
}
lines.push("- Key timeline:".to_string());
for message in messages {
let role = match message.role {
MessageRole::System => "system",
MessageRole::User => "user",
MessageRole::Assistant => "assistant",
MessageRole::Tool => "tool",
};
let content = message
.blocks
.iter()
.map(summarize_block)
.collect::<Vec<_>>()
.join(" | ");
lines.push(format!(" - {role}: {content}"));
}
lines.push("</summary>".to_string());
lines.join("\n")
}
fn merge_compact_summaries(existing_summary: Option<&str>, new_summary: &str) -> String {
let Some(existing_summary) = existing_summary else {
return new_summary.to_string();
};
let previous_highlights = extract_summary_highlights(existing_summary);
let new_formatted_summary = format_compact_summary(new_summary);
let new_highlights = extract_summary_highlights(&new_formatted_summary);
let new_timeline = extract_summary_timeline(&new_formatted_summary);
let mut lines = vec!["<summary>".to_string(), "Conversation summary:".to_string()];
if !previous_highlights.is_empty() {
lines.push("- Previously compacted context:".to_string());
lines.extend(
previous_highlights
.into_iter()
.map(|line| format!(" {line}")),
);
}
if !new_highlights.is_empty() {
lines.push("- Newly compacted context:".to_string());
lines.extend(new_highlights.into_iter().map(|line| format!(" {line}")));
}
if !new_timeline.is_empty() {
lines.push("- Key timeline:".to_string());
lines.extend(new_timeline.into_iter().map(|line| format!(" {line}")));
}
lines.push("</summary>".to_string());
lines.join("\n")
}
fn summarize_block(block: &ContentBlock) -> String {
let raw = match block {
ContentBlock::Text { text } => text.clone(),
ContentBlock::ToolUse { name, input, .. } => format!("tool_use {name}({input})"),
ContentBlock::ToolResult {
tool_name,
output,
is_error,
..
} => format!(
"tool_result {tool_name}: {}{output}",
if *is_error { "error " } else { "" }
),
};
truncate_summary(&raw, 160)
}
fn collect_recent_role_summaries(
messages: &[ConversationMessage],
role: MessageRole,
limit: usize,
) -> Vec<String> {
messages
.iter()
.filter(|message| message.role == role)
.rev()
.filter_map(|message| first_text_block(message))
.take(limit)
.map(|text| truncate_summary(text, 160))
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect()
}
fn infer_pending_work(messages: &[ConversationMessage]) -> Vec<String> {
messages
.iter()
.rev()
.filter_map(first_text_block)
.filter(|text| {
let lowered = text.to_ascii_lowercase();
lowered.contains("todo")
|| lowered.contains("next")
|| lowered.contains("pending")
|| lowered.contains("follow up")
|| lowered.contains("remaining")
})
.take(3)
.map(|text| truncate_summary(text, 160))
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect()
}
fn collect_key_files(messages: &[ConversationMessage]) -> Vec<String> {
let mut files = messages
.iter()
.flat_map(|message| message.blocks.iter())
.map(|block| match block {
ContentBlock::Text { text } => text.as_str(),
ContentBlock::ToolUse { input, .. } => input.as_str(),
ContentBlock::ToolResult { output, .. } => output.as_str(),
})
.flat_map(extract_file_candidates)
.collect::<Vec<_>>();
files.sort();
files.dedup();
files.into_iter().take(8).collect()
}
fn infer_current_work(messages: &[ConversationMessage]) -> Option<String> {
messages
.iter()
.rev()
.filter_map(first_text_block)
.find(|text| !text.trim().is_empty())
.map(|text| truncate_summary(text, 200))
}
fn first_text_block(message: &ConversationMessage) -> Option<&str> {
message.blocks.iter().find_map(|block| match block {
ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.as_str()),
ContentBlock::ToolUse { .. }
| ContentBlock::ToolResult { .. }
| ContentBlock::Text { .. } => None,
})
}
fn has_interesting_extension(candidate: &str) -> bool {
std::path::Path::new(candidate)
.extension()
.and_then(|extension| extension.to_str())
.is_some_and(|extension| {
["rs", "ts", "tsx", "js", "json", "md"]
.iter()
.any(|expected| extension.eq_ignore_ascii_case(expected))
})
}
fn extract_file_candidates(content: &str) -> Vec<String> {
content
.split_whitespace()
.filter_map(|token| {
let candidate = token.trim_matches(|char: char| {
matches!(char, ',' | '.' | ':' | ';' | ')' | '(' | '"' | '\'' | '`')
});
if candidate.contains('/') && has_interesting_extension(candidate) {
Some(candidate.to_string())
} else {
None
}
})
.collect()
}
fn truncate_summary(content: &str, max_chars: usize) -> String {
if content.chars().count() <= max_chars {
return content.to_string();
}
let mut truncated = content.chars().take(max_chars).collect::<String>();
truncated.push('…');
truncated
}
fn estimate_message_tokens(message: &ConversationMessage) -> usize {
message
.blocks
.iter()
.map(|block| match block {
ContentBlock::Text { text } => text.len() / 4 + 1,
ContentBlock::ToolUse { name, input, .. } => (name.len() + input.len()) / 4 + 1,
ContentBlock::ToolResult {
tool_name, output, ..
} => (tool_name.len() + output.len()) / 4 + 1,
})
.sum()
}
fn extract_tag_block(content: &str, tag: &str) -> Option<String> {
let start = format!("<{tag}>");
let end = format!("</{tag}>");
let start_index = content.find(&start)? + start.len();
let end_index = content[start_index..].find(&end)? + start_index;
Some(content[start_index..end_index].to_string())
}
fn strip_tag_block(content: &str, tag: &str) -> String {
let start = format!("<{tag}>");
let end = format!("</{tag}>");
if let (Some(start_index), Some(end_index_rel)) = (content.find(&start), content.find(&end)) {
let end_index = end_index_rel + end.len();
let mut stripped = String::new();
stripped.push_str(&content[..start_index]);
stripped.push_str(&content[end_index..]);
stripped
} else {
content.to_string()
}
}
fn collapse_blank_lines(content: &str) -> String {
let mut result = String::new();
let mut last_blank = false;
for line in content.lines() {
let is_blank = line.trim().is_empty();
if is_blank && last_blank {
continue;
}
result.push_str(line);
result.push('\n');
last_blank = is_blank;
}
result
}
fn extract_existing_compacted_summary(message: &ConversationMessage) -> Option<String> {
if message.role != MessageRole::System {
return None;
}
let text = first_text_block(message)?;
let summary = text.strip_prefix(COMPACT_CONTINUATION_PREAMBLE)?;
let summary = summary
.split_once(&format!("\n\n{COMPACT_RECENT_MESSAGES_NOTE}"))
.map_or(summary, |(value, _)| value);
let summary = summary
.split_once(&format!("\n{COMPACT_DIRECT_RESUME_INSTRUCTION}"))
.map_or(summary, |(value, _)| value);
Some(summary.trim().to_string())
}
fn extract_summary_highlights(summary: &str) -> Vec<String> {
let mut lines = Vec::new();
let mut in_timeline = false;
for line in format_compact_summary(summary).lines() {
let trimmed = line.trim_end();
if trimmed.is_empty() || trimmed == "Summary:" || trimmed == "Conversation summary:" {
continue;
}
if trimmed == "- Key timeline:" {
in_timeline = true;
continue;
}
if in_timeline {
continue;
}
lines.push(trimmed.to_string());
}
lines
}
fn extract_summary_timeline(summary: &str) -> Vec<String> {
let mut lines = Vec::new();
let mut in_timeline = false;
for line in format_compact_summary(summary).lines() {
let trimmed = line.trim_end();
if trimmed == "- Key timeline:" {
in_timeline = true;
continue;
}
if !in_timeline {
continue;
}
if trimmed.is_empty() {
break;
}
lines.push(trimmed.to_string());
}
lines
}
#[cfg(test)]
mod tests {
use super::{
collect_key_files, compact_session, format_compact_summary,
get_compact_continuation_message, infer_pending_work, should_compact, CompactionConfig,
};
use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
#[test]
fn formats_compact_summary_like_upstream() {
let summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>";
assert_eq!(format_compact_summary(summary), "Summary:\nKept work");
}
#[test]
fn leaves_small_sessions_unchanged() {
let mut session = Session::new();
session.messages = vec![ConversationMessage::user_text("hello")];
let result = compact_session(&session, CompactionConfig::default());
assert_eq!(result.removed_message_count, 0);
assert_eq!(result.compacted_session, session);
assert!(result.summary.is_empty());
assert!(result.formatted_summary.is_empty());
}
#[test]
fn compacts_older_messages_into_a_system_summary() {
let mut session = Session::new();
session.messages = vec![
ConversationMessage::user_text("one ".repeat(200)),
ConversationMessage::assistant(vec![ContentBlock::Text {
text: "two ".repeat(200),
}]),
ConversationMessage::tool_result("1", "bash", "ok ".repeat(200), false),
ConversationMessage {
role: MessageRole::Assistant,
blocks: vec![ContentBlock::Text {
text: "recent".to_string(),
}],
usage: None,
},
];
let result = compact_session(
&session,
CompactionConfig {
preserve_recent_messages: 2,
max_estimated_tokens: 1,
},
);
// With the tool-use/tool-result boundary fix, the compaction preserves
// one extra message to avoid an orphaned tool result at the boundary.
// messages[1] (assistant) must be kept along with messages[2] (tool result).
assert!(
result.removed_message_count <= 2,
"expected at most 2 removed, got {}",
result.removed_message_count
);
assert_eq!(
result.compacted_session.messages[0].role,
MessageRole::System
);
assert!(matches!(
&result.compacted_session.messages[0].blocks[0],
ContentBlock::Text { text } if text.contains("Summary:")
));
assert!(result.formatted_summary.contains("Scope:"));
assert!(result.formatted_summary.contains("Key timeline:"));
assert!(should_compact(
&session,
CompactionConfig {
preserve_recent_messages: 2,
max_estimated_tokens: 1,
}
));
// Note: with the tool-use/tool-result boundary guard the compacted session
// may preserve one extra message at the boundary, so token reduction is
// not guaranteed for small sessions. The invariant that matters is that
// the removed_message_count is non-zero (something was compacted).
assert!(
result.removed_message_count > 0,
"compaction must remove at least one message"
);
}
#[test]
fn keeps_previous_compacted_context_when_compacting_again() {
let mut initial_session = Session::new();
initial_session.messages = vec![
ConversationMessage::user_text("Investigate rust/crates/runtime/src/compact.rs"),
ConversationMessage::assistant(vec![ContentBlock::Text {
text: "I will inspect the compact flow.".to_string(),
}]),
ConversationMessage::user_text("Also update rust/crates/runtime/src/conversation.rs"),
ConversationMessage::assistant(vec![ContentBlock::Text {
text: "Next: preserve prior summary context during auto compact.".to_string(),
}]),
];
let config = CompactionConfig {
preserve_recent_messages: 2,
max_estimated_tokens: 1,
};
let first = compact_session(&initial_session, config);
let mut follow_up_messages = first.compacted_session.messages.clone();
follow_up_messages.extend([
ConversationMessage::user_text("Please add regression tests for compaction."),
ConversationMessage::assistant(vec![ContentBlock::Text {
text: "Working on regression coverage now.".to_string(),
}]),
]);
let mut second_session = Session::new();
second_session.messages = follow_up_messages;
let second = compact_session(&second_session, config);
assert!(second
.formatted_summary
.contains("Previously compacted context:"));
assert!(second
.formatted_summary
.contains("Scope: 2 earlier messages compacted"));
assert!(second
.formatted_summary
.contains("Newly compacted context:"));
assert!(second
.formatted_summary
.contains("Also update rust/crates/runtime/src/conversation.rs"));
assert!(matches!(
&second.compacted_session.messages[0].blocks[0],
ContentBlock::Text { text }
if text.contains("Previously compacted context:")
&& text.contains("Newly compacted context:")
));
assert!(matches!(
&second.compacted_session.messages[1].blocks[0],
ContentBlock::Text { text } if text.contains("Please add regression tests for compaction.")
));
}
#[test]
fn ignores_existing_compacted_summary_when_deciding_to_recompact() {
let summary = "<summary>Conversation summary:\n- Scope: earlier work preserved.\n- Key timeline:\n - user: large preserved context\n</summary>";
let mut session = Session::new();
session.messages = vec![
ConversationMessage {
role: MessageRole::System,
blocks: vec![ContentBlock::Text {
text: get_compact_continuation_message(summary, true, true),
}],
usage: None,
},
ConversationMessage::user_text("tiny"),
ConversationMessage::assistant(vec![ContentBlock::Text {
text: "recent".to_string(),
}]),
];
assert!(!should_compact(
&session,
CompactionConfig {
preserve_recent_messages: 2,
max_estimated_tokens: 1,
}
));
}
#[test]
fn truncates_long_blocks_in_summary() {
let summary = super::summarize_block(&ContentBlock::Text {
text: "x".repeat(400),
});
assert!(summary.ends_with('…'));
assert!(summary.chars().count() <= 161);
}
#[test]
fn extracts_key_files_from_message_content() {
let files = collect_key_files(&[ConversationMessage::user_text(
"Update rust/crates/runtime/src/compact.rs and rust/crates/rusty-claude-cli/src/main.rs next.",
)]);
assert!(files.contains(&"rust/crates/runtime/src/compact.rs".to_string()));
assert!(files.contains(&"rust/crates/rusty-claude-cli/src/main.rs".to_string()));
}
/// Regression: compaction must not split an assistant(ToolUse) /
/// user(ToolResult) pair at the boundary. An orphaned tool-result message
/// without the preceding assistant tool_calls causes a 400 on the
/// OpenAI-compat path (gaebal-gajae repro 2026-04-09).
#[test]
fn compaction_does_not_split_tool_use_tool_result_pair() {
use crate::session::{ContentBlock, Session};
let tool_id = "call_abc";
let mut session = Session::default();
// Turn 1: user prompt
session
.push_message(ConversationMessage::user_text("Search for files"))
.unwrap();
// Turn 2: assistant calls a tool
session
.push_message(ConversationMessage::assistant(vec![
ContentBlock::ToolUse {
id: tool_id.to_string(),
name: "search".to_string(),
input: "{\"q\":\"*.rs\"}".to_string(),
},
]))
.unwrap();
// Turn 3: tool result
session
.push_message(ConversationMessage::tool_result(
tool_id,
"search",
"found 5 files",
false,
))
.unwrap();
// Turn 4: assistant final response
session
.push_message(ConversationMessage::assistant(vec![ContentBlock::Text {
text: "Done.".to_string(),
}]))
.unwrap();
// Compact preserving only 1 recent message — without the fix this
// would cut the boundary so that the tool result (turn 3) is first,
// without its preceding assistant tool_calls (turn 2).
let config = CompactionConfig {
preserve_recent_messages: 1,
..CompactionConfig::default()
};
let result = compact_session(&session, config);
// After compaction, no two consecutive messages should have the pattern
// tool_result immediately following a non-assistant message (i.e. an
// orphaned tool result without a preceding assistant ToolUse).
let messages = &result.compacted_session.messages;
for i in 1..messages.len() {
let curr_is_tool_result = messages[i]
.blocks
.first()
.map(|b| matches!(b, ContentBlock::ToolResult { .. }))
.unwrap_or(false);
if curr_is_tool_result {
let prev_has_tool_use = messages[i - 1]
.blocks
.iter()
.any(|b| matches!(b, ContentBlock::ToolUse { .. }));
assert!(
prev_has_tool_use,
"message[{}] is a ToolResult but message[{}] has no ToolUse: {:?}",
i,
i - 1,
&messages[i - 1].blocks
);
}
}
}
#[test]
fn infers_pending_work_from_recent_messages() {
let pending = infer_pending_work(&[
ConversationMessage::user_text("done"),
ConversationMessage::assistant(vec![ContentBlock::Text {
text: "Next: update tests and follow up on remaining CLI polish.".to_string(),
}]),
]);
assert_eq!(pending.len(), 1);
assert!(pending[0].contains("Next: update tests"));
}
}
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,901 @@
use std::collections::BTreeMap;
use std::path::Path;
use crate::config::ConfigError;
use crate::json::JsonValue;
/// Diagnostic emitted when a config file contains a suspect field.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConfigDiagnostic {
pub path: String,
pub field: String,
pub line: Option<usize>,
pub kind: DiagnosticKind,
}
/// Classification of the diagnostic.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DiagnosticKind {
UnknownKey {
suggestion: Option<String>,
},
WrongType {
expected: &'static str,
got: &'static str,
},
Deprecated {
replacement: &'static str,
},
}
impl std::fmt::Display for ConfigDiagnostic {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let location = self
.line
.map_or_else(String::new, |line| format!(" (line {line})"));
match &self.kind {
DiagnosticKind::UnknownKey { suggestion: None } => {
write!(f, "{}: unknown key \"{}\"{location}", self.path, self.field)
}
DiagnosticKind::UnknownKey {
suggestion: Some(hint),
} => {
write!(
f,
"{}: unknown key \"{}\"{location}. Did you mean \"{}\"?",
self.path, self.field, hint
)
}
DiagnosticKind::WrongType { expected, got } => {
write!(
f,
"{}: field \"{}\" must be {expected}, got {got}{location}",
self.path, self.field
)
}
DiagnosticKind::Deprecated { replacement } => {
write!(
f,
"{}: field \"{}\" is deprecated{location}. Use \"{replacement}\" instead",
self.path, self.field
)
}
}
}
}
/// Result of validating a single config file.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ValidationResult {
pub errors: Vec<ConfigDiagnostic>,
pub warnings: Vec<ConfigDiagnostic>,
}
impl ValidationResult {
#[must_use]
pub fn is_ok(&self) -> bool {
self.errors.is_empty()
}
fn merge(&mut self, other: Self) {
self.errors.extend(other.errors);
self.warnings.extend(other.warnings);
}
}
// ---- known-key schema ----
/// Expected type for a config field.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FieldType {
String,
Bool,
Object,
StringArray,
Number,
}
impl FieldType {
fn label(self) -> &'static str {
match self {
Self::String => "a string",
Self::Bool => "a boolean",
Self::Object => "an object",
Self::StringArray => "an array of strings",
Self::Number => "a number",
}
}
fn matches(self, value: &JsonValue) -> bool {
match self {
Self::String => value.as_str().is_some(),
Self::Bool => value.as_bool().is_some(),
Self::Object => value.as_object().is_some(),
Self::StringArray => value
.as_array()
.is_some_and(|arr| arr.iter().all(|v| v.as_str().is_some())),
Self::Number => value.as_i64().is_some(),
}
}
}
fn json_type_label(value: &JsonValue) -> &'static str {
match value {
JsonValue::Null => "null",
JsonValue::Bool(_) => "a boolean",
JsonValue::Number(_) => "a number",
JsonValue::String(_) => "a string",
JsonValue::Array(_) => "an array",
JsonValue::Object(_) => "an object",
}
}
struct FieldSpec {
name: &'static str,
expected: FieldType,
}
struct DeprecatedField {
name: &'static str,
replacement: &'static str,
}
const TOP_LEVEL_FIELDS: &[FieldSpec] = &[
FieldSpec {
name: "$schema",
expected: FieldType::String,
},
FieldSpec {
name: "model",
expected: FieldType::String,
},
FieldSpec {
name: "hooks",
expected: FieldType::Object,
},
FieldSpec {
name: "permissions",
expected: FieldType::Object,
},
FieldSpec {
name: "permissionMode",
expected: FieldType::String,
},
FieldSpec {
name: "mcpServers",
expected: FieldType::Object,
},
FieldSpec {
name: "oauth",
expected: FieldType::Object,
},
FieldSpec {
name: "enabledPlugins",
expected: FieldType::Object,
},
FieldSpec {
name: "plugins",
expected: FieldType::Object,
},
FieldSpec {
name: "sandbox",
expected: FieldType::Object,
},
FieldSpec {
name: "env",
expected: FieldType::Object,
},
FieldSpec {
name: "aliases",
expected: FieldType::Object,
},
FieldSpec {
name: "providerFallbacks",
expected: FieldType::Object,
},
FieldSpec {
name: "trustedRoots",
expected: FieldType::StringArray,
},
];
const HOOKS_FIELDS: &[FieldSpec] = &[
FieldSpec {
name: "PreToolUse",
expected: FieldType::StringArray,
},
FieldSpec {
name: "PostToolUse",
expected: FieldType::StringArray,
},
FieldSpec {
name: "PostToolUseFailure",
expected: FieldType::StringArray,
},
];
const PERMISSIONS_FIELDS: &[FieldSpec] = &[
FieldSpec {
name: "defaultMode",
expected: FieldType::String,
},
FieldSpec {
name: "allow",
expected: FieldType::StringArray,
},
FieldSpec {
name: "deny",
expected: FieldType::StringArray,
},
FieldSpec {
name: "ask",
expected: FieldType::StringArray,
},
];
const PLUGINS_FIELDS: &[FieldSpec] = &[
FieldSpec {
name: "enabled",
expected: FieldType::Object,
},
FieldSpec {
name: "externalDirectories",
expected: FieldType::StringArray,
},
FieldSpec {
name: "installRoot",
expected: FieldType::String,
},
FieldSpec {
name: "registryPath",
expected: FieldType::String,
},
FieldSpec {
name: "bundledRoot",
expected: FieldType::String,
},
FieldSpec {
name: "maxOutputTokens",
expected: FieldType::Number,
},
];
const SANDBOX_FIELDS: &[FieldSpec] = &[
FieldSpec {
name: "enabled",
expected: FieldType::Bool,
},
FieldSpec {
name: "namespaceRestrictions",
expected: FieldType::Bool,
},
FieldSpec {
name: "networkIsolation",
expected: FieldType::Bool,
},
FieldSpec {
name: "filesystemMode",
expected: FieldType::String,
},
FieldSpec {
name: "allowedMounts",
expected: FieldType::StringArray,
},
];
const OAUTH_FIELDS: &[FieldSpec] = &[
FieldSpec {
name: "clientId",
expected: FieldType::String,
},
FieldSpec {
name: "authorizeUrl",
expected: FieldType::String,
},
FieldSpec {
name: "tokenUrl",
expected: FieldType::String,
},
FieldSpec {
name: "callbackPort",
expected: FieldType::Number,
},
FieldSpec {
name: "manualRedirectUrl",
expected: FieldType::String,
},
FieldSpec {
name: "scopes",
expected: FieldType::StringArray,
},
];
const DEPRECATED_FIELDS: &[DeprecatedField] = &[
DeprecatedField {
name: "permissionMode",
replacement: "permissions.defaultMode",
},
DeprecatedField {
name: "enabledPlugins",
replacement: "plugins.enabled",
},
];
// ---- line-number resolution ----
/// Find the 1-based line number where a JSON key first appears in the raw source.
fn find_key_line(source: &str, key: &str) -> Option<usize> {
// Search for `"key"` followed by optional whitespace and a colon.
let needle = format!("\"{key}\"");
let mut search_start = 0;
while let Some(offset) = source[search_start..].find(&needle) {
let absolute = search_start + offset;
let after = absolute + needle.len();
// Verify the next non-whitespace char is `:` to confirm this is a key, not a value.
if source[after..].chars().find(|ch| !ch.is_ascii_whitespace()) == Some(':') {
return Some(source[..absolute].chars().filter(|&ch| ch == '\n').count() + 1);
}
search_start = after;
}
None
}
// ---- core validation ----
fn validate_object_keys(
object: &BTreeMap<String, JsonValue>,
known_fields: &[FieldSpec],
prefix: &str,
source: &str,
path_display: &str,
) -> ValidationResult {
let mut result = ValidationResult {
errors: Vec::new(),
warnings: Vec::new(),
};
let known_names: Vec<&str> = known_fields.iter().map(|f| f.name).collect();
for (key, value) in object {
let field_path = if prefix.is_empty() {
key.clone()
} else {
format!("{prefix}.{key}")
};
if let Some(spec) = known_fields.iter().find(|f| f.name == key) {
// Type check.
if !spec.expected.matches(value) {
result.errors.push(ConfigDiagnostic {
path: path_display.to_string(),
field: field_path,
line: find_key_line(source, key),
kind: DiagnosticKind::WrongType {
expected: spec.expected.label(),
got: json_type_label(value),
},
});
}
} else if DEPRECATED_FIELDS.iter().any(|d| d.name == key) {
// Deprecated key — handled separately, not an unknown-key error.
} else {
// Unknown key.
let suggestion = suggest_field(key, &known_names);
result.errors.push(ConfigDiagnostic {
path: path_display.to_string(),
field: field_path,
line: find_key_line(source, key),
kind: DiagnosticKind::UnknownKey { suggestion },
});
}
}
result
}
fn suggest_field(input: &str, candidates: &[&str]) -> Option<String> {
let input_lower = input.to_ascii_lowercase();
candidates
.iter()
.filter_map(|candidate| {
let distance = simple_edit_distance(&input_lower, &candidate.to_ascii_lowercase());
(distance <= 3).then_some((distance, *candidate))
})
.min_by_key(|(distance, _)| *distance)
.map(|(_, name)| name.to_string())
}
fn simple_edit_distance(left: &str, right: &str) -> usize {
if left.is_empty() {
return right.len();
}
if right.is_empty() {
return left.len();
}
let right_chars: Vec<char> = right.chars().collect();
let mut previous: Vec<usize> = (0..=right_chars.len()).collect();
let mut current = vec![0; right_chars.len() + 1];
for (left_index, left_char) in left.chars().enumerate() {
current[0] = left_index + 1;
for (right_index, right_char) in right_chars.iter().enumerate() {
let cost = usize::from(left_char != *right_char);
current[right_index + 1] = (previous[right_index + 1] + 1)
.min(current[right_index] + 1)
.min(previous[right_index] + cost);
}
previous.clone_from(&current);
}
previous[right_chars.len()]
}
/// Validate a parsed config file's keys and types against the known schema.
///
/// Returns diagnostics (errors and deprecation warnings) without blocking the load.
pub fn validate_config_file(
object: &BTreeMap<String, JsonValue>,
source: &str,
file_path: &Path,
) -> ValidationResult {
let path_display = file_path.display().to_string();
let mut result = validate_object_keys(object, TOP_LEVEL_FIELDS, "", source, &path_display);
// Check deprecated fields.
for deprecated in DEPRECATED_FIELDS {
if object.contains_key(deprecated.name) {
result.warnings.push(ConfigDiagnostic {
path: path_display.clone(),
field: deprecated.name.to_string(),
line: find_key_line(source, deprecated.name),
kind: DiagnosticKind::Deprecated {
replacement: deprecated.replacement,
},
});
}
}
// Validate known nested objects.
if let Some(hooks) = object.get("hooks").and_then(JsonValue::as_object) {
result.merge(validate_object_keys(
hooks,
HOOKS_FIELDS,
"hooks",
source,
&path_display,
));
}
if let Some(permissions) = object.get("permissions").and_then(JsonValue::as_object) {
result.merge(validate_object_keys(
permissions,
PERMISSIONS_FIELDS,
"permissions",
source,
&path_display,
));
}
if let Some(plugins) = object.get("plugins").and_then(JsonValue::as_object) {
result.merge(validate_object_keys(
plugins,
PLUGINS_FIELDS,
"plugins",
source,
&path_display,
));
}
if let Some(sandbox) = object.get("sandbox").and_then(JsonValue::as_object) {
result.merge(validate_object_keys(
sandbox,
SANDBOX_FIELDS,
"sandbox",
source,
&path_display,
));
}
if let Some(oauth) = object.get("oauth").and_then(JsonValue::as_object) {
result.merge(validate_object_keys(
oauth,
OAUTH_FIELDS,
"oauth",
source,
&path_display,
));
}
result
}
/// Check whether a file path uses an unsupported config format (e.g. TOML).
pub fn check_unsupported_format(file_path: &Path) -> Result<(), ConfigError> {
if let Some(ext) = file_path.extension().and_then(|e| e.to_str()) {
if ext.eq_ignore_ascii_case("toml") {
return Err(ConfigError::Parse(format!(
"{}: TOML config files are not supported. Use JSON (settings.json) instead",
file_path.display()
)));
}
}
Ok(())
}
/// Format all diagnostics into a human-readable report.
#[must_use]
pub fn format_diagnostics(result: &ValidationResult) -> String {
let mut lines = Vec::new();
for warning in &result.warnings {
lines.push(format!("warning: {warning}"));
}
for error in &result.errors {
lines.push(format!("error: {error}"));
}
lines.join("\n")
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
fn test_path() -> PathBuf {
PathBuf::from("/test/settings.json")
}
#[test]
fn detects_unknown_top_level_key() {
// given
let source = r#"{"model": "opus", "unknownField": true}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].field, "unknownField");
assert!(matches!(
result.errors[0].kind,
DiagnosticKind::UnknownKey { .. }
));
}
#[test]
fn detects_wrong_type_for_model() {
// given
let source = r#"{"model": 123}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].field, "model");
assert!(matches!(
result.errors[0].kind,
DiagnosticKind::WrongType {
expected: "a string",
got: "a number"
}
));
}
#[test]
fn detects_deprecated_permission_mode() {
// given
let source = r#"{"permissionMode": "plan"}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.warnings.len(), 1);
assert_eq!(result.warnings[0].field, "permissionMode");
assert!(matches!(
result.warnings[0].kind,
DiagnosticKind::Deprecated {
replacement: "permissions.defaultMode"
}
));
}
#[test]
fn detects_deprecated_enabled_plugins() {
// given
let source = r#"{"enabledPlugins": {"tool-guard@builtin": true}}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.warnings.len(), 1);
assert_eq!(result.warnings[0].field, "enabledPlugins");
assert!(matches!(
result.warnings[0].kind,
DiagnosticKind::Deprecated {
replacement: "plugins.enabled"
}
));
}
#[test]
fn reports_line_number_for_unknown_key() {
// given
let source = "{\n \"model\": \"opus\",\n \"badKey\": true\n}";
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].line, Some(3));
assert_eq!(result.errors[0].field, "badKey");
}
#[test]
fn reports_line_number_for_wrong_type() {
// given
let source = "{\n \"model\": 42\n}";
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].line, Some(2));
}
#[test]
fn validates_nested_hooks_keys() {
// given
let source = r#"{"hooks": {"PreToolUse": ["cmd"], "BadHook": ["x"]}}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].field, "hooks.BadHook");
}
#[test]
fn validates_nested_permissions_keys() {
// given
let source = r#"{"permissions": {"allow": ["Read"], "denyAll": true}}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].field, "permissions.denyAll");
}
#[test]
fn validates_nested_sandbox_keys() {
// given
let source = r#"{"sandbox": {"enabled": true, "containerMode": "strict"}}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].field, "sandbox.containerMode");
}
#[test]
fn validates_nested_plugins_keys() {
// given
let source = r#"{"plugins": {"installRoot": "/tmp", "autoUpdate": true}}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].field, "plugins.autoUpdate");
}
#[test]
fn validates_nested_oauth_keys() {
// given
let source = r#"{"oauth": {"clientId": "abc", "secret": "hidden"}}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].field, "oauth.secret");
}
#[test]
fn valid_config_produces_no_diagnostics() {
// given
let source = r#"{
"model": "opus",
"hooks": {"PreToolUse": ["guard"]},
"permissions": {"defaultMode": "plan", "allow": ["Read"]},
"mcpServers": {},
"sandbox": {"enabled": false}
}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert!(result.is_ok());
assert!(result.warnings.is_empty());
}
#[test]
fn suggests_close_field_name() {
// given
let source = r#"{"modle": "opus"}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
match &result.errors[0].kind {
DiagnosticKind::UnknownKey {
suggestion: Some(s),
} => assert_eq!(s, "model"),
other => panic!("expected suggestion, got {other:?}"),
}
}
#[test]
fn format_diagnostics_includes_all_entries() {
// given
let source = r#"{"permissionMode": "plan", "badKey": 1}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
let result = validate_config_file(object, source, &test_path());
// when
let output = format_diagnostics(&result);
// then
assert!(output.contains("warning:"));
assert!(output.contains("error:"));
assert!(output.contains("badKey"));
assert!(output.contains("permissionMode"));
}
#[test]
fn check_unsupported_format_rejects_toml() {
// given
let path = PathBuf::from("/home/.claw/settings.toml");
// when
let result = check_unsupported_format(&path);
// then
assert!(result.is_err());
let message = result.unwrap_err().to_string();
assert!(message.contains("TOML"));
assert!(message.contains("settings.toml"));
}
#[test]
fn check_unsupported_format_allows_json() {
// given
let path = PathBuf::from("/home/.claw/settings.json");
// when / then
assert!(check_unsupported_format(&path).is_ok());
}
#[test]
fn wrong_type_in_nested_sandbox_field() {
// given
let source = r#"{"sandbox": {"enabled": "yes"}}"#;
let parsed = JsonValue::parse(source).expect("valid json");
let object = parsed.as_object().expect("object");
// when
let result = validate_config_file(object, source, &test_path());
// then
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].field, "sandbox.enabled");
assert!(matches!(
result.errors[0].kind,
DiagnosticKind::WrongType {
expected: "a boolean",
got: "a string"
}
));
}
#[test]
fn display_format_unknown_key_with_line() {
// given
let diag = ConfigDiagnostic {
path: "/test/settings.json".to_string(),
field: "badKey".to_string(),
line: Some(5),
kind: DiagnosticKind::UnknownKey { suggestion: None },
};
// when
let output = diag.to_string();
// then
assert_eq!(
output,
r#"/test/settings.json: unknown key "badKey" (line 5)"#
);
}
#[test]
fn display_format_wrong_type_with_line() {
// given
let diag = ConfigDiagnostic {
path: "/test/settings.json".to_string(),
field: "model".to_string(),
line: Some(2),
kind: DiagnosticKind::WrongType {
expected: "a string",
got: "a number",
},
};
// when
let output = diag.to_string();
// then
assert_eq!(
output,
r#"/test/settings.json: field "model" must be a string, got a number (line 2)"#
);
}
#[test]
fn display_format_deprecated_with_line() {
// given
let diag = ConfigDiagnostic {
path: "/test/settings.json".to_string(),
field: "permissionMode".to_string(),
line: Some(3),
kind: DiagnosticKind::Deprecated {
replacement: "permissions.defaultMode",
},
};
// when
let output = diag.to_string();
// then
assert_eq!(
output,
r#"/test/settings.json: field "permissionMode" is deprecated (line 3). Use "permissions.defaultMode" instead"#
);
}
}
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,839 @@
use std::cmp::Reverse;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::time::Instant;
use glob::Pattern;
use regex::RegexBuilder;
use serde::{Deserialize, Serialize};
use walkdir::WalkDir;
/// Maximum file size that can be read (10 MB).
const MAX_READ_SIZE: u64 = 10 * 1024 * 1024;
/// Maximum file size that can be written (10 MB).
const MAX_WRITE_SIZE: usize = 10 * 1024 * 1024;
/// Check whether a file appears to contain binary content by examining
/// the first chunk for NUL bytes.
fn is_binary_file(path: &Path) -> io::Result<bool> {
use std::io::Read;
let mut file = fs::File::open(path)?;
let mut buffer = [0u8; 8192];
let bytes_read = file.read(&mut buffer)?;
Ok(buffer[..bytes_read].contains(&0))
}
/// Validate that a resolved path stays within the given workspace root.
/// Returns the canonical path on success, or an error if the path escapes
/// the workspace boundary (e.g. via `../` traversal or symlink).
#[allow(dead_code)]
fn validate_workspace_boundary(resolved: &Path, workspace_root: &Path) -> io::Result<()> {
if !resolved.starts_with(workspace_root) {
return Err(io::Error::new(
io::ErrorKind::PermissionDenied,
format!(
"path {} escapes workspace boundary {}",
resolved.display(),
workspace_root.display()
),
));
}
Ok(())
}
/// Text payload returned by file-reading operations.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct TextFilePayload {
#[serde(rename = "filePath")]
pub file_path: String,
pub content: String,
#[serde(rename = "numLines")]
pub num_lines: usize,
#[serde(rename = "startLine")]
pub start_line: usize,
#[serde(rename = "totalLines")]
pub total_lines: usize,
}
/// Output envelope for the `read_file` tool.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ReadFileOutput {
#[serde(rename = "type")]
pub kind: String,
pub file: TextFilePayload,
}
/// Structured patch hunk emitted by write and edit operations.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct StructuredPatchHunk {
#[serde(rename = "oldStart")]
pub old_start: usize,
#[serde(rename = "oldLines")]
pub old_lines: usize,
#[serde(rename = "newStart")]
pub new_start: usize,
#[serde(rename = "newLines")]
pub new_lines: usize,
pub lines: Vec<String>,
}
/// Output envelope for full-file write operations.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WriteFileOutput {
#[serde(rename = "type")]
pub kind: String,
#[serde(rename = "filePath")]
pub file_path: String,
pub content: String,
#[serde(rename = "structuredPatch")]
pub structured_patch: Vec<StructuredPatchHunk>,
#[serde(rename = "originalFile")]
pub original_file: Option<String>,
#[serde(rename = "gitDiff")]
pub git_diff: Option<serde_json::Value>,
}
/// Output envelope for targeted string-replacement edits.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct EditFileOutput {
#[serde(rename = "filePath")]
pub file_path: String,
#[serde(rename = "oldString")]
pub old_string: String,
#[serde(rename = "newString")]
pub new_string: String,
#[serde(rename = "originalFile")]
pub original_file: String,
#[serde(rename = "structuredPatch")]
pub structured_patch: Vec<StructuredPatchHunk>,
#[serde(rename = "userModified")]
pub user_modified: bool,
#[serde(rename = "replaceAll")]
pub replace_all: bool,
#[serde(rename = "gitDiff")]
pub git_diff: Option<serde_json::Value>,
}
/// Result of a glob-based filename search.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct GlobSearchOutput {
#[serde(rename = "durationMs")]
pub duration_ms: u128,
#[serde(rename = "numFiles")]
pub num_files: usize,
pub filenames: Vec<String>,
pub truncated: bool,
}
/// Parameters accepted by the grep-style search tool.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct GrepSearchInput {
pub pattern: String,
pub path: Option<String>,
pub glob: Option<String>,
#[serde(rename = "output_mode")]
pub output_mode: Option<String>,
#[serde(rename = "-B")]
pub before: Option<usize>,
#[serde(rename = "-A")]
pub after: Option<usize>,
#[serde(rename = "-C")]
pub context_short: Option<usize>,
pub context: Option<usize>,
#[serde(rename = "-n")]
pub line_numbers: Option<bool>,
#[serde(rename = "-i")]
pub case_insensitive: Option<bool>,
#[serde(rename = "type")]
pub file_type: Option<String>,
pub head_limit: Option<usize>,
pub offset: Option<usize>,
pub multiline: Option<bool>,
}
/// Result payload returned by the grep-style search tool.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct GrepSearchOutput {
pub mode: Option<String>,
#[serde(rename = "numFiles")]
pub num_files: usize,
pub filenames: Vec<String>,
pub content: Option<String>,
#[serde(rename = "numLines")]
pub num_lines: Option<usize>,
#[serde(rename = "numMatches")]
pub num_matches: Option<usize>,
#[serde(rename = "appliedLimit")]
pub applied_limit: Option<usize>,
#[serde(rename = "appliedOffset")]
pub applied_offset: Option<usize>,
}
/// Reads a text file and returns a line-windowed payload.
pub fn read_file(
path: &str,
offset: Option<usize>,
limit: Option<usize>,
) -> io::Result<ReadFileOutput> {
let absolute_path = normalize_path(path)?;
// Check file size before reading
let metadata = fs::metadata(&absolute_path)?;
if metadata.len() > MAX_READ_SIZE {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"file is too large ({} bytes, max {} bytes)",
metadata.len(),
MAX_READ_SIZE
),
));
}
// Detect binary files
if is_binary_file(&absolute_path)? {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"file appears to be binary",
));
}
let content = fs::read_to_string(&absolute_path)?;
let lines: Vec<&str> = content.lines().collect();
let start_index = offset.unwrap_or(0).min(lines.len());
let end_index = limit.map_or(lines.len(), |limit| {
start_index.saturating_add(limit).min(lines.len())
});
let selected = lines[start_index..end_index].join("\n");
Ok(ReadFileOutput {
kind: String::from("text"),
file: TextFilePayload {
file_path: absolute_path.to_string_lossy().into_owned(),
content: selected,
num_lines: end_index.saturating_sub(start_index),
start_line: start_index.saturating_add(1),
total_lines: lines.len(),
},
})
}
/// Replaces a file's contents and returns patch metadata.
pub fn write_file(path: &str, content: &str) -> io::Result<WriteFileOutput> {
if content.len() > MAX_WRITE_SIZE {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"content is too large ({} bytes, max {} bytes)",
content.len(),
MAX_WRITE_SIZE
),
));
}
let absolute_path = normalize_path_allow_missing(path)?;
let original_file = fs::read_to_string(&absolute_path).ok();
if let Some(parent) = absolute_path.parent() {
fs::create_dir_all(parent)?;
}
fs::write(&absolute_path, content)?;
Ok(WriteFileOutput {
kind: if original_file.is_some() {
String::from("update")
} else {
String::from("create")
},
file_path: absolute_path.to_string_lossy().into_owned(),
content: content.to_owned(),
structured_patch: make_patch(original_file.as_deref().unwrap_or(""), content),
original_file,
git_diff: None,
})
}
/// Performs an in-file string replacement and returns patch metadata.
pub fn edit_file(
path: &str,
old_string: &str,
new_string: &str,
replace_all: bool,
) -> io::Result<EditFileOutput> {
let absolute_path = normalize_path(path)?;
let original_file = fs::read_to_string(&absolute_path)?;
if old_string == new_string {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"old_string and new_string must differ",
));
}
if !original_file.contains(old_string) {
return Err(io::Error::new(
io::ErrorKind::NotFound,
"old_string not found in file",
));
}
let updated = if replace_all {
original_file.replace(old_string, new_string)
} else {
original_file.replacen(old_string, new_string, 1)
};
fs::write(&absolute_path, &updated)?;
Ok(EditFileOutput {
file_path: absolute_path.to_string_lossy().into_owned(),
old_string: old_string.to_owned(),
new_string: new_string.to_owned(),
original_file: original_file.clone(),
structured_patch: make_patch(&original_file, &updated),
user_modified: false,
replace_all,
git_diff: None,
})
}
/// Expands a glob pattern and returns matching filenames.
pub fn glob_search(pattern: &str, path: Option<&str>) -> io::Result<GlobSearchOutput> {
let started = Instant::now();
let base_dir = path
.map(normalize_path)
.transpose()?
.unwrap_or(std::env::current_dir()?);
let search_pattern = if Path::new(pattern).is_absolute() {
pattern.to_owned()
} else {
base_dir.join(pattern).to_string_lossy().into_owned()
};
// The `glob` crate does not support brace expansion ({a,b,c}).
// Expand braces into multiple patterns so patterns like
// `Assets/**/*.{cs,uxml,uss}` work correctly.
let expanded = expand_braces(&search_pattern);
let mut seen = std::collections::HashSet::new();
let mut matches = Vec::new();
for pat in &expanded {
let entries = glob::glob(pat)
.map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error.to_string()))?;
for entry in entries.flatten() {
if entry.is_file() && seen.insert(entry.clone()) {
matches.push(entry);
}
}
}
matches.sort_by_key(|path| {
fs::metadata(path)
.and_then(|metadata| metadata.modified())
.ok()
.map(Reverse)
});
let truncated = matches.len() > 100;
let filenames = matches
.into_iter()
.take(100)
.map(|path| path.to_string_lossy().into_owned())
.collect::<Vec<_>>();
Ok(GlobSearchOutput {
duration_ms: started.elapsed().as_millis(),
num_files: filenames.len(),
filenames,
truncated,
})
}
/// Runs a regex search over workspace files with optional context lines.
pub fn grep_search(input: &GrepSearchInput) -> io::Result<GrepSearchOutput> {
let base_path = input
.path
.as_deref()
.map(normalize_path)
.transpose()?
.unwrap_or(std::env::current_dir()?);
let regex = RegexBuilder::new(&input.pattern)
.case_insensitive(input.case_insensitive.unwrap_or(false))
.dot_matches_new_line(input.multiline.unwrap_or(false))
.build()
.map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error.to_string()))?;
let glob_filter = input
.glob
.as_deref()
.map(Pattern::new)
.transpose()
.map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error.to_string()))?;
let file_type = input.file_type.as_deref();
let output_mode = input
.output_mode
.clone()
.unwrap_or_else(|| String::from("files_with_matches"));
let context = input.context.or(input.context_short).unwrap_or(0);
let mut filenames = Vec::new();
let mut content_lines = Vec::new();
let mut total_matches = 0usize;
for file_path in collect_search_files(&base_path)? {
if !matches_optional_filters(&file_path, glob_filter.as_ref(), file_type) {
continue;
}
let Ok(file_contents) = fs::read_to_string(&file_path) else {
continue;
};
if output_mode == "count" {
let count = regex.find_iter(&file_contents).count();
if count > 0 {
filenames.push(file_path.to_string_lossy().into_owned());
total_matches += count;
}
continue;
}
let lines: Vec<&str> = file_contents.lines().collect();
let mut matched_lines = Vec::new();
for (index, line) in lines.iter().enumerate() {
if regex.is_match(line) {
total_matches += 1;
matched_lines.push(index);
}
}
if matched_lines.is_empty() {
continue;
}
filenames.push(file_path.to_string_lossy().into_owned());
if output_mode == "content" {
for index in matched_lines {
let start = index.saturating_sub(input.before.unwrap_or(context));
let end = (index + input.after.unwrap_or(context) + 1).min(lines.len());
for (current, line) in lines.iter().enumerate().take(end).skip(start) {
let prefix = if input.line_numbers.unwrap_or(true) {
format!("{}:{}:", file_path.to_string_lossy(), current + 1)
} else {
format!("{}:", file_path.to_string_lossy())
};
content_lines.push(format!("{prefix}{line}"));
}
}
}
}
let (filenames, applied_limit, applied_offset) =
apply_limit(filenames, input.head_limit, input.offset);
let content_output = if output_mode == "content" {
let (lines, limit, offset) = apply_limit(content_lines, input.head_limit, input.offset);
return Ok(GrepSearchOutput {
mode: Some(output_mode),
num_files: filenames.len(),
filenames,
num_lines: Some(lines.len()),
content: Some(lines.join("\n")),
num_matches: None,
applied_limit: limit,
applied_offset: offset,
});
} else {
None
};
Ok(GrepSearchOutput {
mode: Some(output_mode.clone()),
num_files: filenames.len(),
filenames,
content: content_output,
num_lines: None,
num_matches: (output_mode == "count").then_some(total_matches),
applied_limit,
applied_offset,
})
}
fn collect_search_files(base_path: &Path) -> io::Result<Vec<PathBuf>> {
if base_path.is_file() {
return Ok(vec![base_path.to_path_buf()]);
}
let mut files = Vec::new();
for entry in WalkDir::new(base_path) {
let entry = entry.map_err(|error| io::Error::other(error.to_string()))?;
if entry.file_type().is_file() {
files.push(entry.path().to_path_buf());
}
}
Ok(files)
}
fn matches_optional_filters(
path: &Path,
glob_filter: Option<&Pattern>,
file_type: Option<&str>,
) -> bool {
if let Some(glob_filter) = glob_filter {
let path_string = path.to_string_lossy();
if !glob_filter.matches(&path_string) && !glob_filter.matches_path(path) {
return false;
}
}
if let Some(file_type) = file_type {
let extension = path.extension().and_then(|extension| extension.to_str());
if extension != Some(file_type) {
return false;
}
}
true
}
fn apply_limit<T>(
items: Vec<T>,
limit: Option<usize>,
offset: Option<usize>,
) -> (Vec<T>, Option<usize>, Option<usize>) {
let offset_value = offset.unwrap_or(0);
let mut items = items.into_iter().skip(offset_value).collect::<Vec<_>>();
let explicit_limit = limit.unwrap_or(250);
if explicit_limit == 0 {
return (items, None, (offset_value > 0).then_some(offset_value));
}
let truncated = items.len() > explicit_limit;
items.truncate(explicit_limit);
(
items,
truncated.then_some(explicit_limit),
(offset_value > 0).then_some(offset_value),
)
}
fn make_patch(original: &str, updated: &str) -> Vec<StructuredPatchHunk> {
let mut lines = Vec::new();
for line in original.lines() {
lines.push(format!("-{line}"));
}
for line in updated.lines() {
lines.push(format!("+{line}"));
}
vec![StructuredPatchHunk {
old_start: 1,
old_lines: original.lines().count(),
new_start: 1,
new_lines: updated.lines().count(),
lines,
}]
}
fn normalize_path(path: &str) -> io::Result<PathBuf> {
let candidate = if Path::new(path).is_absolute() {
PathBuf::from(path)
} else {
std::env::current_dir()?.join(path)
};
candidate.canonicalize()
}
fn normalize_path_allow_missing(path: &str) -> io::Result<PathBuf> {
let candidate = if Path::new(path).is_absolute() {
PathBuf::from(path)
} else {
std::env::current_dir()?.join(path)
};
if let Ok(canonical) = candidate.canonicalize() {
return Ok(canonical);
}
if let Some(parent) = candidate.parent() {
let canonical_parent = parent
.canonicalize()
.unwrap_or_else(|_| parent.to_path_buf());
if let Some(name) = candidate.file_name() {
return Ok(canonical_parent.join(name));
}
}
Ok(candidate)
}
/// Read a file with workspace boundary enforcement.
#[allow(dead_code)]
pub fn read_file_in_workspace(
path: &str,
offset: Option<usize>,
limit: Option<usize>,
workspace_root: &Path,
) -> io::Result<ReadFileOutput> {
let absolute_path = normalize_path(path)?;
let canonical_root = workspace_root
.canonicalize()
.unwrap_or_else(|_| workspace_root.to_path_buf());
validate_workspace_boundary(&absolute_path, &canonical_root)?;
read_file(path, offset, limit)
}
/// Write a file with workspace boundary enforcement.
#[allow(dead_code)]
pub fn write_file_in_workspace(
path: &str,
content: &str,
workspace_root: &Path,
) -> io::Result<WriteFileOutput> {
let absolute_path = normalize_path_allow_missing(path)?;
let canonical_root = workspace_root
.canonicalize()
.unwrap_or_else(|_| workspace_root.to_path_buf());
validate_workspace_boundary(&absolute_path, &canonical_root)?;
write_file(path, content)
}
/// Edit a file with workspace boundary enforcement.
#[allow(dead_code)]
pub fn edit_file_in_workspace(
path: &str,
old_string: &str,
new_string: &str,
replace_all: bool,
workspace_root: &Path,
) -> io::Result<EditFileOutput> {
let absolute_path = normalize_path(path)?;
let canonical_root = workspace_root
.canonicalize()
.unwrap_or_else(|_| workspace_root.to_path_buf());
validate_workspace_boundary(&absolute_path, &canonical_root)?;
edit_file(path, old_string, new_string, replace_all)
}
/// Check whether a path is a symlink that resolves outside the workspace.
#[allow(dead_code)]
pub fn is_symlink_escape(path: &Path, workspace_root: &Path) -> io::Result<bool> {
let metadata = fs::symlink_metadata(path)?;
if !metadata.is_symlink() {
return Ok(false);
}
let resolved = path.canonicalize()?;
let canonical_root = workspace_root
.canonicalize()
.unwrap_or_else(|_| workspace_root.to_path_buf());
Ok(!resolved.starts_with(&canonical_root))
}
/// Expand shell-style brace groups in a glob pattern.
///
/// Handles one level of braces: `foo.{a,b,c}` → `["foo.a", "foo.b", "foo.c"]`.
/// Nested braces are not expanded (uncommon in practice).
/// Patterns without braces pass through unchanged.
fn expand_braces(pattern: &str) -> Vec<String> {
let Some(open) = pattern.find('{') else {
return vec![pattern.to_owned()];
};
let Some(close) = pattern[open..].find('}').map(|i| open + i) else {
// Unmatched brace — treat as literal.
return vec![pattern.to_owned()];
};
let prefix = &pattern[..open];
let suffix = &pattern[close + 1..];
let alternatives = &pattern[open + 1..close];
alternatives
.split(',')
.flat_map(|alt| expand_braces(&format!("{prefix}{alt}{suffix}")))
.collect()
}
#[cfg(test)]
mod tests {
use std::time::{SystemTime, UNIX_EPOCH};
use super::{
edit_file, expand_braces, glob_search, grep_search, is_symlink_escape, read_file,
read_file_in_workspace, write_file, GrepSearchInput, MAX_WRITE_SIZE,
};
fn temp_path(name: &str) -> std::path::PathBuf {
let unique = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("time should move forward")
.as_nanos();
std::env::temp_dir().join(format!("clawd-native-{name}-{unique}"))
}
#[test]
fn reads_and_writes_files() {
let path = temp_path("read-write.txt");
let write_output = write_file(path.to_string_lossy().as_ref(), "one\ntwo\nthree")
.expect("write should succeed");
assert_eq!(write_output.kind, "create");
let read_output = read_file(path.to_string_lossy().as_ref(), Some(1), Some(1))
.expect("read should succeed");
assert_eq!(read_output.file.content, "two");
}
#[test]
fn edits_file_contents() {
let path = temp_path("edit.txt");
write_file(path.to_string_lossy().as_ref(), "alpha beta alpha")
.expect("initial write should succeed");
let output = edit_file(path.to_string_lossy().as_ref(), "alpha", "omega", true)
.expect("edit should succeed");
assert!(output.replace_all);
}
#[test]
fn rejects_binary_files() {
let path = temp_path("binary-test.bin");
std::fs::write(&path, b"\x00\x01\x02\x03binary content").expect("write should succeed");
let result = read_file(path.to_string_lossy().as_ref(), None, None);
assert!(result.is_err());
let error = result.unwrap_err();
assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
assert!(error.to_string().contains("binary"));
}
#[test]
fn rejects_oversized_writes() {
let path = temp_path("oversize-write.txt");
let huge = "x".repeat(MAX_WRITE_SIZE + 1);
let result = write_file(path.to_string_lossy().as_ref(), &huge);
assert!(result.is_err());
let error = result.unwrap_err();
assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
assert!(error.to_string().contains("too large"));
}
#[test]
fn enforces_workspace_boundary() {
let workspace = temp_path("workspace-boundary");
std::fs::create_dir_all(&workspace).expect("workspace dir should be created");
let inside = workspace.join("inside.txt");
write_file(inside.to_string_lossy().as_ref(), "safe content")
.expect("write inside workspace should succeed");
// Reading inside workspace should succeed
let result =
read_file_in_workspace(inside.to_string_lossy().as_ref(), None, None, &workspace);
assert!(result.is_ok());
// Reading outside workspace should fail
let outside = temp_path("outside-boundary.txt");
write_file(outside.to_string_lossy().as_ref(), "unsafe content")
.expect("write outside should succeed");
let result =
read_file_in_workspace(outside.to_string_lossy().as_ref(), None, None, &workspace);
assert!(result.is_err());
let error = result.unwrap_err();
assert_eq!(error.kind(), std::io::ErrorKind::PermissionDenied);
assert!(error.to_string().contains("escapes workspace"));
}
#[test]
fn detects_symlink_escape() {
let workspace = temp_path("symlink-workspace");
std::fs::create_dir_all(&workspace).expect("workspace dir should be created");
let outside = temp_path("symlink-target.txt");
std::fs::write(&outside, "target content").expect("target should write");
let link_path = workspace.join("escape-link.txt");
#[cfg(unix)]
{
std::os::unix::fs::symlink(&outside, &link_path).expect("symlink should create");
assert!(is_symlink_escape(&link_path, &workspace).expect("check should succeed"));
}
// Non-symlink file should not be an escape
let normal = workspace.join("normal.txt");
std::fs::write(&normal, "normal content").expect("normal file should write");
assert!(!is_symlink_escape(&normal, &workspace).expect("check should succeed"));
}
#[test]
fn globs_and_greps_directory() {
let dir = temp_path("search-dir");
std::fs::create_dir_all(&dir).expect("directory should be created");
let file = dir.join("demo.rs");
write_file(
file.to_string_lossy().as_ref(),
"fn main() {\n println!(\"hello\");\n}\n",
)
.expect("file write should succeed");
let globbed = glob_search("**/*.rs", Some(dir.to_string_lossy().as_ref()))
.expect("glob should succeed");
assert_eq!(globbed.num_files, 1);
let grep_output = grep_search(&GrepSearchInput {
pattern: String::from("hello"),
path: Some(dir.to_string_lossy().into_owned()),
glob: Some(String::from("**/*.rs")),
output_mode: Some(String::from("content")),
before: None,
after: None,
context_short: None,
context: None,
line_numbers: Some(true),
case_insensitive: Some(false),
file_type: None,
head_limit: Some(10),
offset: Some(0),
multiline: Some(false),
})
.expect("grep should succeed");
assert!(grep_output.content.unwrap_or_default().contains("hello"));
}
#[test]
fn expand_braces_no_braces() {
assert_eq!(expand_braces("*.rs"), vec!["*.rs"]);
}
#[test]
fn expand_braces_single_group() {
let mut result = expand_braces("Assets/**/*.{cs,uxml,uss}");
result.sort();
assert_eq!(
result,
vec!["Assets/**/*.cs", "Assets/**/*.uss", "Assets/**/*.uxml",]
);
}
#[test]
fn expand_braces_nested() {
let mut result = expand_braces("src/{a,b}.{rs,toml}");
result.sort();
assert_eq!(
result,
vec!["src/a.rs", "src/a.toml", "src/b.rs", "src/b.toml"]
);
}
#[test]
fn expand_braces_unmatched() {
assert_eq!(expand_braces("foo.{bar"), vec!["foo.{bar"]);
}
#[test]
fn glob_search_with_braces_finds_files() {
let dir = temp_path("glob-braces");
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("a.rs"), "fn main() {}").unwrap();
std::fs::write(dir.join("b.toml"), "[package]").unwrap();
std::fs::write(dir.join("c.txt"), "hello").unwrap();
let result =
glob_search("*.{rs,toml}", Some(dir.to_str().unwrap())).expect("glob should succeed");
assert_eq!(
result.num_files, 2,
"should match .rs and .toml but not .txt"
);
let _ = std::fs::remove_dir_all(&dir);
}
}
@@ -0,0 +1,324 @@
use std::path::Path;
use std::process::Command;
/// A single git commit entry from the log.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GitCommitEntry {
pub hash: String,
pub subject: String,
}
/// Git-aware context gathered at startup for injection into the system prompt.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GitContext {
pub branch: Option<String>,
pub recent_commits: Vec<GitCommitEntry>,
pub staged_files: Vec<String>,
}
const MAX_RECENT_COMMITS: usize = 5;
impl GitContext {
/// Detect the git context from the given working directory.
///
/// Returns `None` when the directory is not inside a git repository.
#[must_use]
pub fn detect(cwd: &Path) -> Option<Self> {
// Quick gate: is this a git repo at all?
let rev_parse = Command::new("git")
.args(["rev-parse", "--is-inside-work-tree"])
.current_dir(cwd)
.output()
.ok()?;
if !rev_parse.status.success() {
return None;
}
Some(Self {
branch: read_branch(cwd),
recent_commits: read_recent_commits(cwd),
staged_files: read_staged_files(cwd),
})
}
/// Render a human-readable summary suitable for system-prompt injection.
#[must_use]
pub fn render(&self) -> String {
let mut lines = Vec::new();
if let Some(branch) = &self.branch {
lines.push(format!("Git branch: {branch}"));
}
if !self.recent_commits.is_empty() {
lines.push(String::new());
lines.push("Recent commits:".to_string());
for entry in &self.recent_commits {
lines.push(format!(" {} {}", entry.hash, entry.subject));
}
}
if !self.staged_files.is_empty() {
lines.push(String::new());
lines.push("Staged files:".to_string());
for file in &self.staged_files {
lines.push(format!(" {file}"));
}
}
lines.join("\n")
}
}
fn read_branch(cwd: &Path) -> Option<String> {
let output = Command::new("git")
.args(["rev-parse", "--abbrev-ref", "HEAD"])
.current_dir(cwd)
.output()
.ok()?;
if !output.status.success() {
return None;
}
let branch = String::from_utf8(output.stdout).ok()?;
let trimmed = branch.trim();
if trimmed.is_empty() || trimmed == "HEAD" {
None
} else {
Some(trimmed.to_string())
}
}
fn read_recent_commits(cwd: &Path) -> Vec<GitCommitEntry> {
let output = Command::new("git")
.args([
"--no-optional-locks",
"log",
"--oneline",
"-n",
&MAX_RECENT_COMMITS.to_string(),
"--no-decorate",
])
.current_dir(cwd)
.output()
.ok();
let Some(output) = output else {
return Vec::new();
};
if !output.status.success() {
return Vec::new();
}
let stdout = String::from_utf8(output.stdout).unwrap_or_default();
stdout
.lines()
.filter_map(|line| {
let line = line.trim();
if line.is_empty() {
return None;
}
let (hash, subject) = line.split_once(' ')?;
Some(GitCommitEntry {
hash: hash.to_string(),
subject: subject.to_string(),
})
})
.collect()
}
fn read_staged_files(cwd: &Path) -> Vec<String> {
let output = Command::new("git")
.args(["--no-optional-locks", "diff", "--cached", "--name-only"])
.current_dir(cwd)
.output()
.ok();
let Some(output) = output else {
return Vec::new();
};
if !output.status.success() {
return Vec::new();
}
let stdout = String::from_utf8(output.stdout).unwrap_or_default();
stdout
.lines()
.filter(|line| !line.trim().is_empty())
.map(|line| line.trim().to_string())
.collect()
}
#[cfg(test)]
mod tests {
use super::{GitCommitEntry, GitContext};
use std::fs;
use std::process::Command;
use std::time::{SystemTime, UNIX_EPOCH};
fn temp_dir(label: &str) -> std::path::PathBuf {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("time should be after epoch")
.as_nanos();
std::env::temp_dir().join(format!("runtime-git-context-{label}-{nanos}"))
}
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
crate::test_env_lock()
}
fn ensure_valid_cwd() {
if std::env::current_dir().is_err() {
std::env::set_current_dir(env!("CARGO_MANIFEST_DIR"))
.expect("test cwd should be recoverable");
}
}
#[test]
fn returns_none_for_non_git_directory() {
// given
let _guard = env_lock();
ensure_valid_cwd();
let root = temp_dir("non-git");
fs::create_dir_all(&root).expect("create dir");
// when
let context = GitContext::detect(&root);
// then
assert!(context.is_none());
fs::remove_dir_all(root).expect("cleanup");
}
#[test]
fn detects_branch_name_and_commits() {
// given
let _guard = env_lock();
ensure_valid_cwd();
let root = temp_dir("branch-commits");
fs::create_dir_all(&root).expect("create dir");
git(&root, &["init", "--quiet", "--initial-branch=main"]);
git(&root, &["config", "user.email", "[email protected]"]);
git(&root, &["config", "user.name", "Git Context Tests"]);
fs::write(root.join("a.txt"), "a\n").expect("write a");
git(&root, &["add", "a.txt"]);
git(&root, &["commit", "-m", "first commit", "--quiet"]);
fs::write(root.join("b.txt"), "b\n").expect("write b");
git(&root, &["add", "b.txt"]);
git(&root, &["commit", "-m", "second commit", "--quiet"]);
// when
let context = GitContext::detect(&root).expect("should detect git repo");
// then
assert_eq!(context.branch.as_deref(), Some("main"));
assert_eq!(context.recent_commits.len(), 2);
assert_eq!(context.recent_commits[0].subject, "second commit");
assert_eq!(context.recent_commits[1].subject, "first commit");
assert!(context.staged_files.is_empty());
fs::remove_dir_all(root).expect("cleanup");
}
#[test]
fn detects_staged_files() {
// given
let _guard = env_lock();
ensure_valid_cwd();
let root = temp_dir("staged");
fs::create_dir_all(&root).expect("create dir");
git(&root, &["init", "--quiet", "--initial-branch=main"]);
git(&root, &["config", "user.email", "[email protected]"]);
git(&root, &["config", "user.name", "Git Context Tests"]);
fs::write(root.join("init.txt"), "init\n").expect("write init");
git(&root, &["add", "init.txt"]);
git(&root, &["commit", "-m", "initial", "--quiet"]);
fs::write(root.join("staged.txt"), "staged\n").expect("write staged");
git(&root, &["add", "staged.txt"]);
// when
let context = GitContext::detect(&root).expect("should detect git repo");
// then
assert_eq!(context.staged_files, vec!["staged.txt"]);
fs::remove_dir_all(root).expect("cleanup");
}
#[test]
fn render_formats_all_sections() {
// given
let context = GitContext {
branch: Some("feat/test".to_string()),
recent_commits: vec![
GitCommitEntry {
hash: "abc1234".to_string(),
subject: "add feature".to_string(),
},
GitCommitEntry {
hash: "def5678".to_string(),
subject: "fix bug".to_string(),
},
],
staged_files: vec!["src/main.rs".to_string()],
};
// when
let rendered = context.render();
// then
assert!(rendered.contains("Git branch: feat/test"));
assert!(rendered.contains("abc1234 add feature"));
assert!(rendered.contains("def5678 fix bug"));
assert!(rendered.contains("src/main.rs"));
}
#[test]
fn render_omits_empty_sections() {
// given
let context = GitContext {
branch: Some("main".to_string()),
recent_commits: Vec::new(),
staged_files: Vec::new(),
};
// when
let rendered = context.render();
// then
assert!(rendered.contains("Git branch: main"));
assert!(!rendered.contains("Recent commits:"));
assert!(!rendered.contains("Staged files:"));
}
#[test]
fn limits_to_five_recent_commits() {
// given
let _guard = env_lock();
ensure_valid_cwd();
let root = temp_dir("five-commits");
fs::create_dir_all(&root).expect("create dir");
git(&root, &["init", "--quiet", "--initial-branch=main"]);
git(&root, &["config", "user.email", "[email protected]"]);
git(&root, &["config", "user.name", "Git Context Tests"]);
for i in 1..=8 {
let name = format!("file{i}.txt");
fs::write(root.join(&name), format!("{i}\n")).expect("write file");
git(&root, &["add", &name]);
git(&root, &["commit", "-m", &format!("commit {i}"), "--quiet"]);
}
// when
let context = GitContext::detect(&root).expect("should detect git repo");
// then
assert_eq!(context.recent_commits.len(), 5);
assert_eq!(context.recent_commits[0].subject, "commit 8");
assert_eq!(context.recent_commits[4].subject, "commit 4");
fs::remove_dir_all(root).expect("cleanup");
}
fn git(cwd: &std::path::Path, args: &[&str]) {
let status = Command::new("git")
.args(args)
.current_dir(cwd)
.output()
.unwrap_or_else(|_| panic!("git {args:?} should run"))
.status;
assert!(status.success(), "git {args:?} failed");
}
}
@@ -0,0 +1,152 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum GreenLevel {
TargetedTests,
Package,
Workspace,
MergeReady,
}
impl GreenLevel {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::TargetedTests => "targeted_tests",
Self::Package => "package",
Self::Workspace => "workspace",
Self::MergeReady => "merge_ready",
}
}
}
impl std::fmt::Display for GreenLevel {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct GreenContract {
pub required_level: GreenLevel,
}
impl GreenContract {
#[must_use]
pub fn new(required_level: GreenLevel) -> Self {
Self { required_level }
}
#[must_use]
pub fn evaluate(self, observed_level: Option<GreenLevel>) -> GreenContractOutcome {
match observed_level {
Some(level) if level >= self.required_level => GreenContractOutcome::Satisfied {
required_level: self.required_level,
observed_level: level,
},
_ => GreenContractOutcome::Unsatisfied {
required_level: self.required_level,
observed_level,
},
}
}
#[must_use]
pub fn is_satisfied_by(self, observed_level: GreenLevel) -> bool {
observed_level >= self.required_level
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "outcome", rename_all = "snake_case")]
pub enum GreenContractOutcome {
Satisfied {
required_level: GreenLevel,
observed_level: GreenLevel,
},
Unsatisfied {
required_level: GreenLevel,
observed_level: Option<GreenLevel>,
},
}
impl GreenContractOutcome {
#[must_use]
pub fn is_satisfied(&self) -> bool {
matches!(self, Self::Satisfied { .. })
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn given_matching_level_when_evaluating_contract_then_it_is_satisfied() {
// given
let contract = GreenContract::new(GreenLevel::Package);
// when
let outcome = contract.evaluate(Some(GreenLevel::Package));
// then
assert_eq!(
outcome,
GreenContractOutcome::Satisfied {
required_level: GreenLevel::Package,
observed_level: GreenLevel::Package,
}
);
assert!(outcome.is_satisfied());
}
#[test]
fn given_higher_level_when_checking_requirement_then_it_still_satisfies_contract() {
// given
let contract = GreenContract::new(GreenLevel::TargetedTests);
// when
let is_satisfied = contract.is_satisfied_by(GreenLevel::Workspace);
// then
assert!(is_satisfied);
}
#[test]
fn given_lower_level_when_evaluating_contract_then_it_is_unsatisfied() {
// given
let contract = GreenContract::new(GreenLevel::Workspace);
// when
let outcome = contract.evaluate(Some(GreenLevel::Package));
// then
assert_eq!(
outcome,
GreenContractOutcome::Unsatisfied {
required_level: GreenLevel::Workspace,
observed_level: Some(GreenLevel::Package),
}
);
assert!(!outcome.is_satisfied());
}
#[test]
fn given_no_green_level_when_evaluating_contract_then_contract_is_unsatisfied() {
// given
let contract = GreenContract::new(GreenLevel::MergeReady);
// when
let outcome = contract.evaluate(None);
// then
assert_eq!(
outcome,
GreenContractOutcome::Unsatisfied {
required_level: GreenLevel::MergeReady,
observed_level: None,
}
);
}
}
+987
View File
@@ -0,0 +1,987 @@
use std::ffi::OsStr;
use std::io::Write;
use std::process::{Command, Stdio};
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
use std::thread;
use std::time::Duration;
use serde_json::{json, Value};
use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig};
use crate::permissions::PermissionOverride;
pub type HookPermissionDecision = PermissionOverride;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HookEvent {
PreToolUse,
PostToolUse,
PostToolUseFailure,
}
impl HookEvent {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::PreToolUse => "PreToolUse",
Self::PostToolUse => "PostToolUse",
Self::PostToolUseFailure => "PostToolUseFailure",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HookProgressEvent {
Started {
event: HookEvent,
tool_name: String,
command: String,
},
Completed {
event: HookEvent,
tool_name: String,
command: String,
},
Cancelled {
event: HookEvent,
tool_name: String,
command: String,
},
}
pub trait HookProgressReporter {
fn on_event(&mut self, event: &HookProgressEvent);
}
#[derive(Debug, Clone, Default)]
pub struct HookAbortSignal {
aborted: Arc<AtomicBool>,
}
impl HookAbortSignal {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn abort(&self) {
self.aborted.store(true, Ordering::SeqCst);
}
#[must_use]
pub fn is_aborted(&self) -> bool {
self.aborted.load(Ordering::SeqCst)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HookRunResult {
denied: bool,
failed: bool,
cancelled: bool,
messages: Vec<String>,
permission_override: Option<PermissionOverride>,
permission_reason: Option<String>,
updated_input: Option<String>,
}
impl HookRunResult {
#[must_use]
pub fn allow(messages: Vec<String>) -> Self {
Self {
denied: false,
failed: false,
cancelled: false,
messages,
permission_override: None,
permission_reason: None,
updated_input: None,
}
}
#[must_use]
pub fn is_denied(&self) -> bool {
self.denied
}
#[must_use]
pub fn is_failed(&self) -> bool {
self.failed
}
#[must_use]
pub fn is_cancelled(&self) -> bool {
self.cancelled
}
#[must_use]
pub fn messages(&self) -> &[String] {
&self.messages
}
#[must_use]
pub fn permission_override(&self) -> Option<PermissionOverride> {
self.permission_override
}
#[must_use]
pub fn permission_decision(&self) -> Option<HookPermissionDecision> {
self.permission_override
}
#[must_use]
pub fn permission_reason(&self) -> Option<&str> {
self.permission_reason.as_deref()
}
#[must_use]
pub fn updated_input(&self) -> Option<&str> {
self.updated_input.as_deref()
}
#[must_use]
pub fn updated_input_json(&self) -> Option<&str> {
self.updated_input()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct HookRunner {
config: RuntimeHookConfig,
}
impl HookRunner {
#[must_use]
pub fn new(config: RuntimeHookConfig) -> Self {
Self { config }
}
#[must_use]
pub fn from_feature_config(feature_config: &RuntimeFeatureConfig) -> Self {
Self::new(feature_config.hooks().clone())
}
#[must_use]
pub fn run_pre_tool_use(&self, tool_name: &str, tool_input: &str) -> HookRunResult {
self.run_pre_tool_use_with_context(tool_name, tool_input, None, None)
}
#[must_use]
pub fn run_pre_tool_use_with_context(
&self,
tool_name: &str,
tool_input: &str,
abort_signal: Option<&HookAbortSignal>,
reporter: Option<&mut dyn HookProgressReporter>,
) -> HookRunResult {
Self::run_commands(
HookEvent::PreToolUse,
self.config.pre_tool_use(),
tool_name,
tool_input,
None,
false,
abort_signal,
reporter,
)
}
#[must_use]
pub fn run_pre_tool_use_with_signal(
&self,
tool_name: &str,
tool_input: &str,
abort_signal: Option<&HookAbortSignal>,
) -> HookRunResult {
self.run_pre_tool_use_with_context(tool_name, tool_input, abort_signal, None)
}
#[must_use]
pub fn run_post_tool_use(
&self,
tool_name: &str,
tool_input: &str,
tool_output: &str,
is_error: bool,
) -> HookRunResult {
self.run_post_tool_use_with_context(
tool_name,
tool_input,
tool_output,
is_error,
None,
None,
)
}
#[must_use]
pub fn run_post_tool_use_with_context(
&self,
tool_name: &str,
tool_input: &str,
tool_output: &str,
is_error: bool,
abort_signal: Option<&HookAbortSignal>,
reporter: Option<&mut dyn HookProgressReporter>,
) -> HookRunResult {
Self::run_commands(
HookEvent::PostToolUse,
self.config.post_tool_use(),
tool_name,
tool_input,
Some(tool_output),
is_error,
abort_signal,
reporter,
)
}
#[must_use]
pub fn run_post_tool_use_with_signal(
&self,
tool_name: &str,
tool_input: &str,
tool_output: &str,
is_error: bool,
abort_signal: Option<&HookAbortSignal>,
) -> HookRunResult {
self.run_post_tool_use_with_context(
tool_name,
tool_input,
tool_output,
is_error,
abort_signal,
None,
)
}
#[must_use]
pub fn run_post_tool_use_failure(
&self,
tool_name: &str,
tool_input: &str,
tool_error: &str,
) -> HookRunResult {
self.run_post_tool_use_failure_with_context(tool_name, tool_input, tool_error, None, None)
}
#[must_use]
pub fn run_post_tool_use_failure_with_context(
&self,
tool_name: &str,
tool_input: &str,
tool_error: &str,
abort_signal: Option<&HookAbortSignal>,
reporter: Option<&mut dyn HookProgressReporter>,
) -> HookRunResult {
Self::run_commands(
HookEvent::PostToolUseFailure,
self.config.post_tool_use_failure(),
tool_name,
tool_input,
Some(tool_error),
true,
abort_signal,
reporter,
)
}
#[must_use]
pub fn run_post_tool_use_failure_with_signal(
&self,
tool_name: &str,
tool_input: &str,
tool_error: &str,
abort_signal: Option<&HookAbortSignal>,
) -> HookRunResult {
self.run_post_tool_use_failure_with_context(
tool_name,
tool_input,
tool_error,
abort_signal,
None,
)
}
#[allow(clippy::too_many_arguments)]
fn run_commands(
event: HookEvent,
commands: &[String],
tool_name: &str,
tool_input: &str,
tool_output: Option<&str>,
is_error: bool,
abort_signal: Option<&HookAbortSignal>,
mut reporter: Option<&mut dyn HookProgressReporter>,
) -> HookRunResult {
if commands.is_empty() {
return HookRunResult::allow(Vec::new());
}
if abort_signal.is_some_and(HookAbortSignal::is_aborted) {
return HookRunResult {
denied: false,
failed: false,
cancelled: true,
messages: vec![format!(
"{} hook cancelled before execution",
event.as_str()
)],
permission_override: None,
permission_reason: None,
updated_input: None,
};
}
let payload = hook_payload(event, tool_name, tool_input, tool_output, is_error).to_string();
let mut result = HookRunResult::allow(Vec::new());
for command in commands {
if let Some(reporter) = reporter.as_deref_mut() {
reporter.on_event(&HookProgressEvent::Started {
event,
tool_name: tool_name.to_string(),
command: command.clone(),
});
}
match Self::run_command(
command,
event,
tool_name,
tool_input,
tool_output,
is_error,
&payload,
abort_signal,
) {
HookCommandOutcome::Allow { parsed } => {
if let Some(reporter) = reporter.as_deref_mut() {
reporter.on_event(&HookProgressEvent::Completed {
event,
tool_name: tool_name.to_string(),
command: command.clone(),
});
}
merge_parsed_hook_output(&mut result, parsed);
}
HookCommandOutcome::Deny { parsed } => {
if let Some(reporter) = reporter.as_deref_mut() {
reporter.on_event(&HookProgressEvent::Completed {
event,
tool_name: tool_name.to_string(),
command: command.clone(),
});
}
merge_parsed_hook_output(&mut result, parsed);
result.denied = true;
return result;
}
HookCommandOutcome::Failed { parsed } => {
if let Some(reporter) = reporter.as_deref_mut() {
reporter.on_event(&HookProgressEvent::Completed {
event,
tool_name: tool_name.to_string(),
command: command.clone(),
});
}
merge_parsed_hook_output(&mut result, parsed);
result.failed = true;
return result;
}
HookCommandOutcome::Cancelled { message } => {
if let Some(reporter) = reporter.as_deref_mut() {
reporter.on_event(&HookProgressEvent::Cancelled {
event,
tool_name: tool_name.to_string(),
command: command.clone(),
});
}
result.cancelled = true;
result.messages.push(message);
return result;
}
}
}
result
}
#[allow(clippy::too_many_arguments)]
fn run_command(
command: &str,
event: HookEvent,
tool_name: &str,
tool_input: &str,
tool_output: Option<&str>,
is_error: bool,
payload: &str,
abort_signal: Option<&HookAbortSignal>,
) -> HookCommandOutcome {
let mut child = shell_command(command);
child.stdin(Stdio::piped());
child.stdout(Stdio::piped());
child.stderr(Stdio::piped());
child.env("HOOK_EVENT", event.as_str());
child.env("HOOK_TOOL_NAME", tool_name);
child.env("HOOK_TOOL_INPUT", tool_input);
child.env("HOOK_TOOL_IS_ERROR", if is_error { "1" } else { "0" });
if let Some(tool_output) = tool_output {
child.env("HOOK_TOOL_OUTPUT", tool_output);
}
match child.output_with_stdin(payload.as_bytes(), abort_signal) {
Ok(CommandExecution::Finished(output)) => {
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
let parsed = parse_hook_output(&stdout);
let primary_message = parsed.primary_message().map(ToOwned::to_owned);
match output.status.code() {
Some(0) => {
if parsed.deny {
HookCommandOutcome::Deny { parsed }
} else {
HookCommandOutcome::Allow { parsed }
}
}
Some(2) => HookCommandOutcome::Deny {
parsed: parsed.with_fallback_message(format!(
"{} hook denied tool `{tool_name}`",
event.as_str()
)),
},
Some(code) => HookCommandOutcome::Failed {
parsed: parsed.with_fallback_message(format_hook_failure(
command,
code,
primary_message.as_deref(),
stderr.as_str(),
)),
},
None => HookCommandOutcome::Failed {
parsed: parsed.with_fallback_message(format!(
"{} hook `{command}` terminated by signal while handling `{}`",
event.as_str(),
tool_name
)),
},
}
}
Ok(CommandExecution::Cancelled) => HookCommandOutcome::Cancelled {
message: format!(
"{} hook `{command}` cancelled while handling `{tool_name}`",
event.as_str()
),
},
Err(error) => HookCommandOutcome::Failed {
parsed: ParsedHookOutput {
messages: vec![format!(
"{} hook `{command}` failed to start for `{}`: {error}",
event.as_str(),
tool_name
)],
..ParsedHookOutput::default()
},
},
}
}
}
enum HookCommandOutcome {
Allow { parsed: ParsedHookOutput },
Deny { parsed: ParsedHookOutput },
Failed { parsed: ParsedHookOutput },
Cancelled { message: String },
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
struct ParsedHookOutput {
messages: Vec<String>,
deny: bool,
permission_override: Option<PermissionOverride>,
permission_reason: Option<String>,
updated_input: Option<String>,
}
impl ParsedHookOutput {
fn with_fallback_message(mut self, fallback: String) -> Self {
if self.messages.is_empty() {
self.messages.push(fallback);
}
self
}
fn primary_message(&self) -> Option<&str> {
self.messages.first().map(String::as_str)
}
}
fn merge_parsed_hook_output(target: &mut HookRunResult, parsed: ParsedHookOutput) {
target.messages.extend(parsed.messages);
if parsed.permission_override.is_some() {
target.permission_override = parsed.permission_override;
}
if parsed.permission_reason.is_some() {
target.permission_reason = parsed.permission_reason;
}
if parsed.updated_input.is_some() {
target.updated_input = parsed.updated_input;
}
}
fn parse_hook_output(stdout: &str) -> ParsedHookOutput {
if stdout.is_empty() {
return ParsedHookOutput::default();
}
let Ok(Value::Object(root)) = serde_json::from_str::<Value>(stdout) else {
return ParsedHookOutput {
messages: vec![stdout.to_string()],
..ParsedHookOutput::default()
};
};
let mut parsed = ParsedHookOutput::default();
if let Some(message) = root.get("systemMessage").and_then(Value::as_str) {
parsed.messages.push(message.to_string());
}
if let Some(message) = root.get("reason").and_then(Value::as_str) {
parsed.messages.push(message.to_string());
}
if root.get("continue").and_then(Value::as_bool) == Some(false)
|| root.get("decision").and_then(Value::as_str) == Some("block")
{
parsed.deny = true;
}
if let Some(Value::Object(specific)) = root.get("hookSpecificOutput") {
if let Some(Value::String(additional_context)) = specific.get("additionalContext") {
parsed.messages.push(additional_context.clone());
}
if let Some(decision) = specific.get("permissionDecision").and_then(Value::as_str) {
parsed.permission_override = match decision {
"allow" => Some(PermissionOverride::Allow),
"deny" => Some(PermissionOverride::Deny),
"ask" => Some(PermissionOverride::Ask),
_ => None,
};
}
if let Some(reason) = specific
.get("permissionDecisionReason")
.and_then(Value::as_str)
{
parsed.permission_reason = Some(reason.to_string());
}
if let Some(updated_input) = specific.get("updatedInput") {
parsed.updated_input = serde_json::to_string(updated_input).ok();
}
}
if parsed.messages.is_empty() {
parsed.messages.push(stdout.to_string());
}
parsed
}
fn hook_payload(
event: HookEvent,
tool_name: &str,
tool_input: &str,
tool_output: Option<&str>,
is_error: bool,
) -> Value {
match event {
HookEvent::PostToolUseFailure => json!({
"hook_event_name": event.as_str(),
"tool_name": tool_name,
"tool_input": parse_tool_input(tool_input),
"tool_input_json": tool_input,
"tool_error": tool_output,
"tool_result_is_error": true,
}),
_ => json!({
"hook_event_name": event.as_str(),
"tool_name": tool_name,
"tool_input": parse_tool_input(tool_input),
"tool_input_json": tool_input,
"tool_output": tool_output,
"tool_result_is_error": is_error,
}),
}
}
fn parse_tool_input(tool_input: &str) -> Value {
serde_json::from_str(tool_input).unwrap_or_else(|_| json!({ "raw": tool_input }))
}
fn format_hook_failure(command: &str, code: i32, stdout: Option<&str>, stderr: &str) -> String {
let mut message = format!("Hook `{command}` exited with status {code}");
if let Some(stdout) = stdout.filter(|stdout| !stdout.is_empty()) {
message.push_str(": ");
message.push_str(stdout);
} else if !stderr.is_empty() {
message.push_str(": ");
message.push_str(stderr);
}
message
}
fn shell_command(command: &str) -> CommandWithStdin {
#[cfg(windows)]
let mut command_builder = {
let mut command_builder = Command::new("cmd");
command_builder.arg("/C").arg(command);
CommandWithStdin::new(command_builder)
};
#[cfg(not(windows))]
let command_builder = {
let mut command_builder = Command::new("sh");
command_builder.arg("-lc").arg(command);
CommandWithStdin::new(command_builder)
};
command_builder
}
struct CommandWithStdin {
command: Command,
}
impl CommandWithStdin {
fn new(command: Command) -> Self {
Self { command }
}
fn stdin(&mut self, cfg: Stdio) -> &mut Self {
self.command.stdin(cfg);
self
}
fn stdout(&mut self, cfg: Stdio) -> &mut Self {
self.command.stdout(cfg);
self
}
fn stderr(&mut self, cfg: Stdio) -> &mut Self {
self.command.stderr(cfg);
self
}
fn env<K, V>(&mut self, key: K, value: V) -> &mut Self
where
K: AsRef<OsStr>,
V: AsRef<OsStr>,
{
self.command.env(key, value);
self
}
fn output_with_stdin(
&mut self,
stdin: &[u8],
abort_signal: Option<&HookAbortSignal>,
) -> std::io::Result<CommandExecution> {
let mut child = self.command.spawn()?;
if let Some(mut child_stdin) = child.stdin.take() {
child_stdin.write_all(stdin)?;
}
loop {
if abort_signal.is_some_and(HookAbortSignal::is_aborted) {
let _ = child.kill();
let _ = child.wait_with_output();
return Ok(CommandExecution::Cancelled);
}
match child.try_wait()? {
Some(_) => return child.wait_with_output().map(CommandExecution::Finished),
None => thread::sleep(Duration::from_millis(20)),
}
}
}
}
enum CommandExecution {
Finished(std::process::Output),
Cancelled,
}
#[cfg(test)]
mod tests {
use std::thread;
use std::time::Duration;
use super::{
HookAbortSignal, HookEvent, HookProgressEvent, HookProgressReporter, HookRunResult,
HookRunner,
};
use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig};
use crate::permissions::PermissionOverride;
struct RecordingReporter {
events: Vec<HookProgressEvent>,
}
impl HookProgressReporter for RecordingReporter {
fn on_event(&mut self, event: &HookProgressEvent) {
self.events.push(event.clone());
}
}
#[test]
fn allows_exit_code_zero_and_captures_stdout() {
let runner = HookRunner::new(RuntimeHookConfig::new(
vec![shell_snippet("printf 'pre ok'")],
Vec::new(),
Vec::new(),
));
let result = runner.run_pre_tool_use("Read", r#"{"path":"README.md"}"#);
assert_eq!(result, HookRunResult::allow(vec!["pre ok".to_string()]));
}
#[test]
fn denies_exit_code_two() {
let runner = HookRunner::new(RuntimeHookConfig::new(
vec![shell_snippet("printf 'blocked by hook'; exit 2")],
Vec::new(),
Vec::new(),
));
let result = runner.run_pre_tool_use("Bash", r#"{"command":"pwd"}"#);
assert!(result.is_denied());
assert_eq!(result.messages(), &["blocked by hook".to_string()]);
}
#[test]
fn propagates_other_non_zero_statuses_as_failures() {
let runner = HookRunner::from_feature_config(&RuntimeFeatureConfig::default().with_hooks(
RuntimeHookConfig::new(
vec![shell_snippet("printf 'warning hook'; exit 1")],
Vec::new(),
Vec::new(),
),
));
// given
// when
let result = runner.run_pre_tool_use("Edit", r#"{"file":"src/lib.rs"}"#);
// then
assert!(result.is_failed());
assert!(result
.messages()
.iter()
.any(|message| message.contains("warning hook")));
}
#[test]
fn parses_pre_hook_permission_override_and_updated_input() {
let runner = HookRunner::new(RuntimeHookConfig::new(
vec![shell_snippet(
r#"printf '%s' '{"systemMessage":"updated","hookSpecificOutput":{"permissionDecision":"allow","permissionDecisionReason":"hook ok","updatedInput":{"command":"git status"}}}'"#,
)],
Vec::new(),
Vec::new(),
));
let result = runner.run_pre_tool_use("bash", r#"{"command":"pwd"}"#);
assert_eq!(
result.permission_override(),
Some(PermissionOverride::Allow)
);
assert_eq!(result.permission_reason(), Some("hook ok"));
assert_eq!(result.updated_input(), Some(r#"{"command":"git status"}"#));
assert!(result.messages().iter().any(|message| message == "updated"));
}
#[test]
fn runs_post_tool_use_failure_hooks() {
// given
let runner = HookRunner::new(RuntimeHookConfig::new(
Vec::new(),
Vec::new(),
vec![shell_snippet("printf 'failure hook ran'")],
));
// when
let result =
runner.run_post_tool_use_failure("bash", r#"{"command":"false"}"#, "command failed");
// then
assert!(!result.is_denied());
assert_eq!(result.messages(), &["failure hook ran".to_string()]);
}
#[test]
fn stops_running_failure_hooks_after_failure() {
// given
let runner = HookRunner::new(RuntimeHookConfig::new(
Vec::new(),
Vec::new(),
vec![
shell_snippet("printf 'broken failure hook'; exit 1"),
shell_snippet("printf 'later failure hook'"),
],
));
// when
let result =
runner.run_post_tool_use_failure("bash", r#"{"command":"false"}"#, "command failed");
// then
assert!(result.is_failed());
assert!(result
.messages()
.iter()
.any(|message| message.contains("broken failure hook")));
assert!(!result
.messages()
.iter()
.any(|message| message == "later failure hook"));
}
#[test]
fn executes_hooks_in_configured_order() {
// given
let runner = HookRunner::new(RuntimeHookConfig::new(
vec![
shell_snippet("printf 'first'"),
shell_snippet("printf 'second'"),
],
Vec::new(),
Vec::new(),
));
let mut reporter = RecordingReporter { events: Vec::new() };
// when
let result = runner.run_pre_tool_use_with_context(
"Read",
r#"{"path":"README.md"}"#,
None,
Some(&mut reporter),
);
// then
assert_eq!(
result,
HookRunResult::allow(vec!["first".to_string(), "second".to_string()])
);
assert_eq!(reporter.events.len(), 4);
assert!(matches!(
&reporter.events[0],
HookProgressEvent::Started {
event: HookEvent::PreToolUse,
command,
..
} if command == "printf 'first'"
));
assert!(matches!(
&reporter.events[1],
HookProgressEvent::Completed {
event: HookEvent::PreToolUse,
command,
..
} if command == "printf 'first'"
));
assert!(matches!(
&reporter.events[2],
HookProgressEvent::Started {
event: HookEvent::PreToolUse,
command,
..
} if command == "printf 'second'"
));
assert!(matches!(
&reporter.events[3],
HookProgressEvent::Completed {
event: HookEvent::PreToolUse,
command,
..
} if command == "printf 'second'"
));
}
#[test]
fn stops_running_hooks_after_failure() {
// given
let runner = HookRunner::new(RuntimeHookConfig::new(
vec![
shell_snippet("printf 'broken'; exit 1"),
shell_snippet("printf 'later'"),
],
Vec::new(),
Vec::new(),
));
// when
let result = runner.run_pre_tool_use("Edit", r#"{"file":"src/lib.rs"}"#);
// then
assert!(result.is_failed());
assert!(result
.messages()
.iter()
.any(|message| message.contains("broken")));
assert!(!result.messages().iter().any(|message| message == "later"));
}
#[test]
fn abort_signal_cancels_long_running_hook_and_reports_progress() {
let runner = HookRunner::new(RuntimeHookConfig::new(
vec![shell_snippet("sleep 5")],
Vec::new(),
Vec::new(),
));
let abort_signal = HookAbortSignal::new();
let abort_signal_for_thread = abort_signal.clone();
let mut reporter = RecordingReporter { events: Vec::new() };
thread::spawn(move || {
thread::sleep(Duration::from_millis(100));
abort_signal_for_thread.abort();
});
let result = runner.run_pre_tool_use_with_context(
"bash",
r#"{"command":"sleep 5"}"#,
Some(&abort_signal),
Some(&mut reporter),
);
assert!(result.is_cancelled());
assert!(reporter.events.iter().any(|event| matches!(
event,
HookProgressEvent::Started {
event: HookEvent::PreToolUse,
..
}
)));
assert!(reporter.events.iter().any(|event| matches!(
event,
HookProgressEvent::Cancelled {
event: HookEvent::PreToolUse,
..
}
)));
}
#[cfg(windows)]
fn shell_snippet(script: &str) -> String {
script.replace('\'', "\"")
}
#[cfg(not(windows))]
fn shell_snippet(script: &str) -> String {
script.to_string()
}
}
+358
View File
@@ -0,0 +1,358 @@
use std::collections::BTreeMap;
use std::fmt::{Display, Formatter};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum JsonValue {
Null,
Bool(bool),
Number(i64),
String(String),
Array(Vec<JsonValue>),
Object(BTreeMap<String, JsonValue>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JsonError {
message: String,
}
impl JsonError {
#[must_use]
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl Display for JsonError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.message)
}
}
impl std::error::Error for JsonError {}
impl JsonValue {
#[must_use]
pub fn render(&self) -> String {
match self {
Self::Null => "null".to_string(),
Self::Bool(value) => value.to_string(),
Self::Number(value) => value.to_string(),
Self::String(value) => render_string(value),
Self::Array(values) => {
let rendered = values
.iter()
.map(Self::render)
.collect::<Vec<_>>()
.join(",");
format!("[{rendered}]")
}
Self::Object(entries) => {
let rendered = entries
.iter()
.map(|(key, value)| format!("{}:{}", render_string(key), value.render()))
.collect::<Vec<_>>()
.join(",");
format!("{{{rendered}}}")
}
}
}
pub fn parse(source: &str) -> Result<Self, JsonError> {
let mut parser = Parser::new(source);
let value = parser.parse_value()?;
parser.skip_whitespace();
if parser.is_eof() {
Ok(value)
} else {
Err(JsonError::new("unexpected trailing content"))
}
}
#[must_use]
pub fn as_object(&self) -> Option<&BTreeMap<String, JsonValue>> {
match self {
Self::Object(value) => Some(value),
_ => None,
}
}
#[must_use]
pub fn as_array(&self) -> Option<&[JsonValue]> {
match self {
Self::Array(value) => Some(value),
_ => None,
}
}
#[must_use]
pub fn as_str(&self) -> Option<&str> {
match self {
Self::String(value) => Some(value),
_ => None,
}
}
#[must_use]
pub fn as_bool(&self) -> Option<bool> {
match self {
Self::Bool(value) => Some(*value),
_ => None,
}
}
#[must_use]
pub fn as_i64(&self) -> Option<i64> {
match self {
Self::Number(value) => Some(*value),
_ => None,
}
}
}
fn render_string(value: &str) -> String {
let mut rendered = String::with_capacity(value.len() + 2);
rendered.push('"');
for ch in value.chars() {
match ch {
'"' => rendered.push_str("\\\""),
'\\' => rendered.push_str("\\\\"),
'\n' => rendered.push_str("\\n"),
'\r' => rendered.push_str("\\r"),
'\t' => rendered.push_str("\\t"),
'\u{08}' => rendered.push_str("\\b"),
'\u{0C}' => rendered.push_str("\\f"),
control if control.is_control() => push_unicode_escape(&mut rendered, control),
plain => rendered.push(plain),
}
}
rendered.push('"');
rendered
}
fn push_unicode_escape(rendered: &mut String, control: char) {
const HEX: &[u8; 16] = b"0123456789abcdef";
rendered.push_str("\\u");
let value = u32::from(control);
for shift in [12_u32, 8, 4, 0] {
let nibble = ((value >> shift) & 0xF) as usize;
rendered.push(char::from(HEX[nibble]));
}
}
struct Parser<'a> {
chars: Vec<char>,
index: usize,
_source: &'a str,
}
impl<'a> Parser<'a> {
fn new(source: &'a str) -> Self {
Self {
chars: source.chars().collect(),
index: 0,
_source: source,
}
}
fn parse_value(&mut self) -> Result<JsonValue, JsonError> {
self.skip_whitespace();
match self.peek() {
Some('n') => self.parse_literal("null", JsonValue::Null),
Some('t') => self.parse_literal("true", JsonValue::Bool(true)),
Some('f') => self.parse_literal("false", JsonValue::Bool(false)),
Some('"') => self.parse_string().map(JsonValue::String),
Some('[') => self.parse_array(),
Some('{') => self.parse_object(),
Some('-' | '0'..='9') => self.parse_number().map(JsonValue::Number),
Some(other) => Err(JsonError::new(format!("unexpected character: {other}"))),
None => Err(JsonError::new("unexpected end of input")),
}
}
fn parse_literal(&mut self, expected: &str, value: JsonValue) -> Result<JsonValue, JsonError> {
for expected_char in expected.chars() {
if self.next() != Some(expected_char) {
return Err(JsonError::new(format!(
"invalid literal: expected {expected}"
)));
}
}
Ok(value)
}
fn parse_string(&mut self) -> Result<String, JsonError> {
self.expect('"')?;
let mut value = String::new();
while let Some(ch) = self.next() {
match ch {
'"' => return Ok(value),
'\\' => value.push(self.parse_escape()?),
plain => value.push(plain),
}
}
Err(JsonError::new("unterminated string"))
}
fn parse_escape(&mut self) -> Result<char, JsonError> {
match self.next() {
Some('"') => Ok('"'),
Some('\\') => Ok('\\'),
Some('/') => Ok('/'),
Some('b') => Ok('\u{08}'),
Some('f') => Ok('\u{0C}'),
Some('n') => Ok('\n'),
Some('r') => Ok('\r'),
Some('t') => Ok('\t'),
Some('u') => self.parse_unicode_escape(),
Some(other) => Err(JsonError::new(format!("invalid escape sequence: {other}"))),
None => Err(JsonError::new("unexpected end of input in escape sequence")),
}
}
fn parse_unicode_escape(&mut self) -> Result<char, JsonError> {
let mut value = 0_u32;
for _ in 0..4 {
let Some(ch) = self.next() else {
return Err(JsonError::new("unexpected end of input in unicode escape"));
};
value = (value << 4)
| ch.to_digit(16)
.ok_or_else(|| JsonError::new("invalid unicode escape"))?;
}
char::from_u32(value).ok_or_else(|| JsonError::new("invalid unicode scalar value"))
}
fn parse_array(&mut self) -> Result<JsonValue, JsonError> {
self.expect('[')?;
let mut values = Vec::new();
loop {
self.skip_whitespace();
if self.try_consume(']') {
break;
}
values.push(self.parse_value()?);
self.skip_whitespace();
if self.try_consume(']') {
break;
}
self.expect(',')?;
}
Ok(JsonValue::Array(values))
}
fn parse_object(&mut self) -> Result<JsonValue, JsonError> {
self.expect('{')?;
let mut entries = BTreeMap::new();
loop {
self.skip_whitespace();
if self.try_consume('}') {
break;
}
let key = self.parse_string()?;
self.skip_whitespace();
self.expect(':')?;
let value = self.parse_value()?;
entries.insert(key, value);
self.skip_whitespace();
if self.try_consume('}') {
break;
}
self.expect(',')?;
}
Ok(JsonValue::Object(entries))
}
fn parse_number(&mut self) -> Result<i64, JsonError> {
let mut value = String::new();
if self.try_consume('-') {
value.push('-');
}
while let Some(ch @ '0'..='9') = self.peek() {
value.push(ch);
self.index += 1;
}
if value.is_empty() || value == "-" {
return Err(JsonError::new("invalid number"));
}
value
.parse::<i64>()
.map_err(|_| JsonError::new("number out of range"))
}
fn expect(&mut self, expected: char) -> Result<(), JsonError> {
match self.next() {
Some(actual) if actual == expected => Ok(()),
Some(actual) => Err(JsonError::new(format!(
"expected '{expected}', found '{actual}'"
))),
None => Err(JsonError::new(format!(
"expected '{expected}', found end of input"
))),
}
}
fn try_consume(&mut self, expected: char) -> bool {
if self.peek() == Some(expected) {
self.index += 1;
true
} else {
false
}
}
fn skip_whitespace(&mut self) {
while matches!(self.peek(), Some(' ' | '\n' | '\r' | '\t')) {
self.index += 1;
}
}
fn peek(&self) -> Option<char> {
self.chars.get(self.index).copied()
}
fn next(&mut self) -> Option<char> {
let ch = self.peek()?;
self.index += 1;
Some(ch)
}
fn is_eof(&self) -> bool {
self.index >= self.chars.len()
}
}
#[cfg(test)]
mod tests {
use super::{render_string, JsonValue};
use std::collections::BTreeMap;
#[test]
fn renders_and_parses_json_values() {
let mut object = BTreeMap::new();
object.insert("flag".to_string(), JsonValue::Bool(true));
object.insert(
"items".to_string(),
JsonValue::Array(vec![
JsonValue::Number(4),
JsonValue::String("ok".to_string()),
]),
);
let rendered = JsonValue::Object(object).render();
let parsed = JsonValue::parse(&rendered).expect("json should parse");
assert_eq!(parsed.as_object().expect("object").len(), 2);
}
#[test]
fn escapes_control_characters() {
assert_eq!(render_string("a\n\t\"b"), "\"a\\n\\t\\\"b\"");
}
}
@@ -0,0 +1,383 @@
#![allow(clippy::similar_names)]
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum LaneEventName {
#[serde(rename = "lane.started")]
Started,
#[serde(rename = "lane.ready")]
Ready,
#[serde(rename = "lane.prompt_misdelivery")]
PromptMisdelivery,
#[serde(rename = "lane.blocked")]
Blocked,
#[serde(rename = "lane.red")]
Red,
#[serde(rename = "lane.green")]
Green,
#[serde(rename = "lane.commit.created")]
CommitCreated,
#[serde(rename = "lane.pr.opened")]
PrOpened,
#[serde(rename = "lane.merge.ready")]
MergeReady,
#[serde(rename = "lane.finished")]
Finished,
#[serde(rename = "lane.failed")]
Failed,
#[serde(rename = "lane.reconciled")]
Reconciled,
#[serde(rename = "lane.merged")]
Merged,
#[serde(rename = "lane.superseded")]
Superseded,
#[serde(rename = "lane.closed")]
Closed,
#[serde(rename = "branch.stale_against_main")]
BranchStaleAgainstMain,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum LaneEventStatus {
Running,
Ready,
Blocked,
Red,
Green,
Completed,
Failed,
Reconciled,
Merged,
Superseded,
Closed,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum LaneFailureClass {
PromptDelivery,
TrustGate,
BranchDivergence,
Compile,
Test,
PluginStartup,
McpStartup,
McpHandshake,
GatewayRouting,
ToolRuntime,
Infra,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LaneEventBlocker {
#[serde(rename = "failureClass")]
pub failure_class: LaneFailureClass,
pub detail: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LaneCommitProvenance {
pub commit: String,
pub branch: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub worktree: Option<String>,
#[serde(rename = "canonicalCommit", skip_serializing_if = "Option::is_none")]
pub canonical_commit: Option<String>,
#[serde(rename = "supersededBy", skip_serializing_if = "Option::is_none")]
pub superseded_by: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub lineage: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LaneEvent {
pub event: LaneEventName,
pub status: LaneEventStatus,
#[serde(rename = "emittedAt")]
pub emitted_at: String,
#[serde(rename = "failureClass", skip_serializing_if = "Option::is_none")]
pub failure_class: Option<LaneFailureClass>,
#[serde(skip_serializing_if = "Option::is_none")]
pub detail: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub data: Option<Value>,
}
impl LaneEvent {
#[must_use]
pub fn new(
event: LaneEventName,
status: LaneEventStatus,
emitted_at: impl Into<String>,
) -> Self {
Self {
event,
status,
emitted_at: emitted_at.into(),
failure_class: None,
detail: None,
data: None,
}
}
#[must_use]
pub fn started(emitted_at: impl Into<String>) -> Self {
Self::new(LaneEventName::Started, LaneEventStatus::Running, emitted_at)
}
#[must_use]
pub fn finished(emitted_at: impl Into<String>, detail: Option<String>) -> Self {
Self::new(
LaneEventName::Finished,
LaneEventStatus::Completed,
emitted_at,
)
.with_optional_detail(detail)
}
#[must_use]
pub fn commit_created(
emitted_at: impl Into<String>,
detail: Option<String>,
provenance: LaneCommitProvenance,
) -> Self {
Self::new(
LaneEventName::CommitCreated,
LaneEventStatus::Completed,
emitted_at,
)
.with_optional_detail(detail)
.with_data(serde_json::to_value(provenance).expect("commit provenance should serialize"))
}
#[must_use]
pub fn superseded(
emitted_at: impl Into<String>,
detail: Option<String>,
provenance: LaneCommitProvenance,
) -> Self {
Self::new(
LaneEventName::Superseded,
LaneEventStatus::Superseded,
emitted_at,
)
.with_optional_detail(detail)
.with_data(serde_json::to_value(provenance).expect("commit provenance should serialize"))
}
#[must_use]
pub fn blocked(emitted_at: impl Into<String>, blocker: &LaneEventBlocker) -> Self {
Self::new(LaneEventName::Blocked, LaneEventStatus::Blocked, emitted_at)
.with_failure_class(blocker.failure_class)
.with_detail(blocker.detail.clone())
}
#[must_use]
pub fn failed(emitted_at: impl Into<String>, blocker: &LaneEventBlocker) -> Self {
Self::new(LaneEventName::Failed, LaneEventStatus::Failed, emitted_at)
.with_failure_class(blocker.failure_class)
.with_detail(blocker.detail.clone())
}
#[must_use]
pub fn with_failure_class(mut self, failure_class: LaneFailureClass) -> Self {
self.failure_class = Some(failure_class);
self
}
#[must_use]
pub fn with_detail(mut self, detail: impl Into<String>) -> Self {
self.detail = Some(detail.into());
self
}
#[must_use]
pub fn with_optional_detail(mut self, detail: Option<String>) -> Self {
self.detail = detail;
self
}
#[must_use]
pub fn with_data(mut self, data: Value) -> Self {
self.data = Some(data);
self
}
}
#[must_use]
pub fn dedupe_superseded_commit_events(events: &[LaneEvent]) -> Vec<LaneEvent> {
let mut keep = vec![true; events.len()];
let mut latest_by_key = std::collections::BTreeMap::<String, usize>::new();
for (index, event) in events.iter().enumerate() {
if event.event != LaneEventName::CommitCreated {
continue;
}
let Some(data) = event.data.as_ref() else {
continue;
};
let key = data
.get("canonicalCommit")
.or_else(|| data.get("commit"))
.and_then(serde_json::Value::as_str)
.map(str::to_string);
let superseded = data
.get("supersededBy")
.and_then(serde_json::Value::as_str)
.is_some();
if superseded {
keep[index] = false;
continue;
}
if let Some(key) = key {
if let Some(previous) = latest_by_key.insert(key, index) {
keep[previous] = false;
}
}
}
events
.iter()
.cloned()
.zip(keep)
.filter_map(|(event, retain)| retain.then_some(event))
.collect()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
dedupe_superseded_commit_events, LaneCommitProvenance, LaneEvent, LaneEventBlocker,
LaneEventName, LaneEventStatus, LaneFailureClass,
};
#[test]
fn canonical_lane_event_names_serialize_to_expected_wire_values() {
let cases = [
(LaneEventName::Started, "lane.started"),
(LaneEventName::Ready, "lane.ready"),
(LaneEventName::PromptMisdelivery, "lane.prompt_misdelivery"),
(LaneEventName::Blocked, "lane.blocked"),
(LaneEventName::Red, "lane.red"),
(LaneEventName::Green, "lane.green"),
(LaneEventName::CommitCreated, "lane.commit.created"),
(LaneEventName::PrOpened, "lane.pr.opened"),
(LaneEventName::MergeReady, "lane.merge.ready"),
(LaneEventName::Finished, "lane.finished"),
(LaneEventName::Failed, "lane.failed"),
(LaneEventName::Reconciled, "lane.reconciled"),
(LaneEventName::Merged, "lane.merged"),
(LaneEventName::Superseded, "lane.superseded"),
(LaneEventName::Closed, "lane.closed"),
(
LaneEventName::BranchStaleAgainstMain,
"branch.stale_against_main",
),
];
for (event, expected) in cases {
assert_eq!(
serde_json::to_value(event).expect("serialize event"),
json!(expected)
);
}
}
#[test]
fn failure_classes_cover_canonical_taxonomy_wire_values() {
let cases = [
(LaneFailureClass::PromptDelivery, "prompt_delivery"),
(LaneFailureClass::TrustGate, "trust_gate"),
(LaneFailureClass::BranchDivergence, "branch_divergence"),
(LaneFailureClass::Compile, "compile"),
(LaneFailureClass::Test, "test"),
(LaneFailureClass::PluginStartup, "plugin_startup"),
(LaneFailureClass::McpStartup, "mcp_startup"),
(LaneFailureClass::McpHandshake, "mcp_handshake"),
(LaneFailureClass::GatewayRouting, "gateway_routing"),
(LaneFailureClass::ToolRuntime, "tool_runtime"),
(LaneFailureClass::Infra, "infra"),
];
for (failure_class, expected) in cases {
assert_eq!(
serde_json::to_value(failure_class).expect("serialize failure class"),
json!(expected)
);
}
}
#[test]
fn blocked_and_failed_events_reuse_blocker_details() {
let blocker = LaneEventBlocker {
failure_class: LaneFailureClass::McpStartup,
detail: "broken server".to_string(),
};
let blocked = LaneEvent::blocked("2026-04-04T00:00:00Z", &blocker);
let failed = LaneEvent::failed("2026-04-04T00:00:01Z", &blocker);
assert_eq!(blocked.event, LaneEventName::Blocked);
assert_eq!(blocked.status, LaneEventStatus::Blocked);
assert_eq!(blocked.failure_class, Some(LaneFailureClass::McpStartup));
assert_eq!(failed.event, LaneEventName::Failed);
assert_eq!(failed.status, LaneEventStatus::Failed);
assert_eq!(failed.detail.as_deref(), Some("broken server"));
}
#[test]
fn commit_events_can_carry_worktree_and_supersession_metadata() {
let event = LaneEvent::commit_created(
"2026-04-04T00:00:00Z",
Some("commit created".to_string()),
LaneCommitProvenance {
commit: "abc123".to_string(),
branch: "feature/provenance".to_string(),
worktree: Some("wt-a".to_string()),
canonical_commit: Some("abc123".to_string()),
superseded_by: None,
lineage: vec!["abc123".to_string()],
},
);
let event_json = serde_json::to_value(&event).expect("lane event should serialize");
assert_eq!(event_json["event"], "lane.commit.created");
assert_eq!(event_json["data"]["branch"], "feature/provenance");
assert_eq!(event_json["data"]["worktree"], "wt-a");
}
#[test]
fn dedupes_superseded_commit_events_by_canonical_commit() {
let retained = dedupe_superseded_commit_events(&[
LaneEvent::commit_created(
"2026-04-04T00:00:00Z",
Some("old".to_string()),
LaneCommitProvenance {
commit: "old123".to_string(),
branch: "feature/provenance".to_string(),
worktree: Some("wt-a".to_string()),
canonical_commit: Some("canon123".to_string()),
superseded_by: Some("new123".to_string()),
lineage: vec!["old123".to_string(), "new123".to_string()],
},
),
LaneEvent::commit_created(
"2026-04-04T00:00:01Z",
Some("new".to_string()),
LaneCommitProvenance {
commit: "new123".to_string(),
branch: "feature/provenance".to_string(),
worktree: Some("wt-b".to_string()),
canonical_commit: Some("canon123".to_string()),
superseded_by: None,
lineage: vec!["old123".to_string(), "new123".to_string()],
},
),
]);
assert_eq!(retained.len(), 1);
assert_eq!(retained[0].detail.as_deref(), Some("new"));
}
}
+179
View File
@@ -0,0 +1,179 @@
//! Core runtime primitives for the `claw` CLI and supporting crates.
//!
//! This crate owns session persistence, permission evaluation, prompt assembly,
//! MCP plumbing, tool-facing file operations, and the core conversation loop
//! that drives interactive and one-shot turns.
mod bash;
pub mod bash_validation;
mod bootstrap;
pub mod branch_lock;
mod compact;
mod config;
pub mod config_validate;
mod conversation;
mod file_ops;
mod git_context;
pub mod green_contract;
mod hooks;
mod json;
mod lane_events;
pub mod lsp_client;
mod mcp;
mod mcp_client;
pub mod mcp_lifecycle_hardened;
pub mod mcp_server;
mod mcp_stdio;
pub mod mcp_tool_bridge;
mod oauth;
pub mod permission_enforcer;
mod permissions;
pub mod plugin_lifecycle;
mod policy_engine;
mod prompt;
pub mod recovery_recipes;
mod remote;
pub mod sandbox;
mod session;
pub mod session_control;
pub use session_control::SessionStore;
mod sse;
pub mod stale_base;
pub mod stale_branch;
pub mod summary_compression;
pub mod task_packet;
pub mod task_registry;
pub mod team_cron_registry;
#[cfg(test)]
mod trust_resolver;
mod usage;
pub mod worker_boot;
pub use bash::{execute_bash, BashCommandInput, BashCommandOutput};
pub use bootstrap::{BootstrapPhase, BootstrapPlan};
pub use branch_lock::{detect_branch_lock_collisions, BranchLockCollision, BranchLockIntent};
pub use compact::{
compact_session, estimate_session_tokens, format_compact_summary,
get_compact_continuation_message, should_compact, CompactionConfig, CompactionResult,
};
pub use config::{
ConfigEntry, ConfigError, ConfigLoader, ConfigSource, McpConfigCollection,
McpManagedProxyServerConfig, McpOAuthConfig, McpRemoteServerConfig, McpSdkServerConfig,
McpServerConfig, McpStdioServerConfig, McpTransport, McpWebSocketServerConfig, OAuthConfig,
ProviderFallbackConfig, ResolvedPermissionMode, RuntimeConfig, RuntimeFeatureConfig,
RuntimeHookConfig, RuntimePermissionRuleConfig, RuntimePluginConfig, ScopedMcpServerConfig,
CLAW_SETTINGS_SCHEMA_NAME,
};
pub use config_validate::{
check_unsupported_format, format_diagnostics, validate_config_file, ConfigDiagnostic,
DiagnosticKind, ValidationResult,
};
pub use conversation::{
auto_compaction_threshold_from_env, ApiClient, ApiRequest, AssistantEvent, AutoCompactionEvent,
ConversationRuntime, PromptCacheEvent, RuntimeError, StaticToolExecutor, ToolError,
ToolExecutor, TurnSummary,
};
pub use file_ops::{
edit_file, glob_search, grep_search, read_file, write_file, EditFileOutput, GlobSearchOutput,
GrepSearchInput, GrepSearchOutput, ReadFileOutput, StructuredPatchHunk, TextFilePayload,
WriteFileOutput,
};
pub use git_context::{GitCommitEntry, GitContext};
pub use hooks::{
HookAbortSignal, HookEvent, HookProgressEvent, HookProgressReporter, HookRunResult, HookRunner,
};
pub use lane_events::{
dedupe_superseded_commit_events, LaneCommitProvenance, LaneEvent, LaneEventBlocker,
LaneEventName, LaneEventStatus, LaneFailureClass,
};
pub use mcp::{
mcp_server_signature, mcp_tool_name, mcp_tool_prefix, normalize_name_for_mcp,
scoped_mcp_config_hash, unwrap_ccr_proxy_url,
};
pub use mcp_client::{
McpClientAuth, McpClientBootstrap, McpClientTransport, McpManagedProxyTransport,
McpRemoteTransport, McpSdkTransport, McpStdioTransport,
};
pub use mcp_lifecycle_hardened::{
McpDegradedReport, McpErrorSurface, McpFailedServer, McpLifecyclePhase, McpLifecycleState,
McpLifecycleValidator, McpPhaseResult,
};
pub use mcp_server::{McpServer, McpServerSpec, ToolCallHandler, MCP_SERVER_PROTOCOL_VERSION};
pub use mcp_stdio::{
spawn_mcp_stdio_process, JsonRpcError, JsonRpcId, JsonRpcRequest, JsonRpcResponse,
ManagedMcpTool, McpDiscoveryFailure, McpInitializeClientInfo, McpInitializeParams,
McpInitializeResult, McpInitializeServerInfo, McpListResourcesParams, McpListResourcesResult,
McpListToolsParams, McpListToolsResult, McpReadResourceParams, McpReadResourceResult,
McpResource, McpResourceContents, McpServerManager, McpServerManagerError, McpStdioProcess,
McpTool, McpToolCallContent, McpToolCallParams, McpToolCallResult, McpToolDiscoveryReport,
UnsupportedMcpServer,
};
pub use oauth::{
clear_oauth_credentials, code_challenge_s256, credentials_path, generate_pkce_pair,
generate_state, load_oauth_credentials, loopback_redirect_uri, parse_oauth_callback_query,
parse_oauth_callback_request_target, save_oauth_credentials, OAuthAuthorizationRequest,
OAuthCallbackParams, OAuthRefreshRequest, OAuthTokenExchangeRequest, OAuthTokenSet,
PkceChallengeMethod, PkceCodePair,
};
pub use permissions::{
PermissionContext, PermissionMode, PermissionOutcome, PermissionOverride, PermissionPolicy,
PermissionPromptDecision, PermissionPrompter, PermissionRequest,
};
pub use plugin_lifecycle::{
DegradedMode, DiscoveryResult, PluginHealthcheck, PluginLifecycle, PluginLifecycleEvent,
PluginState, ResourceInfo, ServerHealth, ServerStatus, ToolInfo,
};
pub use policy_engine::{
evaluate, DiffScope, GreenLevel, LaneBlocker, LaneContext, PolicyAction, PolicyCondition,
PolicyEngine, PolicyRule, ReconcileReason, ReviewStatus,
};
pub use prompt::{
load_system_prompt, prepend_bullets, ContextFile, ProjectContext, PromptBuildError,
SystemPromptBuilder, FRONTIER_MODEL_NAME, SYSTEM_PROMPT_DYNAMIC_BOUNDARY,
};
pub use recovery_recipes::{
attempt_recovery, recipe_for, EscalationPolicy, FailureScenario, RecoveryContext,
RecoveryEvent, RecoveryRecipe, RecoveryResult, RecoveryStep,
};
pub use remote::{
inherited_upstream_proxy_env, no_proxy_list, read_token, upstream_proxy_ws_url,
RemoteSessionContext, UpstreamProxyBootstrap, UpstreamProxyState, DEFAULT_REMOTE_BASE_URL,
DEFAULT_SESSION_TOKEN_PATH, DEFAULT_SYSTEM_CA_BUNDLE, NO_PROXY_HOSTS, UPSTREAM_PROXY_ENV_KEYS,
};
pub use sandbox::{
build_linux_sandbox_command, detect_container_environment, detect_container_environment_from,
resolve_sandbox_status, resolve_sandbox_status_for_request, ContainerEnvironment,
FilesystemIsolationMode, LinuxSandboxCommand, SandboxConfig, SandboxDetectionInputs,
SandboxRequest, SandboxStatus,
};
pub use session::{
ContentBlock, ConversationMessage, MessageRole, Session, SessionCompaction, SessionError,
SessionFork, SessionPromptEntry,
};
pub use sse::{IncrementalSseParser, SseEvent};
pub use stale_base::{
check_base_commit, format_stale_base_warning, read_claw_base_file, resolve_expected_base,
BaseCommitSource, BaseCommitState,
};
pub use stale_branch::{
apply_policy, check_freshness, BranchFreshness, StaleBranchAction, StaleBranchEvent,
StaleBranchPolicy,
};
pub use task_packet::{validate_packet, TaskPacket, TaskPacketValidationError, ValidatedPacket};
#[cfg(test)]
pub use trust_resolver::{TrustConfig, TrustDecision, TrustEvent, TrustPolicy, TrustResolver};
pub use usage::{
format_usd, pricing_for_model, ModelPricing, TokenUsage, UsageCostEstimate, UsageTracker,
};
pub use worker_boot::{
Worker, WorkerEvent, WorkerEventKind, WorkerEventPayload, WorkerFailure, WorkerFailureKind,
WorkerPromptTarget, WorkerReadySnapshot, WorkerRegistry, WorkerStatus, WorkerTrustResolution,
};
#[cfg(test)]
pub(crate) fn test_env_lock() -> std::sync::MutexGuard<'static, ()> {
static LOCK: std::sync::OnceLock<std::sync::Mutex<()>> = std::sync::OnceLock::new();
LOCK.get_or_init(|| std::sync::Mutex::new(()))
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
@@ -0,0 +1,747 @@
#![allow(clippy::should_implement_trait, clippy::must_use_candidate)]
//! LSP (Language Server Protocol) client registry for tool dispatch.
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use serde::{Deserialize, Serialize};
/// Supported LSP actions.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum LspAction {
Diagnostics,
Hover,
Definition,
References,
Completion,
Symbols,
Format,
}
impl LspAction {
pub fn from_str(s: &str) -> Option<Self> {
match s {
"diagnostics" => Some(Self::Diagnostics),
"hover" => Some(Self::Hover),
"definition" | "goto_definition" => Some(Self::Definition),
"references" | "find_references" => Some(Self::References),
"completion" | "completions" => Some(Self::Completion),
"symbols" | "document_symbols" => Some(Self::Symbols),
"format" | "formatting" => Some(Self::Format),
_ => None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LspDiagnostic {
pub path: String,
pub line: u32,
pub character: u32,
pub severity: String,
pub message: String,
pub source: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LspLocation {
pub path: String,
pub line: u32,
pub character: u32,
pub end_line: Option<u32>,
pub end_character: Option<u32>,
pub preview: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LspHoverResult {
pub content: String,
pub language: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LspCompletionItem {
pub label: String,
pub kind: Option<String>,
pub detail: Option<String>,
pub insert_text: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LspSymbol {
pub name: String,
pub kind: String,
pub path: String,
pub line: u32,
pub character: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum LspServerStatus {
Connected,
Disconnected,
Starting,
Error,
}
impl std::fmt::Display for LspServerStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Connected => write!(f, "connected"),
Self::Disconnected => write!(f, "disconnected"),
Self::Starting => write!(f, "starting"),
Self::Error => write!(f, "error"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LspServerState {
pub language: String,
pub status: LspServerStatus,
pub root_path: Option<String>,
pub capabilities: Vec<String>,
pub diagnostics: Vec<LspDiagnostic>,
}
#[derive(Debug, Clone, Default)]
pub struct LspRegistry {
inner: Arc<Mutex<RegistryInner>>,
}
#[derive(Debug, Default)]
struct RegistryInner {
servers: HashMap<String, LspServerState>,
}
impl LspRegistry {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn register(
&self,
language: &str,
status: LspServerStatus,
root_path: Option<&str>,
capabilities: Vec<String>,
) {
let mut inner = self.inner.lock().expect("lsp registry lock poisoned");
inner.servers.insert(
language.to_owned(),
LspServerState {
language: language.to_owned(),
status,
root_path: root_path.map(str::to_owned),
capabilities,
diagnostics: Vec::new(),
},
);
}
pub fn get(&self, language: &str) -> Option<LspServerState> {
let inner = self.inner.lock().expect("lsp registry lock poisoned");
inner.servers.get(language).cloned()
}
/// Find the appropriate server for a file path based on extension.
pub fn find_server_for_path(&self, path: &str) -> Option<LspServerState> {
let ext = std::path::Path::new(path)
.extension()
.and_then(|e| e.to_str())
.unwrap_or("");
let language = match ext {
"rs" => "rust",
"ts" | "tsx" => "typescript",
"js" | "jsx" => "javascript",
"py" => "python",
"go" => "go",
"java" => "java",
"c" | "h" => "c",
"cpp" | "hpp" | "cc" => "cpp",
"rb" => "ruby",
"lua" => "lua",
_ => return None,
};
self.get(language)
}
/// List all registered servers.
pub fn list_servers(&self) -> Vec<LspServerState> {
let inner = self.inner.lock().expect("lsp registry lock poisoned");
inner.servers.values().cloned().collect()
}
/// Add diagnostics to a server.
pub fn add_diagnostics(
&self,
language: &str,
diagnostics: Vec<LspDiagnostic>,
) -> Result<(), String> {
let mut inner = self.inner.lock().expect("lsp registry lock poisoned");
let server = inner
.servers
.get_mut(language)
.ok_or_else(|| format!("LSP server not found for language: {language}"))?;
server.diagnostics.extend(diagnostics);
Ok(())
}
/// Get diagnostics for a specific file path.
pub fn get_diagnostics(&self, path: &str) -> Vec<LspDiagnostic> {
let inner = self.inner.lock().expect("lsp registry lock poisoned");
inner
.servers
.values()
.flat_map(|s| &s.diagnostics)
.filter(|d| d.path == path)
.cloned()
.collect()
}
/// Clear diagnostics for a language server.
pub fn clear_diagnostics(&self, language: &str) -> Result<(), String> {
let mut inner = self.inner.lock().expect("lsp registry lock poisoned");
let server = inner
.servers
.get_mut(language)
.ok_or_else(|| format!("LSP server not found for language: {language}"))?;
server.diagnostics.clear();
Ok(())
}
/// Disconnect a server.
pub fn disconnect(&self, language: &str) -> Option<LspServerState> {
let mut inner = self.inner.lock().expect("lsp registry lock poisoned");
inner.servers.remove(language)
}
#[must_use]
pub fn len(&self) -> usize {
let inner = self.inner.lock().expect("lsp registry lock poisoned");
inner.servers.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
/// Dispatch an LSP action and return a structured result.
pub fn dispatch(
&self,
action: &str,
path: Option<&str>,
line: Option<u32>,
character: Option<u32>,
_query: Option<&str>,
) -> Result<serde_json::Value, String> {
let lsp_action =
LspAction::from_str(action).ok_or_else(|| format!("unknown LSP action: {action}"))?;
// For diagnostics, we can check existing cached diagnostics
if lsp_action == LspAction::Diagnostics {
if let Some(path) = path {
let diags = self.get_diagnostics(path);
return Ok(serde_json::json!({
"action": "diagnostics",
"path": path,
"diagnostics": diags,
"count": diags.len()
}));
}
// All diagnostics across all servers
let inner = self.inner.lock().expect("lsp registry lock poisoned");
let all_diags: Vec<_> = inner
.servers
.values()
.flat_map(|s| &s.diagnostics)
.collect();
return Ok(serde_json::json!({
"action": "diagnostics",
"diagnostics": all_diags,
"count": all_diags.len()
}));
}
// For other actions, we need a connected server for the given file
let path = path.ok_or("path is required for this LSP action")?;
let server = self
.find_server_for_path(path)
.ok_or_else(|| format!("no LSP server available for path: {path}"))?;
if server.status != LspServerStatus::Connected {
return Err(format!(
"LSP server for '{}' is not connected (status: {})",
server.language, server.status
));
}
// Return structured placeholder — actual LSP JSON-RPC calls would
// go through the real LSP process here.
Ok(serde_json::json!({
"action": action,
"path": path,
"line": line,
"character": character,
"language": server.language,
"status": "dispatched",
"message": format!("LSP {} dispatched to {} server", action, server.language)
}))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn registers_and_retrieves_server() {
let registry = LspRegistry::new();
registry.register(
"rust",
LspServerStatus::Connected,
Some("/workspace"),
vec!["hover".into(), "completion".into()],
);
let server = registry.get("rust").expect("should exist");
assert_eq!(server.language, "rust");
assert_eq!(server.status, LspServerStatus::Connected);
assert_eq!(server.capabilities.len(), 2);
}
#[test]
fn finds_server_by_file_extension() {
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Connected, None, vec![]);
registry.register("typescript", LspServerStatus::Connected, None, vec![]);
let rs_server = registry.find_server_for_path("src/main.rs").unwrap();
assert_eq!(rs_server.language, "rust");
let ts_server = registry.find_server_for_path("src/index.ts").unwrap();
assert_eq!(ts_server.language, "typescript");
assert!(registry.find_server_for_path("data.csv").is_none());
}
#[test]
fn manages_diagnostics() {
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Connected, None, vec![]);
registry
.add_diagnostics(
"rust",
vec![LspDiagnostic {
path: "src/main.rs".into(),
line: 10,
character: 5,
severity: "error".into(),
message: "mismatched types".into(),
source: Some("rust-analyzer".into()),
}],
)
.unwrap();
let diags = registry.get_diagnostics("src/main.rs");
assert_eq!(diags.len(), 1);
assert_eq!(diags[0].message, "mismatched types");
registry.clear_diagnostics("rust").unwrap();
assert!(registry.get_diagnostics("src/main.rs").is_empty());
}
#[test]
fn dispatches_diagnostics_action() {
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Connected, None, vec![]);
registry
.add_diagnostics(
"rust",
vec![LspDiagnostic {
path: "src/lib.rs".into(),
line: 1,
character: 0,
severity: "warning".into(),
message: "unused import".into(),
source: None,
}],
)
.unwrap();
let result = registry
.dispatch("diagnostics", Some("src/lib.rs"), None, None, None)
.unwrap();
assert_eq!(result["count"], 1);
}
#[test]
fn dispatches_hover_action() {
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Connected, None, vec![]);
let result = registry
.dispatch("hover", Some("src/main.rs"), Some(10), Some(5), None)
.unwrap();
assert_eq!(result["action"], "hover");
assert_eq!(result["language"], "rust");
}
#[test]
fn rejects_action_on_disconnected_server() {
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Disconnected, None, vec![]);
assert!(registry
.dispatch("hover", Some("src/main.rs"), Some(1), Some(0), None)
.is_err());
}
#[test]
fn rejects_unknown_action() {
let registry = LspRegistry::new();
assert!(registry
.dispatch("unknown_action", Some("file.rs"), None, None, None)
.is_err());
}
#[test]
fn disconnects_server() {
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Connected, None, vec![]);
assert_eq!(registry.len(), 1);
let removed = registry.disconnect("rust");
assert!(removed.is_some());
assert!(registry.is_empty());
}
#[test]
fn lsp_action_from_str_all_aliases() {
// given
let cases = [
("diagnostics", Some(LspAction::Diagnostics)),
("hover", Some(LspAction::Hover)),
("definition", Some(LspAction::Definition)),
("goto_definition", Some(LspAction::Definition)),
("references", Some(LspAction::References)),
("find_references", Some(LspAction::References)),
("completion", Some(LspAction::Completion)),
("completions", Some(LspAction::Completion)),
("symbols", Some(LspAction::Symbols)),
("document_symbols", Some(LspAction::Symbols)),
("format", Some(LspAction::Format)),
("formatting", Some(LspAction::Format)),
("unknown", None),
];
// when
let resolved: Vec<_> = cases
.into_iter()
.map(|(input, expected)| (input, LspAction::from_str(input), expected))
.collect();
// then
for (input, actual, expected) in resolved {
assert_eq!(actual, expected, "unexpected action resolution for {input}");
}
}
#[test]
fn lsp_server_status_display_all_variants() {
// given
let cases = [
(LspServerStatus::Connected, "connected"),
(LspServerStatus::Disconnected, "disconnected"),
(LspServerStatus::Starting, "starting"),
(LspServerStatus::Error, "error"),
];
// when
let rendered: Vec<_> = cases
.into_iter()
.map(|(status, expected)| (status.to_string(), expected))
.collect();
// then
assert_eq!(
rendered,
vec![
("connected".to_string(), "connected"),
("disconnected".to_string(), "disconnected"),
("starting".to_string(), "starting"),
("error".to_string(), "error"),
]
);
}
#[test]
fn dispatch_diagnostics_without_path_aggregates() {
// given
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Connected, None, vec![]);
registry.register("python", LspServerStatus::Connected, None, vec![]);
registry
.add_diagnostics(
"rust",
vec![LspDiagnostic {
path: "src/lib.rs".into(),
line: 1,
character: 0,
severity: "warning".into(),
message: "unused import".into(),
source: Some("rust-analyzer".into()),
}],
)
.expect("rust diagnostics should add");
registry
.add_diagnostics(
"python",
vec![LspDiagnostic {
path: "script.py".into(),
line: 2,
character: 4,
severity: "error".into(),
message: "undefined name".into(),
source: Some("pyright".into()),
}],
)
.expect("python diagnostics should add");
// when
let result = registry
.dispatch("diagnostics", None, None, None, None)
.expect("aggregate diagnostics should work");
// then
assert_eq!(result["action"], "diagnostics");
assert_eq!(result["count"], 2);
assert_eq!(result["diagnostics"].as_array().map(Vec::len), Some(2));
}
#[test]
fn dispatch_non_diagnostics_requires_path() {
// given
let registry = LspRegistry::new();
// when
let result = registry.dispatch("hover", None, Some(1), Some(0), None);
// then
assert_eq!(
result.expect_err("path should be required"),
"path is required for this LSP action"
);
}
#[test]
fn dispatch_no_server_for_path_errors() {
// given
let registry = LspRegistry::new();
// when
let result = registry.dispatch("hover", Some("notes.md"), Some(1), Some(0), None);
// then
let error = result.expect_err("missing server should fail");
assert!(error.contains("no LSP server available for path: notes.md"));
}
#[test]
fn dispatch_disconnected_server_error_payload() {
// given
let registry = LspRegistry::new();
registry.register("typescript", LspServerStatus::Disconnected, None, vec![]);
// when
let result = registry.dispatch("hover", Some("src/index.ts"), Some(3), Some(2), None);
// then
let error = result.expect_err("disconnected server should fail");
assert!(error.contains("typescript"));
assert!(error.contains("disconnected"));
}
#[test]
fn find_server_for_all_extensions() {
// given
let registry = LspRegistry::new();
for language in [
"rust",
"typescript",
"javascript",
"python",
"go",
"java",
"c",
"cpp",
"ruby",
"lua",
] {
registry.register(language, LspServerStatus::Connected, None, vec![]);
}
let cases = [
("src/main.rs", "rust"),
("src/index.ts", "typescript"),
("src/view.tsx", "typescript"),
("src/app.js", "javascript"),
("src/app.jsx", "javascript"),
("script.py", "python"),
("main.go", "go"),
("Main.java", "java"),
("native.c", "c"),
("native.h", "c"),
("native.cpp", "cpp"),
("native.hpp", "cpp"),
("native.cc", "cpp"),
("script.rb", "ruby"),
("script.lua", "lua"),
];
// when
let resolved: Vec<_> = cases
.into_iter()
.map(|(path, expected)| {
(
path,
registry
.find_server_for_path(path)
.map(|server| server.language),
expected,
)
})
.collect();
// then
for (path, actual, expected) in resolved {
assert_eq!(
actual.as_deref(),
Some(expected),
"unexpected mapping for {path}"
);
}
}
#[test]
fn find_server_for_path_no_extension() {
// given
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Connected, None, vec![]);
// when
let result = registry.find_server_for_path("Makefile");
// then
assert!(result.is_none());
}
#[test]
fn list_servers_with_multiple() {
// given
let registry = LspRegistry::new();
registry.register("rust", LspServerStatus::Connected, None, vec![]);
registry.register("typescript", LspServerStatus::Starting, None, vec![]);
registry.register("python", LspServerStatus::Error, None, vec![]);
// when
let servers = registry.list_servers();
// then
assert_eq!(servers.len(), 3);
assert!(servers.iter().any(|server| server.language == "rust"));
assert!(servers.iter().any(|server| server.language == "typescript"));
assert!(servers.iter().any(|server| server.language == "python"));
}
#[test]
fn get_missing_server_returns_none() {
// given
let registry = LspRegistry::new();
// when
let server = registry.get("missing");
// then
assert!(server.is_none());
}
#[test]
fn add_diagnostics_missing_language_errors() {
// given
let registry = LspRegistry::new();
// when
let result = registry.add_diagnostics("missing", vec![]);
// then
let error = result.expect_err("missing language should fail");
assert!(error.contains("LSP server not found for language: missing"));
}
#[test]
fn get_diagnostics_across_servers() {
// given
let registry = LspRegistry::new();
let shared_path = "shared/file.txt";
registry.register("rust", LspServerStatus::Connected, None, vec![]);
registry.register("python", LspServerStatus::Connected, None, vec![]);
registry
.add_diagnostics(
"rust",
vec![LspDiagnostic {
path: shared_path.into(),
line: 4,
character: 1,
severity: "warning".into(),
message: "warn".into(),
source: None,
}],
)
.expect("rust diagnostics should add");
registry
.add_diagnostics(
"python",
vec![LspDiagnostic {
path: shared_path.into(),
line: 8,
character: 3,
severity: "error".into(),
message: "err".into(),
source: None,
}],
)
.expect("python diagnostics should add");
// when
let diagnostics = registry.get_diagnostics(shared_path);
// then
assert_eq!(diagnostics.len(), 2);
assert!(diagnostics
.iter()
.any(|diagnostic| diagnostic.message == "warn"));
assert!(diagnostics
.iter()
.any(|diagnostic| diagnostic.message == "err"));
}
#[test]
fn clear_diagnostics_missing_language_errors() {
// given
let registry = LspRegistry::new();
// when
let result = registry.clear_diagnostics("missing");
// then
let error = result.expect_err("missing language should fail");
assert!(error.contains("LSP server not found for language: missing"));
}
}
+304
View File
@@ -0,0 +1,304 @@
use crate::config::{McpServerConfig, ScopedMcpServerConfig};
const CLAUDEAI_SERVER_PREFIX: &str = "claude.ai ";
const CCR_PROXY_PATH_MARKERS: [&str; 2] = ["/v2/session_ingress/shttp/mcp/", "/v2/ccr-sessions/"];
#[must_use]
pub fn normalize_name_for_mcp(name: &str) -> String {
let mut normalized = name
.chars()
.map(|ch| match ch {
'a'..='z' | 'A'..='Z' | '0'..='9' | '_' | '-' => ch,
_ => '_',
})
.collect::<String>();
if name.starts_with(CLAUDEAI_SERVER_PREFIX) {
normalized = collapse_underscores(&normalized)
.trim_matches('_')
.to_string();
}
normalized
}
#[must_use]
pub fn mcp_tool_prefix(server_name: &str) -> String {
format!("mcp__{}__", normalize_name_for_mcp(server_name))
}
#[must_use]
pub fn mcp_tool_name(server_name: &str, tool_name: &str) -> String {
format!(
"{}{}",
mcp_tool_prefix(server_name),
normalize_name_for_mcp(tool_name)
)
}
#[must_use]
pub fn unwrap_ccr_proxy_url(url: &str) -> String {
if !CCR_PROXY_PATH_MARKERS
.iter()
.any(|marker| url.contains(marker))
{
return url.to_string();
}
let Some(query_start) = url.find('?') else {
return url.to_string();
};
let query = &url[query_start + 1..];
for pair in query.split('&') {
let mut parts = pair.splitn(2, '=');
if matches!(parts.next(), Some("mcp_url")) {
if let Some(value) = parts.next() {
return percent_decode(value);
}
}
}
url.to_string()
}
#[must_use]
pub fn mcp_server_signature(config: &McpServerConfig) -> Option<String> {
match config {
McpServerConfig::Stdio(config) => {
let mut command = vec![config.command.clone()];
command.extend(config.args.clone());
Some(format!("stdio:{}", render_command_signature(&command)))
}
McpServerConfig::Sse(config) | McpServerConfig::Http(config) => {
Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url)))
}
McpServerConfig::Ws(config) => Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url))),
McpServerConfig::ManagedProxy(config) => {
Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url)))
}
McpServerConfig::Sdk(_) => None,
}
}
#[must_use]
pub fn scoped_mcp_config_hash(config: &ScopedMcpServerConfig) -> String {
let rendered = match &config.config {
McpServerConfig::Stdio(stdio) => format!(
"stdio|{}|{}|{}|{}",
stdio.command,
render_command_signature(&stdio.args),
render_env_signature(&stdio.env),
stdio
.tool_call_timeout_ms
.map_or_else(String::new, |timeout_ms| timeout_ms.to_string())
),
McpServerConfig::Sse(remote) => format!(
"sse|{}|{}|{}|{}",
remote.url,
render_env_signature(&remote.headers),
remote.headers_helper.as_deref().unwrap_or(""),
render_oauth_signature(remote.oauth.as_ref())
),
McpServerConfig::Http(remote) => format!(
"http|{}|{}|{}|{}",
remote.url,
render_env_signature(&remote.headers),
remote.headers_helper.as_deref().unwrap_or(""),
render_oauth_signature(remote.oauth.as_ref())
),
McpServerConfig::Ws(ws) => format!(
"ws|{}|{}|{}",
ws.url,
render_env_signature(&ws.headers),
ws.headers_helper.as_deref().unwrap_or("")
),
McpServerConfig::Sdk(sdk) => format!("sdk|{}", sdk.name),
McpServerConfig::ManagedProxy(proxy) => {
format!("claudeai-proxy|{}|{}", proxy.url, proxy.id)
}
};
stable_hex_hash(&rendered)
}
fn render_command_signature(command: &[String]) -> String {
let escaped = command
.iter()
.map(|part| part.replace('\\', "\\\\").replace('|', "\\|"))
.collect::<Vec<_>>();
format!("[{}]", escaped.join("|"))
}
fn render_env_signature(map: &std::collections::BTreeMap<String, String>) -> String {
map.iter()
.map(|(key, value)| format!("{key}={value}"))
.collect::<Vec<_>>()
.join(";")
}
fn render_oauth_signature(oauth: Option<&crate::config::McpOAuthConfig>) -> String {
oauth.map_or_else(String::new, |oauth| {
format!(
"{}|{}|{}|{}",
oauth.client_id.as_deref().unwrap_or(""),
oauth
.callback_port
.map_or_else(String::new, |port| port.to_string()),
oauth.auth_server_metadata_url.as_deref().unwrap_or(""),
oauth.xaa.map_or_else(String::new, |flag| flag.to_string())
)
})
}
fn stable_hex_hash(value: &str) -> String {
let mut hash = 0xcbf2_9ce4_8422_2325_u64;
for byte in value.as_bytes() {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(0x0100_0000_01b3);
}
format!("{hash:016x}")
}
fn collapse_underscores(value: &str) -> String {
let mut collapsed = String::with_capacity(value.len());
let mut last_was_underscore = false;
for ch in value.chars() {
if ch == '_' {
if !last_was_underscore {
collapsed.push(ch);
}
last_was_underscore = true;
} else {
collapsed.push(ch);
last_was_underscore = false;
}
}
collapsed
}
fn percent_decode(value: &str) -> String {
let bytes = value.as_bytes();
let mut decoded = Vec::with_capacity(bytes.len());
let mut index = 0;
while index < bytes.len() {
match bytes[index] {
b'%' if index + 2 < bytes.len() => {
let hex = &value[index + 1..index + 3];
if let Ok(byte) = u8::from_str_radix(hex, 16) {
decoded.push(byte);
index += 3;
continue;
}
decoded.push(bytes[index]);
index += 1;
}
b'+' => {
decoded.push(b' ');
index += 1;
}
byte => {
decoded.push(byte);
index += 1;
}
}
}
String::from_utf8_lossy(&decoded).into_owned()
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use crate::config::{
ConfigSource, McpRemoteServerConfig, McpServerConfig, McpStdioServerConfig,
McpWebSocketServerConfig, ScopedMcpServerConfig,
};
use super::{
mcp_server_signature, mcp_tool_name, normalize_name_for_mcp, scoped_mcp_config_hash,
unwrap_ccr_proxy_url,
};
#[test]
fn normalizes_server_names_for_mcp_tooling() {
assert_eq!(normalize_name_for_mcp("github.com"), "github_com");
assert_eq!(normalize_name_for_mcp("tool name!"), "tool_name_");
assert_eq!(
normalize_name_for_mcp("claude.ai Example Server!!"),
"claude_ai_Example_Server"
);
assert_eq!(
mcp_tool_name("claude.ai Example Server", "weather tool"),
"mcp__claude_ai_Example_Server__weather_tool"
);
}
#[test]
fn unwraps_ccr_proxy_urls_for_signature_matching() {
let wrapped = "https://api.anthropic.com/v2/session_ingress/shttp/mcp/123?mcp_url=https%3A%2F%2Fvendor.example%2Fmcp&other=1";
assert_eq!(unwrap_ccr_proxy_url(wrapped), "https://vendor.example/mcp");
assert_eq!(
unwrap_ccr_proxy_url("https://vendor.example/mcp"),
"https://vendor.example/mcp"
);
}
#[test]
fn computes_signatures_for_stdio_and_remote_servers() {
let stdio = McpServerConfig::Stdio(McpStdioServerConfig {
command: "uvx".to_string(),
args: vec!["mcp-server".to_string()],
env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]),
tool_call_timeout_ms: None,
});
assert_eq!(
mcp_server_signature(&stdio),
Some("stdio:[uvx|mcp-server]".to_string())
);
let remote = McpServerConfig::Ws(McpWebSocketServerConfig {
url: "https://api.anthropic.com/v2/ccr-sessions/1?mcp_url=wss%3A%2F%2Fvendor.example%2Fmcp".to_string(),
headers: BTreeMap::new(),
headers_helper: None,
});
assert_eq!(
mcp_server_signature(&remote),
Some("url:wss://vendor.example/mcp".to_string())
);
}
#[test]
fn scoped_hash_ignores_scope_but_tracks_config_content() {
let base_config = McpServerConfig::Http(McpRemoteServerConfig {
url: "https://vendor.example/mcp".to_string(),
headers: BTreeMap::from([("Authorization".to_string(), "Bearer token".to_string())]),
headers_helper: Some("helper.sh".to_string()),
oauth: None,
});
let user = ScopedMcpServerConfig {
scope: ConfigSource::User,
config: base_config.clone(),
};
let local = ScopedMcpServerConfig {
scope: ConfigSource::Local,
config: base_config,
};
assert_eq!(
scoped_mcp_config_hash(&user),
scoped_mcp_config_hash(&local)
);
let changed = ScopedMcpServerConfig {
scope: ConfigSource::Local,
config: McpServerConfig::Http(McpRemoteServerConfig {
url: "https://vendor.example/v2/mcp".to_string(),
headers: BTreeMap::new(),
headers_helper: None,
oauth: None,
}),
};
assert_ne!(
scoped_mcp_config_hash(&user),
scoped_mcp_config_hash(&changed)
);
}
}
@@ -0,0 +1,248 @@
use std::collections::BTreeMap;
use crate::config::{McpOAuthConfig, McpServerConfig, ScopedMcpServerConfig};
use crate::mcp::{mcp_server_signature, mcp_tool_prefix, normalize_name_for_mcp};
pub const DEFAULT_MCP_TOOL_CALL_TIMEOUT_MS: u64 = 60_000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum McpClientTransport {
Stdio(McpStdioTransport),
Sse(McpRemoteTransport),
Http(McpRemoteTransport),
WebSocket(McpRemoteTransport),
Sdk(McpSdkTransport),
ManagedProxy(McpManagedProxyTransport),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct McpStdioTransport {
pub command: String,
pub args: Vec<String>,
pub env: BTreeMap<String, String>,
pub tool_call_timeout_ms: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct McpRemoteTransport {
pub url: String,
pub headers: BTreeMap<String, String>,
pub headers_helper: Option<String>,
pub auth: McpClientAuth,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct McpSdkTransport {
pub name: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct McpManagedProxyTransport {
pub url: String,
pub id: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum McpClientAuth {
None,
OAuth(McpOAuthConfig),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct McpClientBootstrap {
pub server_name: String,
pub normalized_name: String,
pub tool_prefix: String,
pub signature: Option<String>,
pub transport: McpClientTransport,
}
impl McpClientBootstrap {
#[must_use]
pub fn from_scoped_config(server_name: &str, config: &ScopedMcpServerConfig) -> Self {
Self {
server_name: server_name.to_string(),
normalized_name: normalize_name_for_mcp(server_name),
tool_prefix: mcp_tool_prefix(server_name),
signature: mcp_server_signature(&config.config),
transport: McpClientTransport::from_config(&config.config),
}
}
}
impl McpClientTransport {
#[must_use]
pub fn from_config(config: &McpServerConfig) -> Self {
match config {
McpServerConfig::Stdio(config) => Self::Stdio(McpStdioTransport {
command: config.command.clone(),
args: config.args.clone(),
env: config.env.clone(),
tool_call_timeout_ms: config.tool_call_timeout_ms,
}),
McpServerConfig::Sse(config) => Self::Sse(McpRemoteTransport {
url: config.url.clone(),
headers: config.headers.clone(),
headers_helper: config.headers_helper.clone(),
auth: McpClientAuth::from_oauth(config.oauth.clone()),
}),
McpServerConfig::Http(config) => Self::Http(McpRemoteTransport {
url: config.url.clone(),
headers: config.headers.clone(),
headers_helper: config.headers_helper.clone(),
auth: McpClientAuth::from_oauth(config.oauth.clone()),
}),
McpServerConfig::Ws(config) => Self::WebSocket(McpRemoteTransport {
url: config.url.clone(),
headers: config.headers.clone(),
headers_helper: config.headers_helper.clone(),
auth: McpClientAuth::None,
}),
McpServerConfig::Sdk(config) => Self::Sdk(McpSdkTransport {
name: config.name.clone(),
}),
McpServerConfig::ManagedProxy(config) => Self::ManagedProxy(McpManagedProxyTransport {
url: config.url.clone(),
id: config.id.clone(),
}),
}
}
}
impl McpStdioTransport {
#[must_use]
pub fn resolved_tool_call_timeout_ms(&self) -> u64 {
self.tool_call_timeout_ms
.unwrap_or(DEFAULT_MCP_TOOL_CALL_TIMEOUT_MS)
}
}
impl McpClientAuth {
#[must_use]
pub fn from_oauth(oauth: Option<McpOAuthConfig>) -> Self {
oauth.map_or(Self::None, Self::OAuth)
}
#[must_use]
pub const fn requires_user_auth(&self) -> bool {
matches!(self, Self::OAuth(_))
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use crate::config::{
ConfigSource, McpOAuthConfig, McpRemoteServerConfig, McpSdkServerConfig, McpServerConfig,
McpStdioServerConfig, McpWebSocketServerConfig, ScopedMcpServerConfig,
};
use super::{McpClientAuth, McpClientBootstrap, McpClientTransport};
#[test]
fn bootstraps_stdio_servers_into_transport_targets() {
let config = ScopedMcpServerConfig {
scope: ConfigSource::User,
config: McpServerConfig::Stdio(McpStdioServerConfig {
command: "uvx".to_string(),
args: vec!["mcp-server".to_string()],
env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]),
tool_call_timeout_ms: Some(15_000),
}),
};
let bootstrap = McpClientBootstrap::from_scoped_config("stdio-server", &config);
assert_eq!(bootstrap.normalized_name, "stdio-server");
assert_eq!(bootstrap.tool_prefix, "mcp__stdio-server__");
assert_eq!(
bootstrap.signature.as_deref(),
Some("stdio:[uvx|mcp-server]")
);
match bootstrap.transport {
McpClientTransport::Stdio(transport) => {
assert_eq!(transport.command, "uvx");
assert_eq!(transport.args, vec!["mcp-server"]);
assert_eq!(
transport.env.get("TOKEN").map(String::as_str),
Some("secret")
);
assert_eq!(transport.tool_call_timeout_ms, Some(15_000));
}
other => panic!("expected stdio transport, got {other:?}"),
}
}
#[test]
fn bootstraps_remote_servers_with_oauth_auth() {
let config = ScopedMcpServerConfig {
scope: ConfigSource::Project,
config: McpServerConfig::Http(McpRemoteServerConfig {
url: "https://vendor.example/mcp".to_string(),
headers: BTreeMap::from([("X-Test".to_string(), "1".to_string())]),
headers_helper: Some("helper.sh".to_string()),
oauth: Some(McpOAuthConfig {
client_id: Some("client-id".to_string()),
callback_port: Some(7777),
auth_server_metadata_url: Some(
"https://issuer.example/.well-known/oauth-authorization-server".to_string(),
),
xaa: Some(true),
}),
}),
};
let bootstrap = McpClientBootstrap::from_scoped_config("remote server", &config);
assert_eq!(bootstrap.normalized_name, "remote_server");
match bootstrap.transport {
McpClientTransport::Http(transport) => {
assert_eq!(transport.url, "https://vendor.example/mcp");
assert_eq!(transport.headers_helper.as_deref(), Some("helper.sh"));
assert!(transport.auth.requires_user_auth());
match transport.auth {
McpClientAuth::OAuth(oauth) => {
assert_eq!(oauth.client_id.as_deref(), Some("client-id"));
}
other @ McpClientAuth::None => panic!("expected oauth auth, got {other:?}"),
}
}
other => panic!("expected http transport, got {other:?}"),
}
}
#[test]
fn bootstraps_websocket_and_sdk_transports_without_oauth() {
let ws = ScopedMcpServerConfig {
scope: ConfigSource::Local,
config: McpServerConfig::Ws(McpWebSocketServerConfig {
url: "wss://vendor.example/mcp".to_string(),
headers: BTreeMap::new(),
headers_helper: None,
}),
};
let sdk = ScopedMcpServerConfig {
scope: ConfigSource::Local,
config: McpServerConfig::Sdk(McpSdkServerConfig {
name: "sdk-server".to_string(),
}),
};
let ws_bootstrap = McpClientBootstrap::from_scoped_config("ws server", &ws);
match ws_bootstrap.transport {
McpClientTransport::WebSocket(transport) => {
assert_eq!(transport.url, "wss://vendor.example/mcp");
assert!(!transport.auth.requires_user_auth());
}
other => panic!("expected websocket transport, got {other:?}"),
}
let sdk_bootstrap = McpClientBootstrap::from_scoped_config("sdk server", &sdk);
assert_eq!(sdk_bootstrap.signature, None);
match sdk_bootstrap.transport {
McpClientTransport::Sdk(transport) => {
assert_eq!(transport.name, "sdk-server");
}
other => panic!("expected sdk transport, got {other:?}"),
}
}
}
@@ -0,0 +1,843 @@
#![allow(clippy::unnested_or_patterns, clippy::map_unwrap_or)]
use std::collections::{BTreeMap, BTreeSet};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
fn now_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum McpLifecyclePhase {
ConfigLoad,
ServerRegistration,
SpawnConnect,
InitializeHandshake,
ToolDiscovery,
ResourceDiscovery,
Ready,
Invocation,
ErrorSurfacing,
Shutdown,
Cleanup,
}
impl McpLifecyclePhase {
#[must_use]
pub fn all() -> [Self; 11] {
[
Self::ConfigLoad,
Self::ServerRegistration,
Self::SpawnConnect,
Self::InitializeHandshake,
Self::ToolDiscovery,
Self::ResourceDiscovery,
Self::Ready,
Self::Invocation,
Self::ErrorSurfacing,
Self::Shutdown,
Self::Cleanup,
]
}
}
impl std::fmt::Display for McpLifecyclePhase {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ConfigLoad => write!(f, "config_load"),
Self::ServerRegistration => write!(f, "server_registration"),
Self::SpawnConnect => write!(f, "spawn_connect"),
Self::InitializeHandshake => write!(f, "initialize_handshake"),
Self::ToolDiscovery => write!(f, "tool_discovery"),
Self::ResourceDiscovery => write!(f, "resource_discovery"),
Self::Ready => write!(f, "ready"),
Self::Invocation => write!(f, "invocation"),
Self::ErrorSurfacing => write!(f, "error_surfacing"),
Self::Shutdown => write!(f, "shutdown"),
Self::Cleanup => write!(f, "cleanup"),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct McpErrorSurface {
pub phase: McpLifecyclePhase,
pub server_name: Option<String>,
pub message: String,
pub context: BTreeMap<String, String>,
pub recoverable: bool,
pub timestamp: u64,
}
impl McpErrorSurface {
#[must_use]
pub fn new(
phase: McpLifecyclePhase,
server_name: Option<String>,
message: impl Into<String>,
context: BTreeMap<String, String>,
recoverable: bool,
) -> Self {
Self {
phase,
server_name,
message: message.into(),
context,
recoverable,
timestamp: now_secs(),
}
}
}
impl std::fmt::Display for McpErrorSurface {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"MCP lifecycle error during {}: {}",
self.phase, self.message
)?;
if let Some(server_name) = &self.server_name {
write!(f, " (server: {server_name})")?;
}
if !self.context.is_empty() {
write!(f, " with context {:?}", self.context)?;
}
if self.recoverable {
write!(f, " [recoverable]")?;
}
Ok(())
}
}
impl std::error::Error for McpErrorSurface {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum McpPhaseResult {
Success {
phase: McpLifecyclePhase,
duration: Duration,
},
Failure {
phase: McpLifecyclePhase,
error: McpErrorSurface,
},
Timeout {
phase: McpLifecyclePhase,
waited: Duration,
error: McpErrorSurface,
},
}
impl McpPhaseResult {
#[must_use]
pub fn phase(&self) -> McpLifecyclePhase {
match self {
Self::Success { phase, .. }
| Self::Failure { phase, .. }
| Self::Timeout { phase, .. } => *phase,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct McpLifecycleState {
current_phase: Option<McpLifecyclePhase>,
phase_errors: BTreeMap<McpLifecyclePhase, Vec<McpErrorSurface>>,
phase_timestamps: BTreeMap<McpLifecyclePhase, u64>,
phase_results: Vec<McpPhaseResult>,
}
impl McpLifecycleState {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn current_phase(&self) -> Option<McpLifecyclePhase> {
self.current_phase
}
#[must_use]
pub fn errors_for_phase(&self, phase: McpLifecyclePhase) -> &[McpErrorSurface] {
self.phase_errors
.get(&phase)
.map(Vec::as_slice)
.unwrap_or(&[])
}
#[must_use]
pub fn results(&self) -> &[McpPhaseResult] {
&self.phase_results
}
#[must_use]
pub fn phase_timestamps(&self) -> &BTreeMap<McpLifecyclePhase, u64> {
&self.phase_timestamps
}
#[must_use]
pub fn phase_timestamp(&self, phase: McpLifecyclePhase) -> Option<u64> {
self.phase_timestamps.get(&phase).copied()
}
fn record_phase(&mut self, phase: McpLifecyclePhase) {
self.current_phase = Some(phase);
self.phase_timestamps.insert(phase, now_secs());
}
fn record_error(&mut self, error: McpErrorSurface) {
self.phase_errors
.entry(error.phase)
.or_default()
.push(error);
}
fn record_result(&mut self, result: McpPhaseResult) {
self.phase_results.push(result);
}
fn can_resume_after_error(&self) -> bool {
match self.phase_results.last() {
Some(McpPhaseResult::Failure { error, .. } | McpPhaseResult::Timeout { error, .. }) => {
error.recoverable
}
_ => false,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct McpFailedServer {
pub server_name: String,
pub phase: McpLifecyclePhase,
pub error: McpErrorSurface,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct McpDegradedReport {
pub working_servers: Vec<String>,
pub failed_servers: Vec<McpFailedServer>,
pub available_tools: Vec<String>,
pub missing_tools: Vec<String>,
}
impl McpDegradedReport {
#[must_use]
pub fn new(
working_servers: Vec<String>,
failed_servers: Vec<McpFailedServer>,
available_tools: Vec<String>,
expected_tools: Vec<String>,
) -> Self {
let working_servers = dedupe_sorted(working_servers);
let available_tools = dedupe_sorted(available_tools);
let available_tool_set: BTreeSet<_> = available_tools.iter().cloned().collect();
let expected_tools = dedupe_sorted(expected_tools);
let missing_tools = expected_tools
.into_iter()
.filter(|tool| !available_tool_set.contains(tool))
.collect();
Self {
working_servers,
failed_servers,
available_tools,
missing_tools,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct McpLifecycleValidator {
state: McpLifecycleState,
}
impl McpLifecycleValidator {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn state(&self) -> &McpLifecycleState {
&self.state
}
#[must_use]
pub fn validate_phase_transition(from: McpLifecyclePhase, to: McpLifecyclePhase) -> bool {
match (from, to) {
(McpLifecyclePhase::ConfigLoad, McpLifecyclePhase::ServerRegistration)
| (McpLifecyclePhase::ServerRegistration, McpLifecyclePhase::SpawnConnect)
| (McpLifecyclePhase::SpawnConnect, McpLifecyclePhase::InitializeHandshake)
| (McpLifecyclePhase::InitializeHandshake, McpLifecyclePhase::ToolDiscovery)
| (McpLifecyclePhase::ToolDiscovery, McpLifecyclePhase::ResourceDiscovery)
| (McpLifecyclePhase::ToolDiscovery, McpLifecyclePhase::Ready)
| (McpLifecyclePhase::ResourceDiscovery, McpLifecyclePhase::Ready)
| (McpLifecyclePhase::Ready, McpLifecyclePhase::Invocation)
| (McpLifecyclePhase::Invocation, McpLifecyclePhase::Ready)
| (McpLifecyclePhase::ErrorSurfacing, McpLifecyclePhase::Ready)
| (McpLifecyclePhase::ErrorSurfacing, McpLifecyclePhase::Shutdown)
| (McpLifecyclePhase::Shutdown, McpLifecyclePhase::Cleanup) => true,
(_, McpLifecyclePhase::Shutdown) => from != McpLifecyclePhase::Cleanup,
(_, McpLifecyclePhase::ErrorSurfacing) => {
from != McpLifecyclePhase::Cleanup && from != McpLifecyclePhase::Shutdown
}
_ => false,
}
}
pub fn run_phase(&mut self, phase: McpLifecyclePhase) -> McpPhaseResult {
let started = Instant::now();
if let Some(current_phase) = self.state.current_phase() {
if current_phase == McpLifecyclePhase::ErrorSurfacing
&& phase == McpLifecyclePhase::Ready
&& !self.state.can_resume_after_error()
{
return self.record_failure(McpErrorSurface::new(
phase,
None,
"cannot return to ready after a non-recoverable MCP lifecycle failure",
BTreeMap::from([
("from".to_string(), current_phase.to_string()),
("to".to_string(), phase.to_string()),
]),
false,
));
}
if !Self::validate_phase_transition(current_phase, phase) {
return self.record_failure(McpErrorSurface::new(
phase,
None,
format!("invalid MCP lifecycle transition from {current_phase} to {phase}"),
BTreeMap::from([
("from".to_string(), current_phase.to_string()),
("to".to_string(), phase.to_string()),
]),
false,
));
}
} else if phase != McpLifecyclePhase::ConfigLoad {
return self.record_failure(McpErrorSurface::new(
phase,
None,
format!("invalid initial MCP lifecycle phase {phase}"),
BTreeMap::from([("phase".to_string(), phase.to_string())]),
false,
));
}
self.state.record_phase(phase);
let result = McpPhaseResult::Success {
phase,
duration: started.elapsed(),
};
self.state.record_result(result.clone());
result
}
pub fn record_failure(&mut self, error: McpErrorSurface) -> McpPhaseResult {
let phase = error.phase;
self.state.record_error(error.clone());
self.state.record_phase(McpLifecyclePhase::ErrorSurfacing);
let result = McpPhaseResult::Failure { phase, error };
self.state.record_result(result.clone());
result
}
pub fn record_timeout(
&mut self,
phase: McpLifecyclePhase,
waited: Duration,
server_name: Option<String>,
mut context: BTreeMap<String, String>,
) -> McpPhaseResult {
context.insert("waited_ms".to_string(), waited.as_millis().to_string());
let error = McpErrorSurface::new(
phase,
server_name,
format!(
"MCP lifecycle phase {phase} timed out after {} ms",
waited.as_millis()
),
context,
true,
);
self.state.record_error(error.clone());
self.state.record_phase(McpLifecyclePhase::ErrorSurfacing);
let result = McpPhaseResult::Timeout {
phase,
waited,
error,
};
self.state.record_result(result.clone());
result
}
}
fn dedupe_sorted(mut values: Vec<String>) -> Vec<String> {
values.sort();
values.dedup();
values
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn phase_display_matches_serde_name() {
// given
let phases = McpLifecyclePhase::all();
// when
let serialized = phases
.into_iter()
.map(|phase| {
(
phase.to_string(),
serde_json::to_value(phase).expect("serialize phase"),
)
})
.collect::<Vec<_>>();
// then
for (display, json_value) in serialized {
assert_eq!(json_value, json!(display));
}
}
#[test]
fn given_startup_path_when_running_to_cleanup_then_each_control_transition_succeeds() {
// given
let mut validator = McpLifecycleValidator::new();
let phases = [
McpLifecyclePhase::ConfigLoad,
McpLifecyclePhase::ServerRegistration,
McpLifecyclePhase::SpawnConnect,
McpLifecyclePhase::InitializeHandshake,
McpLifecyclePhase::ToolDiscovery,
McpLifecyclePhase::ResourceDiscovery,
McpLifecyclePhase::Ready,
McpLifecyclePhase::Invocation,
McpLifecyclePhase::Ready,
McpLifecyclePhase::Shutdown,
McpLifecyclePhase::Cleanup,
];
// when
let results = phases
.into_iter()
.map(|phase| validator.run_phase(phase))
.collect::<Vec<_>>();
// then
assert!(results
.iter()
.all(|result| matches!(result, McpPhaseResult::Success { .. })));
assert_eq!(
validator.state().current_phase(),
Some(McpLifecyclePhase::Cleanup)
);
for phase in [
McpLifecyclePhase::ConfigLoad,
McpLifecyclePhase::ServerRegistration,
McpLifecyclePhase::SpawnConnect,
McpLifecyclePhase::InitializeHandshake,
McpLifecyclePhase::ToolDiscovery,
McpLifecyclePhase::ResourceDiscovery,
McpLifecyclePhase::Ready,
McpLifecyclePhase::Invocation,
McpLifecyclePhase::Shutdown,
McpLifecyclePhase::Cleanup,
] {
assert!(validator.state().phase_timestamp(phase).is_some());
}
}
#[test]
fn given_tool_discovery_when_resource_discovery_is_skipped_then_ready_is_still_allowed() {
// given
let mut validator = McpLifecycleValidator::new();
for phase in [
McpLifecyclePhase::ConfigLoad,
McpLifecyclePhase::ServerRegistration,
McpLifecyclePhase::SpawnConnect,
McpLifecyclePhase::InitializeHandshake,
McpLifecyclePhase::ToolDiscovery,
] {
let result = validator.run_phase(phase);
assert!(matches!(result, McpPhaseResult::Success { .. }));
}
// when
let result = validator.run_phase(McpLifecyclePhase::Ready);
// then
assert!(matches!(result, McpPhaseResult::Success { .. }));
assert_eq!(
validator.state().current_phase(),
Some(McpLifecyclePhase::Ready)
);
}
#[test]
fn validates_expected_phase_transitions() {
// given
let valid_transitions = [
(
McpLifecyclePhase::ConfigLoad,
McpLifecyclePhase::ServerRegistration,
),
(
McpLifecyclePhase::ServerRegistration,
McpLifecyclePhase::SpawnConnect,
),
(
McpLifecyclePhase::SpawnConnect,
McpLifecyclePhase::InitializeHandshake,
),
(
McpLifecyclePhase::InitializeHandshake,
McpLifecyclePhase::ToolDiscovery,
),
(
McpLifecyclePhase::ToolDiscovery,
McpLifecyclePhase::ResourceDiscovery,
),
(McpLifecyclePhase::ToolDiscovery, McpLifecyclePhase::Ready),
(
McpLifecyclePhase::ResourceDiscovery,
McpLifecyclePhase::Ready,
),
(McpLifecyclePhase::Ready, McpLifecyclePhase::Invocation),
(McpLifecyclePhase::Invocation, McpLifecyclePhase::Ready),
(McpLifecyclePhase::Ready, McpLifecyclePhase::Shutdown),
(
McpLifecyclePhase::Invocation,
McpLifecyclePhase::ErrorSurfacing,
),
(
McpLifecyclePhase::ErrorSurfacing,
McpLifecyclePhase::Shutdown,
),
(McpLifecyclePhase::Shutdown, McpLifecyclePhase::Cleanup),
];
// when / then
for (from, to) in valid_transitions {
assert!(McpLifecycleValidator::validate_phase_transition(from, to));
}
assert!(!McpLifecycleValidator::validate_phase_transition(
McpLifecyclePhase::Ready,
McpLifecyclePhase::ConfigLoad,
));
assert!(!McpLifecycleValidator::validate_phase_transition(
McpLifecyclePhase::Cleanup,
McpLifecyclePhase::Ready,
));
}
#[test]
fn given_invalid_transition_when_running_phase_then_structured_failure_is_recorded() {
// given
let mut validator = McpLifecycleValidator::new();
let _ = validator.run_phase(McpLifecyclePhase::ConfigLoad);
let _ = validator.run_phase(McpLifecyclePhase::ServerRegistration);
// when
let result = validator.run_phase(McpLifecyclePhase::Ready);
// then
match result {
McpPhaseResult::Failure { phase, error } => {
assert_eq!(phase, McpLifecyclePhase::Ready);
assert!(!error.recoverable);
assert_eq!(error.phase, McpLifecyclePhase::Ready);
assert_eq!(
error.context.get("from").map(String::as_str),
Some("server_registration")
);
assert_eq!(error.context.get("to").map(String::as_str), Some("ready"));
}
other => panic!("expected failure result, got {other:?}"),
}
assert_eq!(
validator.state().current_phase(),
Some(McpLifecyclePhase::ErrorSurfacing)
);
assert_eq!(
validator
.state()
.errors_for_phase(McpLifecyclePhase::Ready)
.len(),
1
);
}
#[test]
fn given_each_phase_when_failure_is_recorded_then_error_is_tracked_per_phase() {
// given
let mut validator = McpLifecycleValidator::new();
// when / then
for phase in McpLifecyclePhase::all() {
let result = validator.record_failure(McpErrorSurface::new(
phase,
Some("alpha".to_string()),
format!("failure at {phase}"),
BTreeMap::from([("server".to_string(), "alpha".to_string())]),
phase == McpLifecyclePhase::ResourceDiscovery,
));
match result {
McpPhaseResult::Failure {
phase: failed_phase,
error,
} => {
assert_eq!(failed_phase, phase);
assert_eq!(error.phase, phase);
assert_eq!(
error.recoverable,
phase == McpLifecyclePhase::ResourceDiscovery
);
}
other => panic!("expected failure result, got {other:?}"),
}
assert_eq!(validator.state().errors_for_phase(phase).len(), 1);
}
}
#[test]
fn given_spawn_connect_timeout_when_recorded_then_waited_duration_is_preserved() {
// given
let mut validator = McpLifecycleValidator::new();
let waited = Duration::from_millis(250);
// when
let result = validator.record_timeout(
McpLifecyclePhase::SpawnConnect,
waited,
Some("alpha".to_string()),
BTreeMap::from([("attempt".to_string(), "1".to_string())]),
);
// then
match result {
McpPhaseResult::Timeout {
phase,
waited: actual,
error,
} => {
assert_eq!(phase, McpLifecyclePhase::SpawnConnect);
assert_eq!(actual, waited);
assert!(error.recoverable);
assert_eq!(error.server_name.as_deref(), Some("alpha"));
}
other => panic!("expected timeout result, got {other:?}"),
}
let errors = validator
.state()
.errors_for_phase(McpLifecyclePhase::SpawnConnect);
assert_eq!(errors.len(), 1);
assert_eq!(
errors[0].context.get("waited_ms").map(String::as_str),
Some("250")
);
assert_eq!(
validator.state().current_phase(),
Some(McpLifecyclePhase::ErrorSurfacing)
);
}
#[test]
fn given_partial_server_health_when_building_degraded_report_then_missing_tools_are_reported() {
// given
let failed = vec![McpFailedServer {
server_name: "broken".to_string(),
phase: McpLifecyclePhase::InitializeHandshake,
error: McpErrorSurface::new(
McpLifecyclePhase::InitializeHandshake,
Some("broken".to_string()),
"initialize failed",
BTreeMap::from([("reason".to_string(), "broken pipe".to_string())]),
false,
),
}];
// when
let report = McpDegradedReport::new(
vec!["alpha".to_string(), "beta".to_string(), "alpha".to_string()],
failed,
vec![
"alpha.echo".to_string(),
"beta.search".to_string(),
"alpha.echo".to_string(),
],
vec![
"alpha.echo".to_string(),
"beta.search".to_string(),
"broken.fetch".to_string(),
],
);
// then
assert_eq!(
report.working_servers,
vec!["alpha".to_string(), "beta".to_string()]
);
assert_eq!(report.failed_servers.len(), 1);
assert_eq!(report.failed_servers[0].server_name, "broken");
assert_eq!(
report.available_tools,
vec!["alpha.echo".to_string(), "beta.search".to_string()]
);
assert_eq!(report.missing_tools, vec!["broken.fetch".to_string()]);
}
#[test]
fn given_failure_during_resource_discovery_when_shutting_down_then_cleanup_still_succeeds() {
// given
let mut validator = McpLifecycleValidator::new();
for phase in [
McpLifecyclePhase::ConfigLoad,
McpLifecyclePhase::ServerRegistration,
McpLifecyclePhase::SpawnConnect,
McpLifecyclePhase::InitializeHandshake,
McpLifecyclePhase::ToolDiscovery,
] {
let result = validator.run_phase(phase);
assert!(matches!(result, McpPhaseResult::Success { .. }));
}
let _ = validator.record_failure(McpErrorSurface::new(
McpLifecyclePhase::ResourceDiscovery,
Some("alpha".to_string()),
"resource listing failed",
BTreeMap::from([("reason".to_string(), "timeout".to_string())]),
true,
));
// when
let shutdown = validator.run_phase(McpLifecyclePhase::Shutdown);
let cleanup = validator.run_phase(McpLifecyclePhase::Cleanup);
// then
assert!(matches!(shutdown, McpPhaseResult::Success { .. }));
assert!(matches!(cleanup, McpPhaseResult::Success { .. }));
assert_eq!(
validator.state().current_phase(),
Some(McpLifecyclePhase::Cleanup)
);
assert!(validator
.state()
.phase_timestamp(McpLifecyclePhase::ErrorSurfacing)
.is_some());
}
#[test]
fn error_surface_display_includes_phase_server_and_recoverable_flag() {
// given
let error = McpErrorSurface::new(
McpLifecyclePhase::SpawnConnect,
Some("alpha".to_string()),
"process exited early",
BTreeMap::from([("exit_code".to_string(), "1".to_string())]),
true,
);
// when
let rendered = error.to_string();
// then
assert!(rendered.contains("spawn_connect"));
assert!(rendered.contains("process exited early"));
assert!(rendered.contains("server: alpha"));
assert!(rendered.contains("recoverable"));
let trait_object: &dyn std::error::Error = &error;
assert_eq!(trait_object.to_string(), rendered);
}
#[test]
fn given_nonrecoverable_failure_when_returning_to_ready_then_validator_rejects_resume() {
// given
let mut validator = McpLifecycleValidator::new();
for phase in [
McpLifecyclePhase::ConfigLoad,
McpLifecyclePhase::ServerRegistration,
McpLifecyclePhase::SpawnConnect,
McpLifecyclePhase::InitializeHandshake,
McpLifecyclePhase::ToolDiscovery,
McpLifecyclePhase::Ready,
] {
let result = validator.run_phase(phase);
assert!(matches!(result, McpPhaseResult::Success { .. }));
}
let _ = validator.record_failure(McpErrorSurface::new(
McpLifecyclePhase::Invocation,
Some("alpha".to_string()),
"tool call corrupted the session",
BTreeMap::from([("reason".to_string(), "invalid frame".to_string())]),
false,
));
// when
let result = validator.run_phase(McpLifecyclePhase::Ready);
// then
match result {
McpPhaseResult::Failure { phase, error } => {
assert_eq!(phase, McpLifecyclePhase::Ready);
assert!(!error.recoverable);
assert!(error.message.contains("non-recoverable"));
}
other => panic!("expected failure result, got {other:?}"),
}
assert_eq!(
validator.state().current_phase(),
Some(McpLifecyclePhase::ErrorSurfacing)
);
}
#[test]
fn given_recoverable_failure_when_returning_to_ready_then_validator_allows_resume() {
// given
let mut validator = McpLifecycleValidator::new();
for phase in [
McpLifecyclePhase::ConfigLoad,
McpLifecyclePhase::ServerRegistration,
McpLifecyclePhase::SpawnConnect,
McpLifecyclePhase::InitializeHandshake,
McpLifecyclePhase::ToolDiscovery,
McpLifecyclePhase::Ready,
] {
let result = validator.run_phase(phase);
assert!(matches!(result, McpPhaseResult::Success { .. }));
}
let _ = validator.record_failure(McpErrorSurface::new(
McpLifecyclePhase::Invocation,
Some("alpha".to_string()),
"tool call failed but can be retried",
BTreeMap::from([("reason".to_string(), "upstream timeout".to_string())]),
true,
));
// when
let result = validator.run_phase(McpLifecyclePhase::Ready);
// then
assert!(matches!(result, McpPhaseResult::Success { .. }));
assert_eq!(
validator.state().current_phase(),
Some(McpLifecyclePhase::Ready)
);
}
}
@@ -0,0 +1,440 @@
//! Minimal Model Context Protocol (MCP) server.
//!
//! Implements a newline-safe, LSP-framed JSON-RPC server over stdio that
//! answers `initialize`, `tools/list`, and `tools/call` requests. The framing
//! matches the client transport implemented in [`crate::mcp_stdio`] so this
//! server can be driven by either an external MCP client (e.g. Claude
//! Desktop) or `claw`'s own [`McpServerManager`](crate::McpServerManager).
//!
//! The server is intentionally small: it exposes a list of pre-built
//! [`McpTool`] descriptors and delegates `tools/call` to a caller-supplied
//! handler. Tool execution itself lives in the `tools` crate; this module is
//! purely the transport + dispatch loop.
//!
//! [`McpTool`]: crate::mcp_stdio::McpTool
use std::io;
use serde_json::{json, Value as JsonValue};
use tokio::io::{
stdin, stdout, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, Stdin, Stdout,
};
use crate::mcp_stdio::{
JsonRpcError, JsonRpcId, JsonRpcRequest, JsonRpcResponse, McpInitializeResult,
McpInitializeServerInfo, McpListToolsResult, McpTool, McpToolCallContent, McpToolCallParams,
McpToolCallResult,
};
/// Protocol version the server advertises during `initialize`.
///
/// Matches the version used by the built-in client in
/// [`crate::mcp_stdio`], so the two stay in lockstep.
pub const MCP_SERVER_PROTOCOL_VERSION: &str = "2025-03-26";
/// Synchronous handler invoked for every `tools/call` request.
///
/// Returning `Ok(text)` yields a single `text` content block and
/// `isError: false`. Returning `Err(message)` yields a `text` block with the
/// error and `isError: true`, mirroring the error-surfacing convention used
/// elsewhere in claw.
pub type ToolCallHandler =
Box<dyn Fn(&str, &JsonValue) -> Result<String, String> + Send + Sync + 'static>;
/// Configuration for an [`McpServer`] instance.
///
/// Named `McpServerSpec` rather than `McpServerConfig` to avoid colliding
/// with the existing client-side [`crate::config::McpServerConfig`] that
/// describes *remote* MCP servers the runtime connects to.
pub struct McpServerSpec {
/// Name advertised in the `serverInfo` field of the `initialize` response.
pub server_name: String,
/// Version advertised in the `serverInfo` field of the `initialize`
/// response.
pub server_version: String,
/// Tool descriptors returned for `tools/list`.
pub tools: Vec<McpTool>,
/// Handler invoked for `tools/call`.
pub tool_handler: ToolCallHandler,
}
/// Minimal MCP stdio server.
///
/// The server runs a blocking read/dispatch/write loop over the current
/// process's stdin/stdout, terminating cleanly when the peer closes the
/// stream.
pub struct McpServer {
spec: McpServerSpec,
stdin: BufReader<Stdin>,
stdout: Stdout,
}
impl McpServer {
#[must_use]
pub fn new(spec: McpServerSpec) -> Self {
Self {
spec,
stdin: BufReader::new(stdin()),
stdout: stdout(),
}
}
/// Runs the server until the client closes stdin.
///
/// Returns `Ok(())` on clean EOF; any other I/O error is propagated so
/// callers can log and exit non-zero.
pub async fn run(&mut self) -> io::Result<()> {
loop {
let Some(payload) = read_frame(&mut self.stdin).await? else {
return Ok(());
};
// Requests and notifications share a wire format; the absence of
// `id` distinguishes notifications, which must never receive a
// response.
let message: JsonValue = match serde_json::from_slice(&payload) {
Ok(value) => value,
Err(error) => {
// Parse error with null id per JSON-RPC 2.0 §4.2.
let response = JsonRpcResponse::<JsonValue> {
jsonrpc: "2.0".to_string(),
id: JsonRpcId::Null,
result: None,
error: Some(JsonRpcError {
code: -32700,
message: format!("parse error: {error}"),
data: None,
}),
};
write_response(&mut self.stdout, &response).await?;
continue;
}
};
if message.get("id").is_none() {
// Notification: dispatch for side effects only (e.g. log),
// but send no reply.
continue;
}
let request: JsonRpcRequest<JsonValue> = match serde_json::from_value(message) {
Ok(request) => request,
Err(error) => {
let response = JsonRpcResponse::<JsonValue> {
jsonrpc: "2.0".to_string(),
id: JsonRpcId::Null,
result: None,
error: Some(JsonRpcError {
code: -32600,
message: format!("invalid request: {error}"),
data: None,
}),
};
write_response(&mut self.stdout, &response).await?;
continue;
}
};
let response = self.dispatch(request);
write_response(&mut self.stdout, &response).await?;
}
}
fn dispatch(&self, request: JsonRpcRequest<JsonValue>) -> JsonRpcResponse<JsonValue> {
let id = request.id.clone();
match request.method.as_str() {
"initialize" => self.handle_initialize(id),
"tools/list" => self.handle_tools_list(id),
"tools/call" => self.handle_tools_call(id, request.params),
other => JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id,
result: None,
error: Some(JsonRpcError {
code: -32601,
message: format!("method not found: {other}"),
data: None,
}),
},
}
}
fn handle_initialize(&self, id: JsonRpcId) -> JsonRpcResponse<JsonValue> {
let result = McpInitializeResult {
protocol_version: MCP_SERVER_PROTOCOL_VERSION.to_string(),
capabilities: json!({ "tools": {} }),
server_info: McpInitializeServerInfo {
name: self.spec.server_name.clone(),
version: self.spec.server_version.clone(),
},
};
JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id,
result: serde_json::to_value(result).ok(),
error: None,
}
}
fn handle_tools_list(&self, id: JsonRpcId) -> JsonRpcResponse<JsonValue> {
let result = McpListToolsResult {
tools: self.spec.tools.clone(),
next_cursor: None,
};
JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id,
result: serde_json::to_value(result).ok(),
error: None,
}
}
fn handle_tools_call(
&self,
id: JsonRpcId,
params: Option<JsonValue>,
) -> JsonRpcResponse<JsonValue> {
let Some(params) = params else {
return invalid_params_response(id, "missing params for tools/call");
};
let call: McpToolCallParams = match serde_json::from_value(params) {
Ok(value) => value,
Err(error) => {
return invalid_params_response(id, &format!("invalid tools/call params: {error}"));
}
};
let arguments = call.arguments.unwrap_or_else(|| json!({}));
let tool_result = (self.spec.tool_handler)(&call.name, &arguments);
let (text, is_error) = match tool_result {
Ok(text) => (text, false),
Err(message) => (message, true),
};
let mut data = std::collections::BTreeMap::new();
data.insert("text".to_string(), JsonValue::String(text));
let call_result = McpToolCallResult {
content: vec![McpToolCallContent {
kind: "text".to_string(),
data,
}],
structured_content: None,
is_error: Some(is_error),
meta: None,
};
JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id,
result: serde_json::to_value(call_result).ok(),
error: None,
}
}
}
fn invalid_params_response(id: JsonRpcId, message: &str) -> JsonRpcResponse<JsonValue> {
JsonRpcResponse {
jsonrpc: "2.0".to_string(),
id,
result: None,
error: Some(JsonRpcError {
code: -32602,
message: message.to_string(),
data: None,
}),
}
}
/// Reads a single LSP-framed JSON-RPC payload from `reader`.
///
/// Returns `Ok(None)` on clean EOF before any header bytes have been read,
/// matching how [`crate::mcp_stdio::McpStdioProcess`] treats stream closure.
async fn read_frame(reader: &mut BufReader<Stdin>) -> io::Result<Option<Vec<u8>>> {
let mut content_length: Option<usize> = None;
let mut first_header = true;
loop {
let mut line = String::new();
let bytes_read = reader.read_line(&mut line).await?;
if bytes_read == 0 {
if first_header {
return Ok(None);
}
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"MCP stdio stream closed while reading headers",
));
}
first_header = false;
if line == "\r\n" || line == "\n" {
break;
}
let header = line.trim_end_matches(['\r', '\n']);
if let Some((name, value)) = header.split_once(':') {
if name.trim().eq_ignore_ascii_case("Content-Length") {
let parsed = value
.trim()
.parse::<usize>()
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
content_length = Some(parsed);
}
}
}
let content_length = content_length.ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidData, "missing Content-Length header")
})?;
let mut payload = vec![0_u8; content_length];
reader.read_exact(&mut payload).await?;
Ok(Some(payload))
}
async fn write_response(
stdout: &mut Stdout,
response: &JsonRpcResponse<JsonValue>,
) -> io::Result<()> {
let body = serde_json::to_vec(response)
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
let header = format!("Content-Length: {}\r\n\r\n", body.len());
stdout.write_all(header.as_bytes()).await?;
stdout.write_all(&body).await?;
stdout.flush().await
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn dispatch_initialize_returns_server_info() {
let server = McpServer {
spec: McpServerSpec {
server_name: "test".to_string(),
server_version: "9.9.9".to_string(),
tools: Vec::new(),
tool_handler: Box::new(|_, _| Ok(String::new())),
},
stdin: BufReader::new(stdin()),
stdout: stdout(),
};
let request = JsonRpcRequest::<JsonValue> {
jsonrpc: "2.0".to_string(),
id: JsonRpcId::Number(1),
method: "initialize".to_string(),
params: None,
};
let response = server.dispatch(request);
assert_eq!(response.id, JsonRpcId::Number(1));
assert!(response.error.is_none());
let result = response.result.expect("initialize result");
assert_eq!(result["protocolVersion"], MCP_SERVER_PROTOCOL_VERSION);
assert_eq!(result["serverInfo"]["name"], "test");
assert_eq!(result["serverInfo"]["version"], "9.9.9");
}
#[test]
fn dispatch_tools_list_returns_registered_tools() {
let tool = McpTool {
name: "echo".to_string(),
description: Some("Echo".to_string()),
input_schema: Some(json!({"type": "object"})),
annotations: None,
meta: None,
};
let server = McpServer {
spec: McpServerSpec {
server_name: "test".to_string(),
server_version: "0.0.0".to_string(),
tools: vec![tool.clone()],
tool_handler: Box::new(|_, _| Ok(String::new())),
},
stdin: BufReader::new(stdin()),
stdout: stdout(),
};
let request = JsonRpcRequest::<JsonValue> {
jsonrpc: "2.0".to_string(),
id: JsonRpcId::Number(2),
method: "tools/list".to_string(),
params: None,
};
let response = server.dispatch(request);
assert!(response.error.is_none());
let result = response.result.expect("tools/list result");
assert_eq!(result["tools"][0]["name"], "echo");
}
#[test]
fn dispatch_tools_call_wraps_handler_output() {
let server = McpServer {
spec: McpServerSpec {
server_name: "test".to_string(),
server_version: "0.0.0".to_string(),
tools: Vec::new(),
tool_handler: Box::new(|name, args| Ok(format!("called {name} with {args}"))),
},
stdin: BufReader::new(stdin()),
stdout: stdout(),
};
let request = JsonRpcRequest::<JsonValue> {
jsonrpc: "2.0".to_string(),
id: JsonRpcId::Number(3),
method: "tools/call".to_string(),
params: Some(json!({
"name": "echo",
"arguments": {"text": "hi"}
})),
};
let response = server.dispatch(request);
assert!(response.error.is_none());
let result = response.result.expect("tools/call result");
assert_eq!(result["isError"], false);
assert_eq!(result["content"][0]["type"], "text");
assert!(result["content"][0]["text"]
.as_str()
.unwrap()
.starts_with("called echo"));
}
#[test]
fn dispatch_tools_call_surfaces_handler_error() {
let server = McpServer {
spec: McpServerSpec {
server_name: "test".to_string(),
server_version: "0.0.0".to_string(),
tools: Vec::new(),
tool_handler: Box::new(|_, _| Err("boom".to_string())),
},
stdin: BufReader::new(stdin()),
stdout: stdout(),
};
let request = JsonRpcRequest::<JsonValue> {
jsonrpc: "2.0".to_string(),
id: JsonRpcId::Number(4),
method: "tools/call".to_string(),
params: Some(json!({"name": "broken"})),
};
let response = server.dispatch(request);
let result = response.result.expect("tools/call result");
assert_eq!(result["isError"], true);
assert_eq!(result["content"][0]["text"], "boom");
}
#[test]
fn dispatch_unknown_method_returns_method_not_found() {
let server = McpServer {
spec: McpServerSpec {
server_name: "test".to_string(),
server_version: "0.0.0".to_string(),
tools: Vec::new(),
tool_handler: Box::new(|_, _| Ok(String::new())),
},
stdin: BufReader::new(stdin()),
stdout: stdout(),
};
let request = JsonRpcRequest::<JsonValue> {
jsonrpc: "2.0".to_string(),
id: JsonRpcId::Number(5),
method: "nonsense".to_string(),
params: None,
};
let response = server.dispatch(request);
let error = response.error.expect("error payload");
assert_eq!(error.code, -32601);
}
}
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,920 @@
#![allow(
clippy::await_holding_lock,
clippy::doc_markdown,
clippy::match_same_arms,
clippy::must_use_candidate,
clippy::uninlined_format_args,
clippy::unnested_or_patterns
)]
//! Bridge between MCP tool surface (ListMcpResources, ReadMcpResource, McpAuth, MCP)
//! and the existing McpServerManager runtime.
//!
//! Provides a stateful client registry that tool handlers can use to
//! connect to MCP servers and invoke their capabilities.
use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};
use crate::mcp::mcp_tool_name;
use crate::mcp_stdio::McpServerManager;
use serde::{Deserialize, Serialize};
/// Status of a managed MCP server connection.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum McpConnectionStatus {
Disconnected,
Connecting,
Connected,
AuthRequired,
Error,
}
impl std::fmt::Display for McpConnectionStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Disconnected => write!(f, "disconnected"),
Self::Connecting => write!(f, "connecting"),
Self::Connected => write!(f, "connected"),
Self::AuthRequired => write!(f, "auth_required"),
Self::Error => write!(f, "error"),
}
}
}
/// Metadata about an MCP resource.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpResourceInfo {
pub uri: String,
pub name: String,
pub description: Option<String>,
pub mime_type: Option<String>,
}
/// Metadata about an MCP tool exposed by a server.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpToolInfo {
pub name: String,
pub description: Option<String>,
pub input_schema: Option<serde_json::Value>,
}
/// Tracked state of an MCP server connection.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpServerState {
pub server_name: String,
pub status: McpConnectionStatus,
pub tools: Vec<McpToolInfo>,
pub resources: Vec<McpResourceInfo>,
pub server_info: Option<String>,
pub error_message: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct McpToolRegistry {
inner: Arc<Mutex<HashMap<String, McpServerState>>>,
manager: Arc<OnceLock<Arc<Mutex<McpServerManager>>>>,
}
impl McpToolRegistry {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn set_manager(
&self,
manager: Arc<Mutex<McpServerManager>>,
) -> Result<(), Arc<Mutex<McpServerManager>>> {
self.manager.set(manager)
}
pub fn register_server(
&self,
server_name: &str,
status: McpConnectionStatus,
tools: Vec<McpToolInfo>,
resources: Vec<McpResourceInfo>,
server_info: Option<String>,
) {
let mut inner = self.inner.lock().expect("mcp registry lock poisoned");
inner.insert(
server_name.to_owned(),
McpServerState {
server_name: server_name.to_owned(),
status,
tools,
resources,
server_info,
error_message: None,
},
);
}
pub fn get_server(&self, server_name: &str) -> Option<McpServerState> {
let inner = self.inner.lock().expect("mcp registry lock poisoned");
inner.get(server_name).cloned()
}
pub fn list_servers(&self) -> Vec<McpServerState> {
let inner = self.inner.lock().expect("mcp registry lock poisoned");
inner.values().cloned().collect()
}
pub fn list_resources(&self, server_name: &str) -> Result<Vec<McpResourceInfo>, String> {
let inner = self.inner.lock().expect("mcp registry lock poisoned");
match inner.get(server_name) {
Some(state) => {
if state.status != McpConnectionStatus::Connected {
return Err(format!(
"server '{}' is not connected (status: {})",
server_name, state.status
));
}
Ok(state.resources.clone())
}
None => Err(format!("server '{}' not found", server_name)),
}
}
pub fn read_resource(&self, server_name: &str, uri: &str) -> Result<McpResourceInfo, String> {
let inner = self.inner.lock().expect("mcp registry lock poisoned");
let state = inner
.get(server_name)
.ok_or_else(|| format!("server '{}' not found", server_name))?;
if state.status != McpConnectionStatus::Connected {
return Err(format!(
"server '{}' is not connected (status: {})",
server_name, state.status
));
}
state
.resources
.iter()
.find(|r| r.uri == uri)
.cloned()
.ok_or_else(|| format!("resource '{}' not found on server '{}'", uri, server_name))
}
pub fn list_tools(&self, server_name: &str) -> Result<Vec<McpToolInfo>, String> {
let inner = self.inner.lock().expect("mcp registry lock poisoned");
match inner.get(server_name) {
Some(state) => {
if state.status != McpConnectionStatus::Connected {
return Err(format!(
"server '{}' is not connected (status: {})",
server_name, state.status
));
}
Ok(state.tools.clone())
}
None => Err(format!("server '{}' not found", server_name)),
}
}
fn spawn_tool_call(
manager: Arc<Mutex<McpServerManager>>,
qualified_tool_name: String,
arguments: Option<serde_json::Value>,
) -> Result<serde_json::Value, String> {
let join_handle = std::thread::Builder::new()
.name(format!("mcp-tool-call-{qualified_tool_name}"))
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|error| format!("failed to create MCP tool runtime: {error}"))?;
runtime.block_on(async move {
let response = {
let mut manager = manager
.lock()
.map_err(|_| "mcp server manager lock poisoned".to_string())?;
manager
.discover_tools()
.await
.map_err(|error| error.to_string())?;
let response = manager
.call_tool(&qualified_tool_name, arguments)
.await
.map_err(|error| error.to_string());
let shutdown = manager.shutdown().await.map_err(|error| error.to_string());
match (response, shutdown) {
(Ok(response), Ok(())) => Ok(response),
(Err(error), Ok(())) | (Err(error), Err(_)) => Err(error),
(Ok(_), Err(error)) => Err(error),
}
}?;
if let Some(error) = response.error {
return Err(format!(
"MCP server returned JSON-RPC error for tools/call: {} ({})",
error.message, error.code
));
}
let result = response.result.ok_or_else(|| {
"MCP server returned no result for tools/call".to_string()
})?;
serde_json::to_value(result)
.map_err(|error| format!("failed to serialize MCP tool result: {error}"))
})
})
.map_err(|error| format!("failed to spawn MCP tool call thread: {error}"))?;
join_handle.join().map_err(|panic_payload| {
if let Some(message) = panic_payload.downcast_ref::<&str>() {
format!("MCP tool call thread panicked: {message}")
} else if let Some(message) = panic_payload.downcast_ref::<String>() {
format!("MCP tool call thread panicked: {message}")
} else {
"MCP tool call thread panicked".to_string()
}
})?
}
pub fn call_tool(
&self,
server_name: &str,
tool_name: &str,
arguments: &serde_json::Value,
) -> Result<serde_json::Value, String> {
let inner = self.inner.lock().expect("mcp registry lock poisoned");
let state = inner
.get(server_name)
.ok_or_else(|| format!("server '{}' not found", server_name))?;
if state.status != McpConnectionStatus::Connected {
return Err(format!(
"server '{}' is not connected (status: {})",
server_name, state.status
));
}
if !state.tools.iter().any(|t| t.name == tool_name) {
return Err(format!(
"tool '{}' not found on server '{}'",
tool_name, server_name
));
}
drop(inner);
let manager = self
.manager
.get()
.cloned()
.ok_or_else(|| "MCP server manager is not configured".to_string())?;
Self::spawn_tool_call(
manager,
mcp_tool_name(server_name, tool_name),
(!arguments.is_null()).then(|| arguments.clone()),
)
}
/// Set auth status for a server.
pub fn set_auth_status(
&self,
server_name: &str,
status: McpConnectionStatus,
) -> Result<(), String> {
let mut inner = self.inner.lock().expect("mcp registry lock poisoned");
let state = inner
.get_mut(server_name)
.ok_or_else(|| format!("server '{}' not found", server_name))?;
state.status = status;
Ok(())
}
/// Disconnect / remove a server.
pub fn disconnect(&self, server_name: &str) -> Option<McpServerState> {
let mut inner = self.inner.lock().expect("mcp registry lock poisoned");
inner.remove(server_name)
}
/// Number of registered servers.
#[must_use]
pub fn len(&self) -> usize {
let inner = self.inner.lock().expect("mcp registry lock poisoned");
inner.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::fs;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use super::*;
use crate::config::{
ConfigSource, McpServerConfig, McpStdioServerConfig, ScopedMcpServerConfig,
};
fn temp_dir() -> PathBuf {
static NEXT_TEMP_DIR_ID: AtomicU64 = AtomicU64::new(0);
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("time should be after epoch")
.as_nanos();
let unique_id = NEXT_TEMP_DIR_ID.fetch_add(1, Ordering::Relaxed);
std::env::temp_dir().join(format!("runtime-mcp-tool-bridge-{nanos}-{unique_id}"))
}
fn cleanup_script(script_path: &Path) {
if let Some(root) = script_path.parent() {
let _ = fs::remove_dir_all(root);
}
}
fn write_bridge_mcp_server_script() -> PathBuf {
let root = temp_dir();
fs::create_dir_all(&root).expect("temp dir");
let script_path = root.join("bridge-mcp-server.py");
let script = [
"#!/usr/bin/env python3",
"import json, os, sys",
"LABEL = os.environ.get('MCP_SERVER_LABEL', 'server')",
"LOG_PATH = os.environ.get('MCP_LOG_PATH')",
"",
"def log(method):",
" if LOG_PATH:",
" with open(LOG_PATH, 'a', encoding='utf-8') as handle:",
" handle.write(f'{method}\\n')",
"",
"def read_message():",
" header = b''",
r" while not header.endswith(b'\r\n\r\n'):",
" chunk = sys.stdin.buffer.read(1)",
" if not chunk:",
" return None",
" header += chunk",
" length = 0",
r" for line in header.decode().split('\r\n'):",
r" if line.lower().startswith('content-length:'):",
r" length = int(line.split(':', 1)[1].strip())",
" payload = sys.stdin.buffer.read(length)",
" return json.loads(payload.decode())",
"",
"def send_message(message):",
" payload = json.dumps(message).encode()",
r" sys.stdout.buffer.write(f'Content-Length: {len(payload)}\r\n\r\n'.encode() + payload)",
" sys.stdout.buffer.flush()",
"",
"while True:",
" request = read_message()",
" if request is None:",
" break",
" method = request['method']",
" log(method)",
" if method == 'initialize':",
" send_message({",
" 'jsonrpc': '2.0',",
" 'id': request['id'],",
" 'result': {",
" 'protocolVersion': request['params']['protocolVersion'],",
" 'capabilities': {'tools': {}},",
" 'serverInfo': {'name': LABEL, 'version': '1.0.0'}",
" }",
" })",
" elif method == 'tools/list':",
" send_message({",
" 'jsonrpc': '2.0',",
" 'id': request['id'],",
" 'result': {",
" 'tools': [",
" {",
" 'name': 'echo',",
" 'description': f'Echo tool for {LABEL}',",
" 'inputSchema': {",
" 'type': 'object',",
" 'properties': {'text': {'type': 'string'}},",
" 'required': ['text']",
" }",
" }",
" ]",
" }",
" })",
" elif method == 'tools/call':",
" args = request['params'].get('arguments') or {}",
" text = args.get('text', '')",
" send_message({",
" 'jsonrpc': '2.0',",
" 'id': request['id'],",
" 'result': {",
" 'content': [{'type': 'text', 'text': f'{LABEL}:{text}'}],",
" 'structuredContent': {'server': LABEL, 'echoed': text},",
" 'isError': False",
" }",
" })",
" else:",
" send_message({",
" 'jsonrpc': '2.0',",
" 'id': request['id'],",
" 'error': {'code': -32601, 'message': f'unknown method: {method}'},",
" })",
"",
]
.join("\n");
fs::write(&script_path, script).expect("write script");
let mut permissions = fs::metadata(&script_path).expect("metadata").permissions();
permissions.set_mode(0o755);
fs::set_permissions(&script_path, permissions).expect("chmod");
script_path
}
fn manager_server_config(
script_path: &Path,
server_name: &str,
log_path: &Path,
) -> ScopedMcpServerConfig {
ScopedMcpServerConfig {
scope: ConfigSource::Local,
config: McpServerConfig::Stdio(McpStdioServerConfig {
command: "python3".to_string(),
args: vec![script_path.to_string_lossy().into_owned()],
env: BTreeMap::from([
("MCP_SERVER_LABEL".to_string(), server_name.to_string()),
(
"MCP_LOG_PATH".to_string(),
log_path.to_string_lossy().into_owned(),
),
]),
tool_call_timeout_ms: Some(1_000),
}),
}
}
#[test]
fn registers_and_retrieves_server() {
let registry = McpToolRegistry::new();
registry.register_server(
"test-server",
McpConnectionStatus::Connected,
vec![McpToolInfo {
name: "greet".into(),
description: Some("Greet someone".into()),
input_schema: None,
}],
vec![McpResourceInfo {
uri: "res://data".into(),
name: "Data".into(),
description: None,
mime_type: Some("application/json".into()),
}],
Some("TestServer v1.0".into()),
);
let server = registry.get_server("test-server").expect("should exist");
assert_eq!(server.status, McpConnectionStatus::Connected);
assert_eq!(server.tools.len(), 1);
assert_eq!(server.resources.len(), 1);
}
#[test]
fn lists_resources_from_connected_server() {
let registry = McpToolRegistry::new();
registry.register_server(
"srv",
McpConnectionStatus::Connected,
vec![],
vec![McpResourceInfo {
uri: "res://alpha".into(),
name: "Alpha".into(),
description: None,
mime_type: None,
}],
None,
);
let resources = registry.list_resources("srv").expect("should succeed");
assert_eq!(resources.len(), 1);
assert_eq!(resources[0].uri, "res://alpha");
}
#[test]
fn rejects_resource_listing_for_disconnected_server() {
let registry = McpToolRegistry::new();
registry.register_server(
"srv",
McpConnectionStatus::Disconnected,
vec![],
vec![],
None,
);
assert!(registry.list_resources("srv").is_err());
}
#[test]
fn reads_specific_resource() {
let registry = McpToolRegistry::new();
registry.register_server(
"srv",
McpConnectionStatus::Connected,
vec![],
vec![McpResourceInfo {
uri: "res://data".into(),
name: "Data".into(),
description: Some("Test data".into()),
mime_type: Some("text/plain".into()),
}],
None,
);
let resource = registry
.read_resource("srv", "res://data")
.expect("should find");
assert_eq!(resource.name, "Data");
assert!(registry.read_resource("srv", "res://missing").is_err());
}
#[test]
fn given_connected_server_without_manager_when_calling_tool_then_it_errors() {
let registry = McpToolRegistry::new();
registry.register_server(
"srv",
McpConnectionStatus::Connected,
vec![McpToolInfo {
name: "greet".into(),
description: None,
input_schema: None,
}],
vec![],
None,
);
let error = registry
.call_tool("srv", "greet", &serde_json::json!({"name": "world"}))
.expect_err("should require a configured manager");
assert!(error.contains("MCP server manager is not configured"));
// Unknown tool should fail
assert!(registry
.call_tool("srv", "missing", &serde_json::json!({}))
.is_err());
}
#[test]
fn given_connected_server_with_manager_when_calling_tool_then_it_returns_live_result() {
let script_path = write_bridge_mcp_server_script();
let root = script_path.parent().expect("script parent");
let log_path = root.join("bridge.log");
let servers = BTreeMap::from([(
"alpha".to_string(),
manager_server_config(&script_path, "alpha", &log_path),
)]);
let manager = Arc::new(Mutex::new(McpServerManager::from_servers(&servers)));
let registry = McpToolRegistry::new();
registry.register_server(
"alpha",
McpConnectionStatus::Connected,
vec![McpToolInfo {
name: "echo".into(),
description: Some("Echo tool for alpha".into()),
input_schema: Some(serde_json::json!({
"type": "object",
"properties": {"text": {"type": "string"}},
"required": ["text"]
})),
}],
vec![],
Some("bridge test server".into()),
);
registry
.set_manager(Arc::clone(&manager))
.expect("manager should only be set once");
let result = registry
.call_tool("alpha", "echo", &serde_json::json!({"text": "hello"}))
.expect("should return live MCP result");
assert_eq!(
result["structuredContent"]["server"],
serde_json::json!("alpha")
);
assert_eq!(
result["structuredContent"]["echoed"],
serde_json::json!("hello")
);
assert_eq!(
result["content"][0]["text"],
serde_json::json!("alpha:hello")
);
let log = fs::read_to_string(&log_path).expect("read log");
assert_eq!(
log.lines().collect::<Vec<_>>(),
vec!["initialize", "tools/list", "tools/call"]
);
cleanup_script(&script_path);
}
#[test]
fn rejects_tool_call_on_disconnected_server() {
let registry = McpToolRegistry::new();
registry.register_server(
"srv",
McpConnectionStatus::AuthRequired,
vec![McpToolInfo {
name: "greet".into(),
description: None,
input_schema: None,
}],
vec![],
None,
);
assert!(registry
.call_tool("srv", "greet", &serde_json::json!({}))
.is_err());
}
#[test]
fn sets_auth_and_disconnects() {
let registry = McpToolRegistry::new();
registry.register_server(
"srv",
McpConnectionStatus::AuthRequired,
vec![],
vec![],
None,
);
registry
.set_auth_status("srv", McpConnectionStatus::Connected)
.expect("should succeed");
let state = registry.get_server("srv").unwrap();
assert_eq!(state.status, McpConnectionStatus::Connected);
let removed = registry.disconnect("srv");
assert!(removed.is_some());
assert!(registry.is_empty());
}
#[test]
fn rejects_operations_on_missing_server() {
let registry = McpToolRegistry::new();
assert!(registry.list_resources("missing").is_err());
assert!(registry.read_resource("missing", "uri").is_err());
assert!(registry.list_tools("missing").is_err());
assert!(registry
.call_tool("missing", "tool", &serde_json::json!({}))
.is_err());
assert!(registry
.set_auth_status("missing", McpConnectionStatus::Connected)
.is_err());
}
#[test]
fn mcp_connection_status_display_all_variants() {
// given
let cases = [
(McpConnectionStatus::Disconnected, "disconnected"),
(McpConnectionStatus::Connecting, "connecting"),
(McpConnectionStatus::Connected, "connected"),
(McpConnectionStatus::AuthRequired, "auth_required"),
(McpConnectionStatus::Error, "error"),
];
// when
let rendered: Vec<_> = cases
.into_iter()
.map(|(status, expected)| (status.to_string(), expected))
.collect();
// then
assert_eq!(
rendered,
vec![
("disconnected".to_string(), "disconnected"),
("connecting".to_string(), "connecting"),
("connected".to_string(), "connected"),
("auth_required".to_string(), "auth_required"),
("error".to_string(), "error"),
]
);
}
#[test]
fn list_servers_returns_all_registered() {
// given
let registry = McpToolRegistry::new();
registry.register_server(
"alpha",
McpConnectionStatus::Connected,
vec![],
vec![],
None,
);
registry.register_server(
"beta",
McpConnectionStatus::Connecting,
vec![],
vec![],
None,
);
// when
let servers = registry.list_servers();
// then
assert_eq!(servers.len(), 2);
assert!(servers.iter().any(|server| server.server_name == "alpha"));
assert!(servers.iter().any(|server| server.server_name == "beta"));
}
#[test]
fn list_tools_from_connected_server() {
// given
let registry = McpToolRegistry::new();
registry.register_server(
"srv",
McpConnectionStatus::Connected,
vec![McpToolInfo {
name: "inspect".into(),
description: Some("Inspect data".into()),
input_schema: Some(serde_json::json!({"type": "object"})),
}],
vec![],
None,
);
// when
let tools = registry.list_tools("srv").expect("tools should list");
// then
assert_eq!(tools.len(), 1);
assert_eq!(tools[0].name, "inspect");
}
#[test]
fn list_tools_rejects_disconnected_server() {
// given
let registry = McpToolRegistry::new();
registry.register_server(
"srv",
McpConnectionStatus::AuthRequired,
vec![],
vec![],
None,
);
// when
let result = registry.list_tools("srv");
// then
let error = result.expect_err("non-connected server should fail");
assert!(error.contains("not connected"));
assert!(error.contains("auth_required"));
}
#[test]
fn list_tools_rejects_missing_server() {
// given
let registry = McpToolRegistry::new();
// when
let result = registry.list_tools("missing");
// then
assert_eq!(
result.expect_err("missing server should fail"),
"server 'missing' not found"
);
}
#[test]
fn get_server_returns_none_for_missing() {
// given
let registry = McpToolRegistry::new();
// when
let server = registry.get_server("missing");
// then
assert!(server.is_none());
}
#[test]
fn call_tool_payload_structure() {
let script_path = write_bridge_mcp_server_script();
let root = script_path.parent().expect("script parent");
let log_path = root.join("payload.log");
let servers = BTreeMap::from([(
"srv".to_string(),
manager_server_config(&script_path, "srv", &log_path),
)]);
let registry = McpToolRegistry::new();
let arguments = serde_json::json!({"text": "world"});
registry.register_server(
"srv",
McpConnectionStatus::Connected,
vec![McpToolInfo {
name: "echo".into(),
description: Some("Echo tool for srv".into()),
input_schema: Some(serde_json::json!({
"type": "object",
"properties": {"text": {"type": "string"}},
"required": ["text"]
})),
}],
vec![],
None,
);
registry
.set_manager(Arc::new(Mutex::new(McpServerManager::from_servers(
&servers,
))))
.expect("manager should only be set once");
let result = registry
.call_tool("srv", "echo", &arguments)
.expect("tool should return live payload");
assert_eq!(result["structuredContent"]["server"], "srv");
assert_eq!(result["structuredContent"]["echoed"], "world");
assert_eq!(result["content"][0]["text"], "srv:world");
cleanup_script(&script_path);
}
#[test]
fn upsert_overwrites_existing_server() {
// given
let registry = McpToolRegistry::new();
registry.register_server("srv", McpConnectionStatus::Connecting, vec![], vec![], None);
// when
registry.register_server(
"srv",
McpConnectionStatus::Connected,
vec![McpToolInfo {
name: "inspect".into(),
description: None,
input_schema: None,
}],
vec![],
Some("Inspector".into()),
);
let state = registry.get_server("srv").expect("server should exist");
// then
assert_eq!(state.status, McpConnectionStatus::Connected);
assert_eq!(state.tools.len(), 1);
assert_eq!(state.server_info.as_deref(), Some("Inspector"));
}
#[test]
fn disconnect_missing_returns_none() {
// given
let registry = McpToolRegistry::new();
// when
let removed = registry.disconnect("missing");
// then
assert!(removed.is_none());
}
#[test]
fn len_and_is_empty_transitions() {
// given
let registry = McpToolRegistry::new();
// when
registry.register_server(
"alpha",
McpConnectionStatus::Connected,
vec![],
vec![],
None,
);
registry.register_server("beta", McpConnectionStatus::Connected, vec![], vec![], None);
let after_create = registry.len();
registry.disconnect("alpha");
let after_first_remove = registry.len();
registry.disconnect("beta");
// then
assert_eq!(after_create, 2);
assert_eq!(after_first_remove, 1);
assert_eq!(registry.len(), 0);
assert!(registry.is_empty());
}
}
+603
View File
@@ -0,0 +1,603 @@
use std::collections::BTreeMap;
use std::fs::{self, File};
use std::io::{self, Read};
use std::path::PathBuf;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::config::OAuthConfig;
/// Persisted OAuth access token bundle used by the CLI.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct OAuthTokenSet {
pub access_token: String,
pub refresh_token: Option<String>,
pub expires_at: Option<u64>,
pub scopes: Vec<String>,
}
/// PKCE verifier/challenge pair generated for an OAuth authorization flow.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PkceCodePair {
pub verifier: String,
pub challenge: String,
pub challenge_method: PkceChallengeMethod,
}
/// Challenge algorithms supported by the local PKCE helpers.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PkceChallengeMethod {
S256,
}
impl PkceChallengeMethod {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::S256 => "S256",
}
}
}
/// Parameters needed to build an authorization URL for browser-based login.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OAuthAuthorizationRequest {
pub authorize_url: String,
pub client_id: String,
pub redirect_uri: String,
pub scopes: Vec<String>,
pub state: String,
pub code_challenge: String,
pub code_challenge_method: PkceChallengeMethod,
pub extra_params: BTreeMap<String, String>,
}
/// Request body for exchanging an OAuth authorization code for tokens.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OAuthTokenExchangeRequest {
pub grant_type: &'static str,
pub code: String,
pub redirect_uri: String,
pub client_id: String,
pub code_verifier: String,
pub state: String,
}
/// Request body for refreshing an existing OAuth token set.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OAuthRefreshRequest {
pub grant_type: &'static str,
pub refresh_token: String,
pub client_id: String,
pub scopes: Vec<String>,
}
/// Parsed query parameters returned to the local OAuth callback endpoint.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OAuthCallbackParams {
pub code: Option<String>,
pub state: Option<String>,
pub error: Option<String>,
pub error_description: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct StoredOAuthCredentials {
access_token: String,
#[serde(default)]
refresh_token: Option<String>,
#[serde(default)]
expires_at: Option<u64>,
#[serde(default)]
scopes: Vec<String>,
}
impl From<OAuthTokenSet> for StoredOAuthCredentials {
fn from(value: OAuthTokenSet) -> Self {
Self {
access_token: value.access_token,
refresh_token: value.refresh_token,
expires_at: value.expires_at,
scopes: value.scopes,
}
}
}
impl From<StoredOAuthCredentials> for OAuthTokenSet {
fn from(value: StoredOAuthCredentials) -> Self {
Self {
access_token: value.access_token,
refresh_token: value.refresh_token,
expires_at: value.expires_at,
scopes: value.scopes,
}
}
}
impl OAuthAuthorizationRequest {
#[must_use]
pub fn from_config(
config: &OAuthConfig,
redirect_uri: impl Into<String>,
state: impl Into<String>,
pkce: &PkceCodePair,
) -> Self {
Self {
authorize_url: config.authorize_url.clone(),
client_id: config.client_id.clone(),
redirect_uri: redirect_uri.into(),
scopes: config.scopes.clone(),
state: state.into(),
code_challenge: pkce.challenge.clone(),
code_challenge_method: pkce.challenge_method,
extra_params: BTreeMap::new(),
}
}
#[must_use]
pub fn with_extra_param(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.extra_params.insert(key.into(), value.into());
self
}
#[must_use]
pub fn build_url(&self) -> String {
let mut params = vec![
("response_type", "code".to_string()),
("client_id", self.client_id.clone()),
("redirect_uri", self.redirect_uri.clone()),
("scope", self.scopes.join(" ")),
("state", self.state.clone()),
("code_challenge", self.code_challenge.clone()),
(
"code_challenge_method",
self.code_challenge_method.as_str().to_string(),
),
];
params.extend(
self.extra_params
.iter()
.map(|(key, value)| (key.as_str(), value.clone())),
);
let query = params
.into_iter()
.map(|(key, value)| format!("{}={}", percent_encode(key), percent_encode(&value)))
.collect::<Vec<_>>()
.join("&");
format!(
"{}{}{}",
self.authorize_url,
if self.authorize_url.contains('?') {
'&'
} else {
'?'
},
query
)
}
}
impl OAuthTokenExchangeRequest {
#[must_use]
pub fn from_config(
config: &OAuthConfig,
code: impl Into<String>,
state: impl Into<String>,
verifier: impl Into<String>,
redirect_uri: impl Into<String>,
) -> Self {
Self {
grant_type: "authorization_code",
code: code.into(),
redirect_uri: redirect_uri.into(),
client_id: config.client_id.clone(),
code_verifier: verifier.into(),
state: state.into(),
}
}
#[must_use]
pub fn form_params(&self) -> BTreeMap<&str, String> {
BTreeMap::from([
("grant_type", self.grant_type.to_string()),
("code", self.code.clone()),
("redirect_uri", self.redirect_uri.clone()),
("client_id", self.client_id.clone()),
("code_verifier", self.code_verifier.clone()),
("state", self.state.clone()),
])
}
}
impl OAuthRefreshRequest {
#[must_use]
pub fn from_config(
config: &OAuthConfig,
refresh_token: impl Into<String>,
scopes: Option<Vec<String>>,
) -> Self {
Self {
grant_type: "refresh_token",
refresh_token: refresh_token.into(),
client_id: config.client_id.clone(),
scopes: scopes.unwrap_or_else(|| config.scopes.clone()),
}
}
#[must_use]
pub fn form_params(&self) -> BTreeMap<&str, String> {
BTreeMap::from([
("grant_type", self.grant_type.to_string()),
("refresh_token", self.refresh_token.clone()),
("client_id", self.client_id.clone()),
("scope", self.scopes.join(" ")),
])
}
}
pub fn generate_pkce_pair() -> io::Result<PkceCodePair> {
let verifier = generate_random_token(32)?;
Ok(PkceCodePair {
challenge: code_challenge_s256(&verifier),
verifier,
challenge_method: PkceChallengeMethod::S256,
})
}
pub fn generate_state() -> io::Result<String> {
generate_random_token(32)
}
#[must_use]
pub fn code_challenge_s256(verifier: &str) -> String {
let digest = Sha256::digest(verifier.as_bytes());
base64url_encode(&digest)
}
#[must_use]
pub fn loopback_redirect_uri(port: u16) -> String {
format!("http://localhost:{port}/callback")
}
pub fn credentials_path() -> io::Result<PathBuf> {
Ok(credentials_home_dir()?.join("credentials.json"))
}
pub fn load_oauth_credentials() -> io::Result<Option<OAuthTokenSet>> {
let path = credentials_path()?;
let root = read_credentials_root(&path)?;
let Some(oauth) = root.get("oauth") else {
return Ok(None);
};
if oauth.is_null() {
return Ok(None);
}
let stored = serde_json::from_value::<StoredOAuthCredentials>(oauth.clone())
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
Ok(Some(stored.into()))
}
pub fn save_oauth_credentials(token_set: &OAuthTokenSet) -> io::Result<()> {
let path = credentials_path()?;
let mut root = read_credentials_root(&path)?;
root.insert(
"oauth".to_string(),
serde_json::to_value(StoredOAuthCredentials::from(token_set.clone()))
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?,
);
write_credentials_root(&path, &root)
}
pub fn clear_oauth_credentials() -> io::Result<()> {
let path = credentials_path()?;
let mut root = read_credentials_root(&path)?;
root.remove("oauth");
write_credentials_root(&path, &root)
}
pub fn parse_oauth_callback_request_target(target: &str) -> Result<OAuthCallbackParams, String> {
let (path, query) = target
.split_once('?')
.map_or((target, ""), |(path, query)| (path, query));
if path != "/callback" {
return Err(format!("unexpected callback path: {path}"));
}
parse_oauth_callback_query(query)
}
pub fn parse_oauth_callback_query(query: &str) -> Result<OAuthCallbackParams, String> {
let mut params = BTreeMap::new();
for pair in query.split('&').filter(|pair| !pair.is_empty()) {
let (key, value) = pair
.split_once('=')
.map_or((pair, ""), |(key, value)| (key, value));
params.insert(percent_decode(key)?, percent_decode(value)?);
}
Ok(OAuthCallbackParams {
code: params.get("code").cloned(),
state: params.get("state").cloned(),
error: params.get("error").cloned(),
error_description: params.get("error_description").cloned(),
})
}
fn generate_random_token(bytes: usize) -> io::Result<String> {
let mut buffer = vec![0_u8; bytes];
File::open("/dev/urandom")?.read_exact(&mut buffer)?;
Ok(base64url_encode(&buffer))
}
fn credentials_home_dir() -> io::Result<PathBuf> {
if let Some(path) = std::env::var_os("CLAW_CONFIG_HOME") {
return Ok(PathBuf::from(path));
}
let home = std::env::var_os("HOME")
.or_else(|| std::env::var_os("USERPROFILE"))
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::NotFound,
"HOME is not set (on Windows, set USERPROFILE or HOME, \
or use CLAW_CONFIG_HOME to point directly at the config directory)",
)
})?;
Ok(PathBuf::from(home).join(".claw"))
}
fn read_credentials_root(path: &PathBuf) -> io::Result<Map<String, Value>> {
match fs::read_to_string(path) {
Ok(contents) => {
if contents.trim().is_empty() {
return Ok(Map::new());
}
serde_json::from_str::<Value>(&contents)
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?
.as_object()
.cloned()
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidData,
"credentials file must contain a JSON object",
)
})
}
Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(Map::new()),
Err(error) => Err(error),
}
}
fn write_credentials_root(path: &PathBuf, root: &Map<String, Value>) -> io::Result<()> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
let rendered = serde_json::to_string_pretty(&Value::Object(root.clone()))
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
let temp_path = path.with_extension("json.tmp");
fs::write(&temp_path, format!("{rendered}\n"))?;
fs::rename(temp_path, path)
}
fn base64url_encode(bytes: &[u8]) -> String {
const TABLE: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
let mut output = String::new();
let mut index = 0;
while index + 3 <= bytes.len() {
let block = (u32::from(bytes[index]) << 16)
| (u32::from(bytes[index + 1]) << 8)
| u32::from(bytes[index + 2]);
output.push(TABLE[((block >> 18) & 0x3F) as usize] as char);
output.push(TABLE[((block >> 12) & 0x3F) as usize] as char);
output.push(TABLE[((block >> 6) & 0x3F) as usize] as char);
output.push(TABLE[(block & 0x3F) as usize] as char);
index += 3;
}
match bytes.len().saturating_sub(index) {
1 => {
let block = u32::from(bytes[index]) << 16;
output.push(TABLE[((block >> 18) & 0x3F) as usize] as char);
output.push(TABLE[((block >> 12) & 0x3F) as usize] as char);
}
2 => {
let block = (u32::from(bytes[index]) << 16) | (u32::from(bytes[index + 1]) << 8);
output.push(TABLE[((block >> 18) & 0x3F) as usize] as char);
output.push(TABLE[((block >> 12) & 0x3F) as usize] as char);
output.push(TABLE[((block >> 6) & 0x3F) as usize] as char);
}
_ => {}
}
output
}
fn percent_encode(value: &str) -> String {
let mut encoded = String::new();
for byte in value.bytes() {
match byte {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
encoded.push(char::from(byte));
}
_ => {
use std::fmt::Write as _;
let _ = write!(&mut encoded, "%{byte:02X}");
}
}
}
encoded
}
fn percent_decode(value: &str) -> Result<String, String> {
let mut decoded = Vec::with_capacity(value.len());
let bytes = value.as_bytes();
let mut index = 0;
while index < bytes.len() {
match bytes[index] {
b'%' if index + 2 < bytes.len() => {
let hi = decode_hex(bytes[index + 1])?;
let lo = decode_hex(bytes[index + 2])?;
decoded.push((hi << 4) | lo);
index += 3;
}
b'+' => {
decoded.push(b' ');
index += 1;
}
byte => {
decoded.push(byte);
index += 1;
}
}
}
String::from_utf8(decoded).map_err(|error| error.to_string())
}
fn decode_hex(byte: u8) -> Result<u8, String> {
match byte {
b'0'..=b'9' => Ok(byte - b'0'),
b'a'..=b'f' => Ok(byte - b'a' + 10),
b'A'..=b'F' => Ok(byte - b'A' + 10),
_ => Err(format!("invalid percent byte: {byte}")),
}
}
#[cfg(test)]
mod tests {
use std::time::{SystemTime, UNIX_EPOCH};
use super::{
clear_oauth_credentials, code_challenge_s256, credentials_path, generate_pkce_pair,
generate_state, load_oauth_credentials, loopback_redirect_uri, parse_oauth_callback_query,
parse_oauth_callback_request_target, save_oauth_credentials, OAuthAuthorizationRequest,
OAuthConfig, OAuthRefreshRequest, OAuthTokenExchangeRequest, OAuthTokenSet,
};
fn sample_config() -> OAuthConfig {
OAuthConfig {
client_id: "runtime-client".to_string(),
authorize_url: "https://console.test/oauth/authorize".to_string(),
token_url: "https://console.test/oauth/token".to_string(),
callback_port: Some(4545),
manual_redirect_url: Some("https://console.test/oauth/callback".to_string()),
scopes: vec!["org:read".to_string(), "user:write".to_string()],
}
}
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
crate::test_env_lock()
}
fn temp_config_home() -> std::path::PathBuf {
std::env::temp_dir().join(format!(
"runtime-oauth-test-{}-{}",
std::process::id(),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("time")
.as_nanos()
))
}
#[test]
fn s256_challenge_matches_expected_vector() {
assert_eq!(
code_challenge_s256("dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"),
"E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"
);
}
#[test]
fn generates_pkce_pair_and_state() {
let pair = generate_pkce_pair().expect("pkce pair");
let state = generate_state().expect("state");
assert!(!pair.verifier.is_empty());
assert!(!pair.challenge.is_empty());
assert!(!state.is_empty());
}
#[test]
fn builds_authorize_url_and_form_requests() {
let config = sample_config();
let pair = generate_pkce_pair().expect("pkce");
let url = OAuthAuthorizationRequest::from_config(
&config,
loopback_redirect_uri(4545),
"state-123",
&pair,
)
.with_extra_param("login_hint", "[email protected]")
.build_url();
assert!(url.starts_with("https://console.test/oauth/authorize?"));
assert!(url.contains("response_type=code"));
assert!(url.contains("client_id=runtime-client"));
assert!(url.contains("scope=org%3Aread%20user%3Awrite"));
assert!(url.contains("login_hint=user%40example.com"));
let exchange = OAuthTokenExchangeRequest::from_config(
&config,
"auth-code",
"state-123",
pair.verifier,
loopback_redirect_uri(4545),
);
assert_eq!(
exchange.form_params().get("grant_type").map(String::as_str),
Some("authorization_code")
);
let refresh = OAuthRefreshRequest::from_config(&config, "refresh-token", None);
assert_eq!(
refresh.form_params().get("scope").map(String::as_str),
Some("org:read user:write")
);
}
#[test]
fn oauth_credentials_round_trip_and_clear_preserves_other_fields() {
let _guard = env_lock();
let config_home = temp_config_home();
std::env::set_var("CLAW_CONFIG_HOME", &config_home);
let path = credentials_path().expect("credentials path");
std::fs::create_dir_all(path.parent().expect("parent")).expect("create parent");
std::fs::write(&path, "{\"other\":\"value\"}\n").expect("seed credentials");
let token_set = OAuthTokenSet {
access_token: "access-token".to_string(),
refresh_token: Some("refresh-token".to_string()),
expires_at: Some(123),
scopes: vec!["scope:a".to_string()],
};
save_oauth_credentials(&token_set).expect("save credentials");
assert_eq!(
load_oauth_credentials().expect("load credentials"),
Some(token_set)
);
let saved = std::fs::read_to_string(&path).expect("read saved file");
assert!(saved.contains("\"other\": \"value\""));
assert!(saved.contains("\"oauth\""));
clear_oauth_credentials().expect("clear credentials");
assert_eq!(load_oauth_credentials().expect("load cleared"), None);
let cleared = std::fs::read_to_string(&path).expect("read cleared file");
assert!(cleared.contains("\"other\": \"value\""));
assert!(!cleared.contains("\"oauth\""));
std::env::remove_var("CLAW_CONFIG_HOME");
std::fs::remove_dir_all(config_home).expect("cleanup temp dir");
}
#[test]
fn parses_callback_query_and_target() {
let params =
parse_oauth_callback_query("code=abc123&state=state-1&error_description=needs%20login")
.expect("parse query");
assert_eq!(params.code.as_deref(), Some("abc123"));
assert_eq!(params.state.as_deref(), Some("state-1"));
assert_eq!(params.error_description.as_deref(), Some("needs login"));
let params = parse_oauth_callback_request_target("/callback?code=abc&state=xyz")
.expect("parse callback target");
assert_eq!(params.code.as_deref(), Some("abc"));
assert_eq!(params.state.as_deref(), Some("xyz"));
assert!(parse_oauth_callback_request_target("/wrong?code=abc").is_err());
}
}
@@ -0,0 +1,551 @@
#![allow(
clippy::match_wildcard_for_single_variants,
clippy::must_use_candidate,
clippy::uninlined_format_args
)]
//! Permission enforcement layer that gates tool execution based on the
//! active `PermissionPolicy`.
use crate::permissions::{PermissionMode, PermissionOutcome, PermissionPolicy};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "outcome")]
pub enum EnforcementResult {
/// Tool execution is allowed.
Allowed,
/// Tool execution was denied due to insufficient permissions.
Denied {
tool: String,
active_mode: String,
required_mode: String,
reason: String,
},
}
#[derive(Debug, Clone, PartialEq)]
pub struct PermissionEnforcer {
policy: PermissionPolicy,
}
impl PermissionEnforcer {
#[must_use]
pub fn new(policy: PermissionPolicy) -> Self {
Self { policy }
}
/// Check whether a tool can be executed under the current permission policy.
/// Auto-denies when prompting is required but no prompter is provided.
pub fn check(&self, tool_name: &str, input: &str) -> EnforcementResult {
// When the active mode is Prompt, defer to the caller's interactive
// prompt flow rather than hard-denying (the enforcer has no prompter).
if self.policy.active_mode() == PermissionMode::Prompt {
return EnforcementResult::Allowed;
}
let outcome = self.policy.authorize(tool_name, input, None);
match outcome {
PermissionOutcome::Allow => EnforcementResult::Allowed,
PermissionOutcome::Deny { reason } => {
let active_mode = self.policy.active_mode();
let required_mode = self.policy.required_mode_for(tool_name);
EnforcementResult::Denied {
tool: tool_name.to_owned(),
active_mode: active_mode.as_str().to_owned(),
required_mode: required_mode.as_str().to_owned(),
reason,
}
}
}
}
#[must_use]
pub fn is_allowed(&self, tool_name: &str, input: &str) -> bool {
matches!(self.check(tool_name, input), EnforcementResult::Allowed)
}
#[must_use]
pub fn active_mode(&self) -> PermissionMode {
self.policy.active_mode()
}
/// Classify a file operation against workspace boundaries.
pub fn check_file_write(&self, path: &str, workspace_root: &str) -> EnforcementResult {
let mode = self.policy.active_mode();
match mode {
PermissionMode::ReadOnly => EnforcementResult::Denied {
tool: "write_file".to_owned(),
active_mode: mode.as_str().to_owned(),
required_mode: PermissionMode::WorkspaceWrite.as_str().to_owned(),
reason: format!("file writes are not allowed in '{}' mode", mode.as_str()),
},
PermissionMode::WorkspaceWrite => {
if is_within_workspace(path, workspace_root) {
EnforcementResult::Allowed
} else {
EnforcementResult::Denied {
tool: "write_file".to_owned(),
active_mode: mode.as_str().to_owned(),
required_mode: PermissionMode::DangerFullAccess.as_str().to_owned(),
reason: format!(
"path '{}' is outside workspace root '{}'",
path, workspace_root
),
}
}
}
// Allow and DangerFullAccess permit all writes
PermissionMode::Allow | PermissionMode::DangerFullAccess => EnforcementResult::Allowed,
PermissionMode::Prompt => EnforcementResult::Denied {
tool: "write_file".to_owned(),
active_mode: mode.as_str().to_owned(),
required_mode: PermissionMode::WorkspaceWrite.as_str().to_owned(),
reason: "file write requires confirmation in prompt mode".to_owned(),
},
}
}
/// Check if a bash command should be allowed based on current mode.
pub fn check_bash(&self, command: &str) -> EnforcementResult {
let mode = self.policy.active_mode();
match mode {
PermissionMode::ReadOnly => {
if is_read_only_command(command) {
EnforcementResult::Allowed
} else {
EnforcementResult::Denied {
tool: "bash".to_owned(),
active_mode: mode.as_str().to_owned(),
required_mode: PermissionMode::WorkspaceWrite.as_str().to_owned(),
reason: format!(
"command may modify state; not allowed in '{}' mode",
mode.as_str()
),
}
}
}
PermissionMode::Prompt => EnforcementResult::Denied {
tool: "bash".to_owned(),
active_mode: mode.as_str().to_owned(),
required_mode: PermissionMode::DangerFullAccess.as_str().to_owned(),
reason: "bash requires confirmation in prompt mode".to_owned(),
},
// WorkspaceWrite, Allow, DangerFullAccess: permit bash
_ => EnforcementResult::Allowed,
}
}
}
/// Simple workspace boundary check via string prefix.
fn is_within_workspace(path: &str, workspace_root: &str) -> bool {
let normalized = if path.starts_with('/') {
path.to_owned()
} else {
format!("{workspace_root}/{path}")
};
let root = if workspace_root.ends_with('/') {
workspace_root.to_owned()
} else {
format!("{workspace_root}/")
};
normalized.starts_with(&root) || normalized == workspace_root.trim_end_matches('/')
}
/// Conservative heuristic: is this bash command read-only?
fn is_read_only_command(command: &str) -> bool {
let first_token = command
.split_whitespace()
.next()
.unwrap_or("")
.rsplit('/')
.next()
.unwrap_or("");
matches!(
first_token,
"cat"
| "head"
| "tail"
| "less"
| "more"
| "wc"
| "ls"
| "find"
| "grep"
| "rg"
| "awk"
| "sed"
| "echo"
| "printf"
| "which"
| "where"
| "whoami"
| "pwd"
| "env"
| "printenv"
| "date"
| "cal"
| "df"
| "du"
| "free"
| "uptime"
| "uname"
| "file"
| "stat"
| "diff"
| "sort"
| "uniq"
| "tr"
| "cut"
| "paste"
| "tee"
| "xargs"
| "test"
| "true"
| "false"
| "type"
| "readlink"
| "realpath"
| "basename"
| "dirname"
| "sha256sum"
| "md5sum"
| "b3sum"
| "xxd"
| "hexdump"
| "od"
| "strings"
| "tree"
| "jq"
| "yq"
| "python3"
| "python"
| "node"
| "ruby"
| "cargo"
| "rustc"
| "git"
| "gh"
) && !command.contains("-i ")
&& !command.contains("--in-place")
&& !command.contains(" > ")
&& !command.contains(" >> ")
}
#[cfg(test)]
mod tests {
use super::*;
fn make_enforcer(mode: PermissionMode) -> PermissionEnforcer {
let policy = PermissionPolicy::new(mode);
PermissionEnforcer::new(policy)
}
#[test]
fn allow_mode_permits_everything() {
let enforcer = make_enforcer(PermissionMode::Allow);
assert!(enforcer.is_allowed("bash", ""));
assert!(enforcer.is_allowed("write_file", ""));
assert!(enforcer.is_allowed("edit_file", ""));
assert_eq!(
enforcer.check_file_write("/outside/path", "/workspace"),
EnforcementResult::Allowed
);
assert_eq!(enforcer.check_bash("rm -rf /"), EnforcementResult::Allowed);
}
#[test]
fn read_only_denies_writes() {
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
.with_tool_requirement("read_file", PermissionMode::ReadOnly)
.with_tool_requirement("grep_search", PermissionMode::ReadOnly)
.with_tool_requirement("write_file", PermissionMode::WorkspaceWrite);
let enforcer = PermissionEnforcer::new(policy);
assert!(enforcer.is_allowed("read_file", ""));
assert!(enforcer.is_allowed("grep_search", ""));
// write_file requires WorkspaceWrite but we're in ReadOnly
let result = enforcer.check("write_file", "");
assert!(matches!(result, EnforcementResult::Denied { .. }));
let result = enforcer.check_file_write("/workspace/file.rs", "/workspace");
assert!(matches!(result, EnforcementResult::Denied { .. }));
}
#[test]
fn read_only_allows_read_commands() {
let enforcer = make_enforcer(PermissionMode::ReadOnly);
assert_eq!(
enforcer.check_bash("cat src/main.rs"),
EnforcementResult::Allowed
);
assert_eq!(
enforcer.check_bash("grep -r 'pattern' ."),
EnforcementResult::Allowed
);
assert_eq!(enforcer.check_bash("ls -la"), EnforcementResult::Allowed);
}
#[test]
fn read_only_denies_write_commands() {
let enforcer = make_enforcer(PermissionMode::ReadOnly);
let result = enforcer.check_bash("rm file.txt");
assert!(matches!(result, EnforcementResult::Denied { .. }));
}
#[test]
fn workspace_write_allows_within_workspace() {
let enforcer = make_enforcer(PermissionMode::WorkspaceWrite);
let result = enforcer.check_file_write("/workspace/src/main.rs", "/workspace");
assert_eq!(result, EnforcementResult::Allowed);
}
#[test]
fn workspace_write_denies_outside_workspace() {
let enforcer = make_enforcer(PermissionMode::WorkspaceWrite);
let result = enforcer.check_file_write("/etc/passwd", "/workspace");
assert!(matches!(result, EnforcementResult::Denied { .. }));
}
#[test]
fn prompt_mode_denies_without_prompter() {
let enforcer = make_enforcer(PermissionMode::Prompt);
let result = enforcer.check_bash("echo test");
assert!(matches!(result, EnforcementResult::Denied { .. }));
let result = enforcer.check_file_write("/workspace/file.rs", "/workspace");
assert!(matches!(result, EnforcementResult::Denied { .. }));
}
#[test]
fn workspace_boundary_check() {
assert!(is_within_workspace("/workspace/src/main.rs", "/workspace"));
assert!(is_within_workspace("/workspace", "/workspace"));
assert!(!is_within_workspace("/etc/passwd", "/workspace"));
assert!(!is_within_workspace("/workspacex/hack", "/workspace"));
}
#[test]
fn read_only_command_heuristic() {
assert!(is_read_only_command("cat file.txt"));
assert!(is_read_only_command("grep pattern file"));
assert!(is_read_only_command("git log --oneline"));
assert!(!is_read_only_command("rm file.txt"));
assert!(!is_read_only_command("echo test > file.txt"));
assert!(!is_read_only_command("sed -i 's/a/b/' file"));
}
#[test]
fn active_mode_returns_policy_mode() {
// given
let modes = [
PermissionMode::ReadOnly,
PermissionMode::WorkspaceWrite,
PermissionMode::DangerFullAccess,
PermissionMode::Prompt,
PermissionMode::Allow,
];
// when
let active_modes: Vec<_> = modes
.into_iter()
.map(|mode| make_enforcer(mode).active_mode())
.collect();
// then
assert_eq!(active_modes, modes);
}
#[test]
fn danger_full_access_permits_file_writes_and_bash() {
// given
let enforcer = make_enforcer(PermissionMode::DangerFullAccess);
// when
let file_result = enforcer.check_file_write("/outside/workspace/file.txt", "/workspace");
let bash_result = enforcer.check_bash("rm -rf /tmp/scratch");
// then
assert_eq!(file_result, EnforcementResult::Allowed);
assert_eq!(bash_result, EnforcementResult::Allowed);
}
#[test]
fn check_denied_payload_contains_tool_and_modes() {
// given
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
.with_tool_requirement("write_file", PermissionMode::WorkspaceWrite);
let enforcer = PermissionEnforcer::new(policy);
// when
let result = enforcer.check("write_file", "{}");
// then
match result {
EnforcementResult::Denied {
tool,
active_mode,
required_mode,
reason,
} => {
assert_eq!(tool, "write_file");
assert_eq!(active_mode, "read-only");
assert_eq!(required_mode, "workspace-write");
assert!(reason.contains("requires workspace-write permission"));
}
other => panic!("expected denied result, got {other:?}"),
}
}
#[test]
fn workspace_write_relative_path_resolved() {
// given
let enforcer = make_enforcer(PermissionMode::WorkspaceWrite);
// when
let result = enforcer.check_file_write("src/main.rs", "/workspace");
// then
assert_eq!(result, EnforcementResult::Allowed);
}
#[test]
fn workspace_root_with_trailing_slash() {
// given
let enforcer = make_enforcer(PermissionMode::WorkspaceWrite);
// when
let result = enforcer.check_file_write("/workspace/src/main.rs", "/workspace/");
// then
assert_eq!(result, EnforcementResult::Allowed);
}
#[test]
fn workspace_root_equality() {
// given
let root = "/workspace/";
// when
let equal_to_root = is_within_workspace("/workspace", root);
// then
assert!(equal_to_root);
}
#[test]
fn bash_heuristic_full_path_prefix() {
// given
let full_path_command = "/usr/bin/cat Cargo.toml";
let git_path_command = "/usr/local/bin/git status";
// when
let cat_result = is_read_only_command(full_path_command);
let git_result = is_read_only_command(git_path_command);
// then
assert!(cat_result);
assert!(git_result);
}
#[test]
fn bash_heuristic_redirects_block_read_only_commands() {
// given
let overwrite = "cat Cargo.toml > out.txt";
let append = "echo test >> out.txt";
// when
let overwrite_result = is_read_only_command(overwrite);
let append_result = is_read_only_command(append);
// then
assert!(!overwrite_result);
assert!(!append_result);
}
#[test]
fn bash_heuristic_in_place_flag_blocks() {
// given
let interactive_python = "python -i script.py";
let in_place_sed = "sed --in-place 's/a/b/' file.txt";
// when
let interactive_result = is_read_only_command(interactive_python);
let in_place_result = is_read_only_command(in_place_sed);
// then
assert!(!interactive_result);
assert!(!in_place_result);
}
#[test]
fn bash_heuristic_empty_command() {
// given
let empty = "";
let whitespace = " ";
// when
let empty_result = is_read_only_command(empty);
let whitespace_result = is_read_only_command(whitespace);
// then
assert!(!empty_result);
assert!(!whitespace_result);
}
#[test]
fn prompt_mode_check_bash_denied_payload_fields() {
// given
let enforcer = make_enforcer(PermissionMode::Prompt);
// when
let result = enforcer.check_bash("git status");
// then
match result {
EnforcementResult::Denied {
tool,
active_mode,
required_mode,
reason,
} => {
assert_eq!(tool, "bash");
assert_eq!(active_mode, "prompt");
assert_eq!(required_mode, "danger-full-access");
assert_eq!(reason, "bash requires confirmation in prompt mode");
}
other => panic!("expected denied result, got {other:?}"),
}
}
#[test]
fn read_only_check_file_write_denied_payload() {
// given
let enforcer = make_enforcer(PermissionMode::ReadOnly);
// when
let result = enforcer.check_file_write("/workspace/file.txt", "/workspace");
// then
match result {
EnforcementResult::Denied {
tool,
active_mode,
required_mode,
reason,
} => {
assert_eq!(tool, "write_file");
assert_eq!(active_mode, "read-only");
assert_eq!(required_mode, "workspace-write");
assert!(reason.contains("file writes are not allowed"));
}
other => panic!("expected denied result, got {other:?}"),
}
}
}
@@ -0,0 +1,683 @@
use std::collections::BTreeMap;
use serde_json::Value;
use crate::config::RuntimePermissionRuleConfig;
/// Permission level assigned to a tool invocation or runtime session.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum PermissionMode {
ReadOnly,
WorkspaceWrite,
DangerFullAccess,
Prompt,
Allow,
}
impl PermissionMode {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::ReadOnly => "read-only",
Self::WorkspaceWrite => "workspace-write",
Self::DangerFullAccess => "danger-full-access",
Self::Prompt => "prompt",
Self::Allow => "allow",
}
}
}
/// Hook-provided override applied before standard permission evaluation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PermissionOverride {
Allow,
Deny,
Ask,
}
/// Additional permission context supplied by hooks or higher-level orchestration.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct PermissionContext {
override_decision: Option<PermissionOverride>,
override_reason: Option<String>,
}
impl PermissionContext {
#[must_use]
pub fn new(
override_decision: Option<PermissionOverride>,
override_reason: Option<String>,
) -> Self {
Self {
override_decision,
override_reason,
}
}
#[must_use]
pub fn override_decision(&self) -> Option<PermissionOverride> {
self.override_decision
}
#[must_use]
pub fn override_reason(&self) -> Option<&str> {
self.override_reason.as_deref()
}
}
/// Full authorization request presented to a permission prompt.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PermissionRequest {
pub tool_name: String,
pub input: String,
pub current_mode: PermissionMode,
pub required_mode: PermissionMode,
pub reason: Option<String>,
}
/// User-facing decision returned by a [`PermissionPrompter`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PermissionPromptDecision {
Allow,
Deny { reason: String },
}
/// Prompting interface used when policy requires interactive approval.
pub trait PermissionPrompter {
fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision;
}
/// Final authorization result after evaluating static rules and prompts.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PermissionOutcome {
Allow,
Deny { reason: String },
}
/// Evaluates permission mode requirements plus allow/deny/ask rules.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PermissionPolicy {
active_mode: PermissionMode,
tool_requirements: BTreeMap<String, PermissionMode>,
allow_rules: Vec<PermissionRule>,
deny_rules: Vec<PermissionRule>,
ask_rules: Vec<PermissionRule>,
}
impl PermissionPolicy {
#[must_use]
pub fn new(active_mode: PermissionMode) -> Self {
Self {
active_mode,
tool_requirements: BTreeMap::new(),
allow_rules: Vec::new(),
deny_rules: Vec::new(),
ask_rules: Vec::new(),
}
}
#[must_use]
pub fn with_tool_requirement(
mut self,
tool_name: impl Into<String>,
required_mode: PermissionMode,
) -> Self {
self.tool_requirements
.insert(tool_name.into(), required_mode);
self
}
#[must_use]
pub fn with_permission_rules(mut self, config: &RuntimePermissionRuleConfig) -> Self {
self.allow_rules = config
.allow()
.iter()
.map(|rule| PermissionRule::parse(rule))
.collect();
self.deny_rules = config
.deny()
.iter()
.map(|rule| PermissionRule::parse(rule))
.collect();
self.ask_rules = config
.ask()
.iter()
.map(|rule| PermissionRule::parse(rule))
.collect();
self
}
#[must_use]
pub fn active_mode(&self) -> PermissionMode {
self.active_mode
}
#[must_use]
pub fn required_mode_for(&self, tool_name: &str) -> PermissionMode {
self.tool_requirements
.get(tool_name)
.copied()
.unwrap_or(PermissionMode::DangerFullAccess)
}
#[must_use]
pub fn authorize(
&self,
tool_name: &str,
input: &str,
prompter: Option<&mut dyn PermissionPrompter>,
) -> PermissionOutcome {
self.authorize_with_context(tool_name, input, &PermissionContext::default(), prompter)
}
#[must_use]
#[allow(clippy::too_many_lines)]
pub fn authorize_with_context(
&self,
tool_name: &str,
input: &str,
context: &PermissionContext,
prompter: Option<&mut dyn PermissionPrompter>,
) -> PermissionOutcome {
if let Some(rule) = Self::find_matching_rule(&self.deny_rules, tool_name, input) {
return PermissionOutcome::Deny {
reason: format!(
"Permission to use {tool_name} has been denied by rule '{}'",
rule.raw
),
};
}
let current_mode = self.active_mode();
let required_mode = self.required_mode_for(tool_name);
let ask_rule = Self::find_matching_rule(&self.ask_rules, tool_name, input);
let allow_rule = Self::find_matching_rule(&self.allow_rules, tool_name, input);
match context.override_decision() {
Some(PermissionOverride::Deny) => {
return PermissionOutcome::Deny {
reason: context.override_reason().map_or_else(
|| format!("tool '{tool_name}' denied by hook"),
ToOwned::to_owned,
),
};
}
Some(PermissionOverride::Ask) => {
let reason = context.override_reason().map_or_else(
|| format!("tool '{tool_name}' requires approval due to hook guidance"),
ToOwned::to_owned,
);
return Self::prompt_or_deny(
tool_name,
input,
current_mode,
required_mode,
Some(reason),
prompter,
);
}
Some(PermissionOverride::Allow) => {
if let Some(rule) = ask_rule {
let reason = format!(
"tool '{tool_name}' requires approval due to ask rule '{}'",
rule.raw
);
return Self::prompt_or_deny(
tool_name,
input,
current_mode,
required_mode,
Some(reason),
prompter,
);
}
if allow_rule.is_some()
|| current_mode == PermissionMode::Allow
|| current_mode >= required_mode
{
return PermissionOutcome::Allow;
}
}
None => {}
}
if let Some(rule) = ask_rule {
let reason = format!(
"tool '{tool_name}' requires approval due to ask rule '{}'",
rule.raw
);
return Self::prompt_or_deny(
tool_name,
input,
current_mode,
required_mode,
Some(reason),
prompter,
);
}
if allow_rule.is_some()
|| current_mode == PermissionMode::Allow
|| current_mode >= required_mode
{
return PermissionOutcome::Allow;
}
if current_mode == PermissionMode::Prompt
|| (current_mode == PermissionMode::WorkspaceWrite
&& required_mode == PermissionMode::DangerFullAccess)
{
let reason = Some(format!(
"tool '{tool_name}' requires approval to escalate from {} to {}",
current_mode.as_str(),
required_mode.as_str()
));
return Self::prompt_or_deny(
tool_name,
input,
current_mode,
required_mode,
reason,
prompter,
);
}
PermissionOutcome::Deny {
reason: format!(
"tool '{tool_name}' requires {} permission; current mode is {}",
required_mode.as_str(),
current_mode.as_str()
),
}
}
fn prompt_or_deny(
tool_name: &str,
input: &str,
current_mode: PermissionMode,
required_mode: PermissionMode,
reason: Option<String>,
mut prompter: Option<&mut dyn PermissionPrompter>,
) -> PermissionOutcome {
let request = PermissionRequest {
tool_name: tool_name.to_string(),
input: input.to_string(),
current_mode,
required_mode,
reason: reason.clone(),
};
match prompter.as_mut() {
Some(prompter) => match prompter.decide(&request) {
PermissionPromptDecision::Allow => PermissionOutcome::Allow,
PermissionPromptDecision::Deny { reason } => PermissionOutcome::Deny { reason },
},
None => PermissionOutcome::Deny {
reason: reason.unwrap_or_else(|| {
format!(
"tool '{tool_name}' requires approval to run while mode is {}",
current_mode.as_str()
)
}),
},
}
}
fn find_matching_rule<'a>(
rules: &'a [PermissionRule],
tool_name: &str,
input: &str,
) -> Option<&'a PermissionRule> {
rules.iter().find(|rule| rule.matches(tool_name, input))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct PermissionRule {
raw: String,
tool_name: String,
matcher: PermissionRuleMatcher,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum PermissionRuleMatcher {
Any,
Exact(String),
Prefix(String),
}
impl PermissionRule {
fn parse(raw: &str) -> Self {
let trimmed = raw.trim();
let open = find_first_unescaped(trimmed, '(');
let close = find_last_unescaped(trimmed, ')');
if let (Some(open), Some(close)) = (open, close) {
if close == trimmed.len() - 1 && open < close {
let tool_name = trimmed[..open].trim();
let content = &trimmed[open + 1..close];
if !tool_name.is_empty() {
let matcher = parse_rule_matcher(content);
return Self {
raw: trimmed.to_string(),
tool_name: tool_name.to_string(),
matcher,
};
}
}
}
Self {
raw: trimmed.to_string(),
tool_name: trimmed.to_string(),
matcher: PermissionRuleMatcher::Any,
}
}
fn matches(&self, tool_name: &str, input: &str) -> bool {
if self.tool_name != tool_name {
return false;
}
match &self.matcher {
PermissionRuleMatcher::Any => true,
PermissionRuleMatcher::Exact(expected) => {
extract_permission_subject(input).is_some_and(|candidate| candidate == *expected)
}
PermissionRuleMatcher::Prefix(prefix) => extract_permission_subject(input)
.is_some_and(|candidate| candidate.starts_with(prefix)),
}
}
}
fn parse_rule_matcher(content: &str) -> PermissionRuleMatcher {
let unescaped = unescape_rule_content(content.trim());
if unescaped.is_empty() || unescaped == "*" {
PermissionRuleMatcher::Any
} else if let Some(prefix) = unescaped.strip_suffix(":*") {
PermissionRuleMatcher::Prefix(prefix.to_string())
} else {
PermissionRuleMatcher::Exact(unescaped)
}
}
fn unescape_rule_content(content: &str) -> String {
content
.replace(r"\(", "(")
.replace(r"\)", ")")
.replace(r"\\", r"\")
}
fn find_first_unescaped(value: &str, needle: char) -> Option<usize> {
let mut escaped = false;
for (idx, ch) in value.char_indices() {
if ch == '\\' {
escaped = !escaped;
continue;
}
if ch == needle && !escaped {
return Some(idx);
}
escaped = false;
}
None
}
fn find_last_unescaped(value: &str, needle: char) -> Option<usize> {
let chars = value.char_indices().collect::<Vec<_>>();
for (pos, (idx, ch)) in chars.iter().enumerate().rev() {
if *ch != needle {
continue;
}
let mut backslashes = 0;
for (_, prev) in chars[..pos].iter().rev() {
if *prev == '\\' {
backslashes += 1;
} else {
break;
}
}
if backslashes % 2 == 0 {
return Some(*idx);
}
}
None
}
fn extract_permission_subject(input: &str) -> Option<String> {
let parsed = serde_json::from_str::<Value>(input).ok();
if let Some(Value::Object(object)) = parsed {
for key in [
"command",
"path",
"file_path",
"filePath",
"notebook_path",
"notebookPath",
"url",
"pattern",
"code",
"message",
] {
if let Some(value) = object.get(key).and_then(Value::as_str) {
return Some(value.to_string());
}
}
}
(!input.trim().is_empty()).then(|| input.to_string())
}
#[cfg(test)]
mod tests {
use super::{
PermissionContext, PermissionMode, PermissionOutcome, PermissionOverride, PermissionPolicy,
PermissionPromptDecision, PermissionPrompter, PermissionRequest,
};
use crate::config::RuntimePermissionRuleConfig;
struct RecordingPrompter {
seen: Vec<PermissionRequest>,
allow: bool,
}
impl PermissionPrompter for RecordingPrompter {
fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision {
self.seen.push(request.clone());
if self.allow {
PermissionPromptDecision::Allow
} else {
PermissionPromptDecision::Deny {
reason: "not now".to_string(),
}
}
}
}
#[test]
fn allows_tools_when_active_mode_meets_requirement() {
let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
.with_tool_requirement("read_file", PermissionMode::ReadOnly)
.with_tool_requirement("write_file", PermissionMode::WorkspaceWrite);
assert_eq!(
policy.authorize("read_file", "{}", None),
PermissionOutcome::Allow
);
assert_eq!(
policy.authorize("write_file", "{}", None),
PermissionOutcome::Allow
);
}
#[test]
fn denies_read_only_escalations_without_prompt() {
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
.with_tool_requirement("write_file", PermissionMode::WorkspaceWrite)
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
assert!(matches!(
policy.authorize("write_file", "{}", None),
PermissionOutcome::Deny { reason } if reason.contains("requires workspace-write permission")
));
assert!(matches!(
policy.authorize("bash", "{}", None),
PermissionOutcome::Deny { reason } if reason.contains("requires danger-full-access permission")
));
}
#[test]
fn prompts_for_workspace_write_to_danger_full_access_escalation() {
let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
let mut prompter = RecordingPrompter {
seen: Vec::new(),
allow: true,
};
let outcome = policy.authorize("bash", "echo hi", Some(&mut prompter));
assert_eq!(outcome, PermissionOutcome::Allow);
assert_eq!(prompter.seen.len(), 1);
assert_eq!(prompter.seen[0].tool_name, "bash");
assert_eq!(
prompter.seen[0].current_mode,
PermissionMode::WorkspaceWrite
);
assert_eq!(
prompter.seen[0].required_mode,
PermissionMode::DangerFullAccess
);
}
#[test]
fn honors_prompt_rejection_reason() {
let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
let mut prompter = RecordingPrompter {
seen: Vec::new(),
allow: false,
};
assert!(matches!(
policy.authorize("bash", "echo hi", Some(&mut prompter)),
PermissionOutcome::Deny { reason } if reason == "not now"
));
}
#[test]
fn applies_rule_based_denials_and_allows() {
let rules = RuntimePermissionRuleConfig::new(
vec!["bash(git:*)".to_string()],
vec!["bash(rm -rf:*)".to_string()],
Vec::new(),
);
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
.with_tool_requirement("bash", PermissionMode::DangerFullAccess)
.with_permission_rules(&rules);
assert_eq!(
policy.authorize("bash", r#"{"command":"git status"}"#, None),
PermissionOutcome::Allow
);
assert!(matches!(
policy.authorize("bash", r#"{"command":"rm -rf /tmp/x"}"#, None),
PermissionOutcome::Deny { reason } if reason.contains("denied by rule")
));
}
#[test]
fn ask_rules_force_prompt_even_when_mode_allows() {
let rules = RuntimePermissionRuleConfig::new(
Vec::new(),
Vec::new(),
vec!["bash(git:*)".to_string()],
);
let policy = PermissionPolicy::new(PermissionMode::DangerFullAccess)
.with_tool_requirement("bash", PermissionMode::DangerFullAccess)
.with_permission_rules(&rules);
let mut prompter = RecordingPrompter {
seen: Vec::new(),
allow: true,
};
let outcome = policy.authorize("bash", r#"{"command":"git status"}"#, Some(&mut prompter));
assert_eq!(outcome, PermissionOutcome::Allow);
assert_eq!(prompter.seen.len(), 1);
assert!(prompter.seen[0]
.reason
.as_deref()
.is_some_and(|reason| reason.contains("ask rule")));
}
#[test]
fn hook_allow_still_respects_ask_rules() {
let rules = RuntimePermissionRuleConfig::new(
Vec::new(),
Vec::new(),
vec!["bash(git:*)".to_string()],
);
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
.with_tool_requirement("bash", PermissionMode::DangerFullAccess)
.with_permission_rules(&rules);
let context = PermissionContext::new(
Some(PermissionOverride::Allow),
Some("hook approved".to_string()),
);
let mut prompter = RecordingPrompter {
seen: Vec::new(),
allow: true,
};
let outcome = policy.authorize_with_context(
"bash",
r#"{"command":"git status"}"#,
&context,
Some(&mut prompter),
);
assert_eq!(outcome, PermissionOutcome::Allow);
assert_eq!(prompter.seen.len(), 1);
}
#[test]
fn hook_deny_short_circuits_permission_flow() {
let policy = PermissionPolicy::new(PermissionMode::DangerFullAccess)
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
let context = PermissionContext::new(
Some(PermissionOverride::Deny),
Some("blocked by hook".to_string()),
);
assert_eq!(
policy.authorize_with_context("bash", "{}", &context, None),
PermissionOutcome::Deny {
reason: "blocked by hook".to_string(),
}
);
}
#[test]
fn hook_ask_forces_prompt() {
let policy = PermissionPolicy::new(PermissionMode::DangerFullAccess)
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
let context = PermissionContext::new(
Some(PermissionOverride::Ask),
Some("hook requested confirmation".to_string()),
);
let mut prompter = RecordingPrompter {
seen: Vec::new(),
allow: true,
};
let outcome = policy.authorize_with_context("bash", "{}", &context, Some(&mut prompter));
assert_eq!(outcome, PermissionOutcome::Allow);
assert_eq!(prompter.seen.len(), 1);
assert_eq!(
prompter.seen[0].reason.as_deref(),
Some("hook requested confirmation")
);
}
}
Loaded 100 of 268 files, more files were not shown because too many files have changed in this diff. Show more