Compare commits
12
Commits
ba0da84d59
..
rust
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d8d77824f4 | ||
|
|
4a04faf926 | ||
|
|
0782159ecd | ||
|
|
0d12ca9854 | ||
|
|
2639093fe4 | ||
|
|
16bd7df1b4 | ||
|
|
d6604ba0d9 | ||
|
|
07cdeef2ea | ||
|
|
f6a2b05045 | ||
|
|
84ab33a2d5 | ||
|
|
5d0e4e3456 | ||
|
|
c6643ca134 |
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"permissions": {
|
||||
"defaultMode": "dontAsk"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
ANTHROPIC_API_KEY="9494feba6f7c45f48c3dfc35a85ffd89.2WUCscxcSp92ETNg"
|
||||
ANTHROPIC_BASE_URL="https://open.bigmodel.cn/api/anthropic"
|
||||
CLAW_MODEL="glm-5"
|
||||
@@ -0,0 +1,36 @@
|
||||
name: CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
||||
jobs:
|
||||
rust:
|
||||
name: ${{ matrix.os }}
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os:
|
||||
- ubuntu-latest
|
||||
- macos-latest
|
||||
|
||||
steps:
|
||||
- name: Check out repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: Run cargo check
|
||||
run: cargo check --workspace
|
||||
|
||||
- name: Run cargo test
|
||||
run: cargo test --workspace
|
||||
|
||||
- name: Run release build
|
||||
run: cargo build --release
|
||||
@@ -1,3 +1,6 @@
|
||||
target/
|
||||
.omx/
|
||||
.clawd-agents/
|
||||
# Claw Code local artifacts
|
||||
.claw/
|
||||
.claude/
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
# CLAUDE.md
|
||||
|
||||
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
|
||||
|
||||
## Project Overview
|
||||
|
||||
Claw Code is a local coding-agent CLI tool written in safe Rust. It is a clean-room implementation inspired by Claude Code, providing an interactive REPL, one-shot prompts, workspace-aware tools, local agent workflows, and plugin support. The project name and references throughout use "claw" / "Claw Code".
|
||||
|
||||
## Build & Run Commands
|
||||
|
||||
```bash
|
||||
# Build release binary (produces target/release/claw)
|
||||
cargo build --release -p claw-cli
|
||||
|
||||
# Run from source (interactive REPL)
|
||||
cargo run --bin claw --
|
||||
|
||||
# Run one-shot prompt
|
||||
cargo run --bin claw -- prompt "summarize this workspace"
|
||||
|
||||
# Install locally
|
||||
cargo install --path crates/claw-cli --locked
|
||||
|
||||
# Run the HTTP server binary
|
||||
cargo run --bin claw-server
|
||||
```
|
||||
|
||||
## Verification Commands
|
||||
|
||||
```bash
|
||||
# Format check
|
||||
cargo fmt
|
||||
|
||||
# Lint (workspace-level clippy with deny warnings)
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
|
||||
# Run all tests
|
||||
cargo test --workspace
|
||||
|
||||
# Run tests for a specific crate
|
||||
cargo test -p api
|
||||
cargo test -p runtime
|
||||
|
||||
# Run a single test by name
|
||||
cargo test -p <crate> -- <test_name>
|
||||
|
||||
# Integration tests in crates/api/tests/ use mock TCP servers (no network needed)
|
||||
# One smoke test is #[ignore] — run with: cargo test -p api -- --ignored
|
||||
```
|
||||
|
||||
## Workspace Architecture
|
||||
|
||||
Cargo workspace with `resolver = "2"`. All crates live under `crates/`.
|
||||
|
||||
### Crate Dependency Graph
|
||||
|
||||
```
|
||||
claw-cli ──→ api, runtime, tools, commands, plugins, compat-harness
|
||||
server ──→ api, runtime, tools, plugins, commands
|
||||
tools ──→ api, runtime, plugins
|
||||
commands ──→ runtime, plugins
|
||||
api ──→ runtime
|
||||
runtime ──→ lsp, plugins
|
||||
plugins ──→ (standalone, serde only)
|
||||
lsp ──→ (standalone)
|
||||
compat-harness ──→ commands, tools, runtime
|
||||
```
|
||||
|
||||
### Core Crates
|
||||
|
||||
- **`claw-cli`** — User-facing binary (`claw`). REPL loop with markdown rendering (pulldown-cmark + syntect), argument parsing, OAuth flow. Entry point: `crates/claw-cli/src/main.rs`.
|
||||
|
||||
- **`runtime`** — Session management, conversation runtime, permissions, system prompt construction, context compaction, MCP stdio management, and hook execution. Key types: `Session` (versioned message history), `ConversationRuntime<C, T>` (generic over `ApiClient` + `ToolExecutor` traits), `PermissionMode`, `McpServerManager`, `HookRunner`.
|
||||
|
||||
- **`api`** — HTTP client for LLM providers with SSE streaming. `ClawApiClient` (Anthropic-compatible), `OpenAiCompatClient`, and `Provider` trait. `ProviderKind` enum distinguishes ClawApi, Xai, OpenAi. Request/response types: `MessageRequest`, `StreamEvent`, `ToolDefinition`.
|
||||
|
||||
- **`tools`** — Built-in tool definitions and dispatch. `GlobalToolRegistry` is a lazy-static singleton. Tools: Read, Write, Edit, Glob, Grep, Bash, LSP, Task*, Cron*, Worktree*. Each tool has a `ToolSpec` with JSON schema.
|
||||
|
||||
- **`commands`** — Slash command registry and handlers (`/help`, `/config`, `/compact`, `/resume`, `/plugins`, `/agents`, `/doctor`, etc.). `SlashCommandSpec` defines each command's name, aliases, description, and category.
|
||||
|
||||
- **`plugins`** — Plugin discovery and lifecycle. `PluginManager` loads builtin, bundled, and external (from `~/.claw/plugins/`) plugins. Plugins can provide additional tools via `PluginTool`.
|
||||
|
||||
- **`server`** — Axum-based HTTP server (`claw-server`). REST endpoints for session CRUD + SSE event streaming. `AppState` holds shared session store.
|
||||
|
||||
- **`lsp`** — Language Server Protocol types and process management for code intelligence features.
|
||||
|
||||
- **`compat-harness`** — Extracts command/tool/bootstrap-plan manifests from upstream TypeScript source files (for compatibility tracking). Uses `CLAUDE_CODE_UPSTREAM` env var to locate the upstream repo.
|
||||
|
||||
## Key Architectural Patterns
|
||||
|
||||
- **Trait-based abstraction**: `ApiClient`, `ToolExecutor`, `Provider` traits enable swappable implementations. `ConversationRuntime` is generic over client and executor.
|
||||
- **Static registries**: `GlobalToolRegistry` and slash command specs use lazy-static initialization with compile-time definitions.
|
||||
- **SSE streaming**: API responses stream through `MessageStream` (async iterator) to the terminal renderer or server SSE endpoints.
|
||||
- **Permission model**: `PermissionMode` enum — ReadOnly, WorkspaceWrite, DangerFullAccess. Configurable via `.claw.json` (`permissions.defaultMode`).
|
||||
- **Hook system**: Pre/post tool execution hooks via `HookRunner` in the runtime crate.
|
||||
|
||||
## Configuration & Environment
|
||||
|
||||
- `.claw.json` — Project-level config (permissions, etc.)
|
||||
- `.claw/` — Project directory for hooks, plugins, local settings
|
||||
- `~/.claw/` — User-level config directory
|
||||
- `.env` — API keys (gitignored): `ANTHROPIC_API_KEY`, `ANTHROPIC_BASE_URL`, `XAI_API_KEY`, `XAI_BASE_URL`
|
||||
- `CLAW.md` — Workspace instructions loaded into the system prompt (analogous to CLAUDE.md but for claw itself)
|
||||
|
||||
## Lint Rules
|
||||
|
||||
- `unsafe_code` is **forbidden** at workspace level
|
||||
- Clippy `all` + `pedantic` lints are warnings; some pedantic lints are allowed (`module_name_repetitions`, `missing_panics_doc`, `missing_errors_doc`)
|
||||
- CI runs: `cargo check --workspace`, `cargo test --workspace`, `cargo build --release` on Ubuntu and macOS
|
||||
|
||||
## Language
|
||||
|
||||
Code comments, commit messages, and documentation are primarily in Chinese (中文). UI strings and exported symbol names are in English.
|
||||
@@ -0,0 +1,15 @@
|
||||
# CLAW.md
|
||||
|
||||
This file provides guidance to Claw Code (clawcode.dev) when working with code in this repository.
|
||||
|
||||
## Detected stack
|
||||
- Languages: Rust.
|
||||
- Frameworks: none detected from the supported starter markers.
|
||||
|
||||
## Verification
|
||||
- Run Rust verification from the repo root: `cargo fmt`, `cargo clippy --workspace --all-targets -- -D warnings`, `cargo test --workspace`
|
||||
|
||||
## Working agreement
|
||||
- Prefer small, reviewable changes and keep generated bootstrap files aligned with actual repo workflows.
|
||||
- Keep shared defaults in `.claw.json`; reserve `.claw/settings.local.json` for machine-local overrides.
|
||||
- Do not overwrite existing `CLAW.md` content automatically; update it intentionally when repo workflows change.
|
||||
@@ -0,0 +1,43 @@
|
||||
# 贡献指南
|
||||
|
||||
感谢你为 Claw Code 做出贡献。
|
||||
|
||||
## 开发设置
|
||||
|
||||
- 安装稳定的 Rust 工具链。
|
||||
- 在此 Rust 工作区的仓库根目录下进行开发。如果你从父仓库根目录开始,请先执行 `cd rust/`。
|
||||
|
||||
## 构建
|
||||
|
||||
```bash
|
||||
cargo build
|
||||
cargo build --release
|
||||
```
|
||||
|
||||
## 测试与验证
|
||||
|
||||
在开启 Pull Request 之前,请运行完整的 Rust 验证集:
|
||||
|
||||
```bash
|
||||
cargo fmt --all --check
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
cargo check --workspace
|
||||
cargo test --workspace
|
||||
```
|
||||
|
||||
如果你更改了行为,请在同一个 Pull Request 中添加或更新相关的测试。
|
||||
|
||||
## 代码风格
|
||||
|
||||
- 遵循所修改 crate 中的现有模式,而不是引入新的风格。
|
||||
- 使用 `rustfmt` 格式化代码。
|
||||
- 确保你修改的工作区目标的 `clippy` 检查通过。
|
||||
- 优先采用针对性的 diff,而不是顺便进行的重构。
|
||||
|
||||
## Pull Request
|
||||
|
||||
- 从 `main` 分支拉取新分支。
|
||||
- 确保每个 Pull Request 的范围仅限于一个明确的更改。
|
||||
- 说明更改动机、实现摘要以及你运行的验证。
|
||||
- 在请求审查之前,确保本地检查已通过。
|
||||
- 如果审查反馈导致行为更改,请重新运行相关的验证命令。
|
||||
Generated
+572
-11
@@ -17,6 +17,15 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "android_system_properties"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "api"
|
||||
version = "0.1.0"
|
||||
@@ -25,15 +34,96 @@ dependencies = [
|
||||
"runtime",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"telemetry",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
|
||||
dependencies = [
|
||||
"async-stream-impl",
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream-impl"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "atomic-waker"
|
||||
version = "1.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
|
||||
|
||||
[[package]]
|
||||
name = "autocfg"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.8.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8"
|
||||
dependencies = [
|
||||
"axum-core",
|
||||
"bytes",
|
||||
"form_urlencoded",
|
||||
"futures-util",
|
||||
"http",
|
||||
"http-body",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"itoa",
|
||||
"matchit",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"serde_core",
|
||||
"serde_json",
|
||||
"serde_path_to_error",
|
||||
"serde_urlencoded",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.5.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"http",
|
||||
"http-body",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
"sync_wrapper",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "base64"
|
||||
version = "0.22.1"
|
||||
@@ -49,6 +139,12 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "1.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "2.11.0"
|
||||
@@ -98,11 +194,54 @@ version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
|
||||
|
||||
[[package]]
|
||||
name = "chrono"
|
||||
version = "0.4.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0"
|
||||
dependencies = [
|
||||
"iana-time-zone",
|
||||
"js-sys",
|
||||
"num-traits",
|
||||
"wasm-bindgen",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "claw-cli"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"api",
|
||||
"commands",
|
||||
"compat-harness",
|
||||
"crossterm",
|
||||
"dotenvy",
|
||||
"plugins",
|
||||
"pulldown-cmark",
|
||||
"runtime",
|
||||
"rustyline",
|
||||
"serde_json",
|
||||
"syntect",
|
||||
"tokio",
|
||||
"tools",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "clipboard-win"
|
||||
version = "5.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bde03770d3df201d4fb868f2c9c59e66a3e4e2bd06692a0fe701e7103c7e84d4"
|
||||
dependencies = [
|
||||
"error-code",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "commands"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"plugins",
|
||||
"runtime",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -114,6 +253,12 @@ dependencies = [
|
||||
"tools",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation-sys"
|
||||
version = "0.8.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
|
||||
|
||||
[[package]]
|
||||
name = "cpufeatures"
|
||||
version = "0.2.17"
|
||||
@@ -138,11 +283,11 @@ version = "0.28.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"crossterm_winapi",
|
||||
"mio",
|
||||
"parking_lot",
|
||||
"rustix",
|
||||
"rustix 0.38.44",
|
||||
"signal-hook",
|
||||
"signal-hook-mio",
|
||||
"winapi",
|
||||
@@ -197,6 +342,30 @@ dependencies = [
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dotenvy"
|
||||
version = "0.15.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b"
|
||||
|
||||
[[package]]
|
||||
name = "either"
|
||||
version = "1.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
|
||||
|
||||
[[package]]
|
||||
name = "endian-type"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d"
|
||||
|
||||
[[package]]
|
||||
name = "env_home"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe"
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
@@ -213,6 +382,23 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "error-code"
|
||||
version = "3.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59"
|
||||
|
||||
[[package]]
|
||||
name = "fd-lock"
|
||||
version = "4.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ce92ff622d6dadf7349484f42c93271a0d49b7cc4d466a936405bacbe10aa78"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"rustix 1.1.4",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.9"
|
||||
@@ -229,6 +415,15 @@ dependencies = [
|
||||
"miniz_oxide",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fluent-uri"
|
||||
version = "0.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "17c704e9dbe1ddd863da1e6ff3567795087b1eb201ce80d8fa81162e1516500d"
|
||||
dependencies = [
|
||||
"bitflags 1.3.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fnv"
|
||||
version = "1.0.7"
|
||||
@@ -266,6 +461,17 @@ version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
|
||||
|
||||
[[package]]
|
||||
name = "futures-macro"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-sink"
|
||||
version = "0.3.32"
|
||||
@@ -286,6 +492,7 @@ checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-io",
|
||||
"futures-macro",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"memchr",
|
||||
@@ -351,6 +558,15 @@ version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
|
||||
|
||||
[[package]]
|
||||
name = "home"
|
||||
version = "0.5.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "1.4.0"
|
||||
@@ -390,6 +606,12 @@ version = "1.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
|
||||
|
||||
[[package]]
|
||||
name = "httpdate"
|
||||
version = "1.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
|
||||
|
||||
[[package]]
|
||||
name = "hyper"
|
||||
version = "1.9.0"
|
||||
@@ -403,6 +625,7 @@ dependencies = [
|
||||
"http",
|
||||
"http-body",
|
||||
"httparse",
|
||||
"httpdate",
|
||||
"itoa",
|
||||
"pin-project-lite",
|
||||
"smallvec",
|
||||
@@ -450,6 +673,30 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone"
|
||||
version = "0.1.65"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
|
||||
dependencies = [
|
||||
"android_system_properties",
|
||||
"core-foundation-sys",
|
||||
"iana-time-zone-haiku",
|
||||
"js-sys",
|
||||
"log",
|
||||
"wasm-bindgen",
|
||||
"windows-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone-haiku"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_collections"
|
||||
version = "2.1.1"
|
||||
@@ -614,6 +861,12 @@ version = "0.4.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab"
|
||||
|
||||
[[package]]
|
||||
name = "linux-raw-sys"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
|
||||
|
||||
[[package]]
|
||||
name = "litemap"
|
||||
version = "0.8.1"
|
||||
@@ -641,12 +894,48 @@ version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
|
||||
|
||||
[[package]]
|
||||
name = "lsp"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"lsp-types",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lsp-types"
|
||||
version = "0.97.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "53353550a17c04ac46c585feb189c2db82154fc84b79c7a66c96c2c644f66071"
|
||||
dependencies = [
|
||||
"bitflags 1.3.2",
|
||||
"fluent-uri",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_repr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "matchit"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
|
||||
|
||||
[[package]]
|
||||
name = "mime"
|
||||
version = "0.3.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
|
||||
|
||||
[[package]]
|
||||
name = "miniz_oxide"
|
||||
version = "0.8.9"
|
||||
@@ -669,12 +958,51 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mock-anthropic-service"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"api",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nibble_vec"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43"
|
||||
dependencies = [
|
||||
"smallvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nix"
|
||||
version = "0.29.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"cfg-if",
|
||||
"cfg_aliases",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-conv"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967"
|
||||
|
||||
[[package]]
|
||||
name = "num-traits"
|
||||
version = "0.2.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "once_cell"
|
||||
version = "1.21.4"
|
||||
@@ -687,7 +1015,7 @@ version = "6.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "336b9c63443aceef14bea841b899035ae3abe89b7c486aaf4c5bd8aafedac3f0"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"onig_sys",
|
||||
@@ -757,6 +1085,14 @@ dependencies = [
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "plugins"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "potential_utf"
|
||||
version = "0.1.4"
|
||||
@@ -796,7 +1132,7 @@ version = "0.13.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c3a14896dfa883796f1cb410461aef38810ea05f2b2c33c5aded3649095fdad"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"getopts",
|
||||
"memchr",
|
||||
"pulldown-cmark-escape",
|
||||
@@ -888,6 +1224,16 @@ version = "5.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
|
||||
|
||||
[[package]]
|
||||
name = "radix_trie"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd"
|
||||
dependencies = [
|
||||
"endian-type",
|
||||
"nibble_vec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.9.2"
|
||||
@@ -923,7 +1269,7 @@ version = "0.5.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -985,12 +1331,14 @@ dependencies = [
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams",
|
||||
"web-sys",
|
||||
"webpki-roots",
|
||||
]
|
||||
@@ -1014,12 +1362,15 @@ name = "runtime"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"glob",
|
||||
"plugins",
|
||||
"regex",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"telemetry",
|
||||
"tokio",
|
||||
"walkdir",
|
||||
"which",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1034,11 +1385,24 @@ version = "0.38.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.52.0",
|
||||
"linux-raw-sys 0.4.15",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "1.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys 0.12.1",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1090,17 +1454,40 @@ dependencies = [
|
||||
"commands",
|
||||
"compat-harness",
|
||||
"crossterm",
|
||||
"mock-anthropic-service",
|
||||
"plugins",
|
||||
"pulldown-cmark",
|
||||
"reqwest",
|
||||
"runtime",
|
||||
"rustyline",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"syntect",
|
||||
"tokio",
|
||||
"tools",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustyline"
|
||||
version = "15.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2ee1e066dc922e513bda599c6ccb5f3bb2b0ea5870a579448f2622993f0a9a2f"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"cfg-if",
|
||||
"clipboard-win",
|
||||
"fd-lock",
|
||||
"home",
|
||||
"libc",
|
||||
"log",
|
||||
"memchr",
|
||||
"nix",
|
||||
"radix_trie",
|
||||
"unicode-segmentation",
|
||||
"unicode-width",
|
||||
"utf8parse",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ryu"
|
||||
version = "1.0.23"
|
||||
@@ -1165,6 +1552,28 @@ dependencies = [
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_path_to_error"
|
||||
version = "0.1.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"serde",
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_repr"
|
||||
version = "0.1.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_urlencoded"
|
||||
version = "0.7.1"
|
||||
@@ -1177,6 +1586,27 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "server"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"api",
|
||||
"async-stream",
|
||||
"axum",
|
||||
"chrono",
|
||||
"commands",
|
||||
"dotenvy",
|
||||
"plugins",
|
||||
"reqwest",
|
||||
"runtime",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tools",
|
||||
"tower",
|
||||
"tower-http",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.10.9"
|
||||
@@ -1317,6 +1747,14 @@ dependencies = [
|
||||
"yaml-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "telemetry"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "2.0.18"
|
||||
@@ -1430,11 +1868,27 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-util"
|
||||
version = "0.7.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tools"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"api",
|
||||
"commands",
|
||||
"flate2",
|
||||
"plugins",
|
||||
"reqwest",
|
||||
"runtime",
|
||||
"serde",
|
||||
@@ -1455,6 +1909,7 @@ dependencies = [
|
||||
"tokio",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1463,7 +1918,7 @@ version = "0.6.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http",
|
||||
@@ -1493,6 +1948,7 @@ version = "0.1.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100"
|
||||
dependencies = [
|
||||
"log",
|
||||
"pin-project-lite",
|
||||
"tracing-core",
|
||||
]
|
||||
@@ -1530,6 +1986,12 @@ version = "1.0.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-segmentation"
|
||||
version = "1.13.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-width"
|
||||
version = "0.2.2"
|
||||
@@ -1560,6 +2022,12 @@ version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "utf8parse"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
|
||||
|
||||
[[package]]
|
||||
name = "version_check"
|
||||
version = "0.9.5"
|
||||
@@ -1655,6 +2123,19 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-streams"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.93"
|
||||
@@ -1684,6 +2165,18 @@ dependencies = [
|
||||
"rustls-pki-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "which"
|
||||
version = "7.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "24d643ce3fd3e5b54854602a080f34fb10ab75e0b813ee32d00ca2b44fa74762"
|
||||
dependencies = [
|
||||
"either",
|
||||
"env_home",
|
||||
"rustix 1.1.4",
|
||||
"winsafe",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winapi"
|
||||
version = "0.3.9"
|
||||
@@ -1715,12 +2208,65 @@ version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
|
||||
|
||||
[[package]]
|
||||
name = "windows-core"
|
||||
version = "0.62.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
|
||||
dependencies = [
|
||||
"windows-implement",
|
||||
"windows-interface",
|
||||
"windows-link",
|
||||
"windows-result",
|
||||
"windows-strings",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-implement"
|
||||
version = "0.60.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-interface"
|
||||
version = "0.59.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "windows-result"
|
||||
version = "0.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-strings"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.52.0"
|
||||
@@ -1730,6 +2276,15 @@ dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.59.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
|
||||
dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.60.2"
|
||||
@@ -1877,6 +2432,12 @@ version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650"
|
||||
|
||||
[[package]]
|
||||
name = "winsafe"
|
||||
version = "0.0.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904"
|
||||
|
||||
[[package]]
|
||||
name = "wit-bindgen"
|
||||
version = "0.51.0"
|
||||
|
||||
+36
@@ -8,6 +8,10 @@ edition = "2021"
|
||||
license = "MIT"
|
||||
publish = false
|
||||
|
||||
[workspace.dependencies]
|
||||
lsp-types = "0.97"
|
||||
serde_json = "1"
|
||||
|
||||
[workspace.lints.rust]
|
||||
unsafe_code = "forbid"
|
||||
|
||||
@@ -17,3 +21,35 @@ pedantic = { level = "warn", priority = -1 }
|
||||
module_name_repetitions = "allow"
|
||||
missing_panics_doc = "allow"
|
||||
missing_errors_doc = "allow"
|
||||
uninlined_format_args = "allow"
|
||||
map_unwrap_or = "allow"
|
||||
doc_markdown = "allow"
|
||||
redundant_pattern_matching = "allow"
|
||||
items_after_statements = "allow"
|
||||
must_use_candidate = "allow"
|
||||
io_other_error = "allow"
|
||||
implicit_clone = "allow"
|
||||
redundant_closure_for_method_calls = "allow"
|
||||
unnecessary_map_or = "allow"
|
||||
manual_let_else = "allow"
|
||||
format_push_string = "allow"
|
||||
match_same_arms = "allow"
|
||||
similar_names = "allow"
|
||||
needless_continue = "allow"
|
||||
assigning_clones = "allow"
|
||||
cast_possible_truncation = "allow"
|
||||
cast_possible_wrap = "allow"
|
||||
cast_sign_loss = "allow"
|
||||
cmp_owned = "allow"
|
||||
collapsible_if = "allow"
|
||||
too_many_lines = "allow"
|
||||
wildcard_in_or_patterns = "allow"
|
||||
explicit_counter_loop = "allow"
|
||||
manual_repeat_n = "allow"
|
||||
manual_str_repeat = "allow"
|
||||
needless_borrow = "allow"
|
||||
needless_pass_by_value = "allow"
|
||||
single_match_else = "allow"
|
||||
too_many_arguments = "allow"
|
||||
unnested_or_patterns = "allow"
|
||||
unused_self = "allow"
|
||||
|
||||
@@ -1,230 +1,121 @@
|
||||
# Rusty Claude CLI
|
||||
# Claw Code
|
||||
|
||||
`rust/` contains the Rust workspace for the integrated `rusty-claude-cli` deliverable.
|
||||
It is intended to be something you can clone, build, and run directly.
|
||||
Claw Code 是一个使用安全 Rust 实现的本地编程代理(coding-agent)命令行工具。它的设计灵感来自 **Claude Code**,并作为一个**净室实现(clean-room implementation)**开发:旨在提供强大的本地代理体验,但它**不是** Claude Code 的直接移植或复制。
|
||||
|
||||
## Workspace layout
|
||||
Rust 工作区是当前主要的产品界面。`claw` 二进制文件在单个工作区内提供交互式会话、单次提示、工作区感知工具、本地代理工作流以及支持插件的操作。
|
||||
|
||||
```text
|
||||
rust/
|
||||
├── Cargo.toml
|
||||
├── Cargo.lock
|
||||
├── README.md
|
||||
└── crates/
|
||||
├── api/ # Anthropic API client + SSE streaming support
|
||||
├── commands/ # Shared slash-command metadata/help surfaces
|
||||
├── compat-harness/ # Upstream TS manifest extraction harness
|
||||
├── runtime/ # Session/runtime/config/prompt orchestration
|
||||
├── rusty-claude-cli/ # Main CLI binary
|
||||
└── tools/ # Built-in tool implementations
|
||||
## 当前状态
|
||||
|
||||
- **版本:** `0.1.0`
|
||||
- **发布阶段:** 初始公开发布,源码编译分发
|
||||
- **主要实现:** 本仓库中的 Rust 工作区
|
||||
- **平台焦点:** macOS 和 Linux 开发工作站
|
||||
|
||||
## 安装、构建与运行
|
||||
|
||||
### 准备工作
|
||||
|
||||
- Rust 稳定版工具链
|
||||
- Cargo
|
||||
- 你想使用的模型的提供商凭据
|
||||
|
||||
你可以通过环境变量或在项目根目录创建 **`.env`** 文件来配置 API 密钥:
|
||||
|
||||
**配置 Claude (推荐):**
|
||||
```bash
|
||||
ANTHROPIC_API_KEY="..."
|
||||
# 使用兼容的端点时可选
|
||||
export ANTHROPIC_BASE_URL="https://api.anthropic.com"
|
||||
```
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Rust toolchain installed (`rustup`, stable toolchain)
|
||||
- Network access and Anthropic credentials for live prompt/REPL usage
|
||||
|
||||
## Build
|
||||
|
||||
From the repository root:
|
||||
Grok 模型:
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo build --release -p rusty-claude-cli
|
||||
export XAI_API_KEY="..."
|
||||
# 使用兼容的端点时可选
|
||||
export XAI_BASE_URL="https://api.x.ai"
|
||||
```
|
||||
|
||||
The optimized binary will be written to:
|
||||
也可以使用 OAuth 登录:
|
||||
|
||||
```bash
|
||||
./target/release/rusty-claude-cli
|
||||
cargo run --bin claw -- login
|
||||
```
|
||||
|
||||
## Test
|
||||
|
||||
Run the verified workspace test suite used for release-readiness:
|
||||
### 本地安装
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo test --workspace --exclude compat-harness
|
||||
cargo install --path crates/claw-cli --locked
|
||||
```
|
||||
|
||||
## Quick start
|
||||
|
||||
### Show help
|
||||
### 从源码构建
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- --help
|
||||
cargo build --release -p claw-cli
|
||||
```
|
||||
|
||||
### Print version
|
||||
### 运行
|
||||
|
||||
在工作区内运行:
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- --version
|
||||
cargo run --bin claw -- --help
|
||||
cargo run --bin claw --
|
||||
cargo run --bin claw -- prompt "总结此工作区"
|
||||
cargo run --bin claw -- --model sonnet "审查最新更改"
|
||||
```
|
||||
|
||||
### Login with OAuth
|
||||
|
||||
Configure `settings.json` with an `oauth` block containing `clientId`, `authorizeUrl`, `tokenUrl`, optional `callbackPort`, and optional `scopes`, then run:
|
||||
运行发布版本:
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- login
|
||||
./target/release/claw
|
||||
./target/release/claw prompt "解释 crates/runtime"
|
||||
```
|
||||
|
||||
This opens the browser, listens on the configured localhost callback, exchanges the auth code for tokens, and stores OAuth credentials in `~/.claude/credentials.json` (or `$CLAUDE_CONFIG_HOME/credentials.json`).
|
||||
## 支持的功能
|
||||
|
||||
### Logout
|
||||
- 交互式 REPL 和单次提示执行
|
||||
- 已保存会话的检查和恢复流程
|
||||
- 内置工作区工具:shell、文件读/写/编辑、搜索、网页获取/搜索、待办事项和笔记本更新
|
||||
- 斜杠命令:状态、压缩、配置检查、差异(diff)、导出、会话管理和版本报告
|
||||
- 本地代理和技能发现:通过 `claw agents` 和 `claw skills`
|
||||
- 通过命令行和斜杠命令界面发现并管理插件
|
||||
- OAuth 登录/注销,以及从命令行选择模型/提供商
|
||||
- 工作区感知的指令/配置加载(`CLAW.md`、配置文件、权限、插件设置)
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- logout
|
||||
```
|
||||
## 当前限制
|
||||
|
||||
This removes only the stored OAuth credentials and preserves unrelated JSON fields in `credentials.json`.
|
||||
- 目前公开发布**仅限源码构建**;此工作区尚未设置 crates.io 发布
|
||||
- GitHub CI 验证 `cargo check`、`cargo test` 和发布构建,但尚未提供自动化的发布打包
|
||||
- 当前 CI 目标为 Ubuntu 和 macOS;Windows 的发布就绪性仍待建立
|
||||
- 一些实时提供商集成覆盖是可选的,因为它们需要外部凭据 and 网络访问
|
||||
- 命令界面可能会在 `0.x` 系列期间继续演进
|
||||
|
||||
### Self-update
|
||||
## 实现现状
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- self-update
|
||||
```
|
||||
Rust 工作区是当前的产品实现。目前包含以下 crate:
|
||||
|
||||
The command checks the latest GitHub release for `instructkr/clawd-code`, compares it to the current binary version, downloads the matching binary asset plus checksum manifest, verifies SHA-256, replaces the current executable, and prints the release changelog. If no published release or matching asset exists, it exits safely with an explanatory message.
|
||||
- `claw-cli` — 面向用户的二进制文件
|
||||
- `api` — 提供商客户端和流式处理
|
||||
- `runtime` — 会话、配置、权限、提示词和运行时循环
|
||||
- `tools` — 内置工具实现
|
||||
- `commands` — 斜杠命令注册和处理程序
|
||||
- `plugins` — 插件发现、注册和生命周期支持
|
||||
- `lsp` — 语言服务器协议支持类型和进程助手
|
||||
- `server` 和 `compat-harness` — 支持服务和兼容性工具
|
||||
|
||||
## Usage examples
|
||||
## 路线图
|
||||
|
||||
### 1) Prompt mode
|
||||
- 发布打包好的构件,用于公共安装
|
||||
- 添加可重复的发布工作流和长期维护的变更日志(changelog)规范
|
||||
- 将平台验证扩展到当前 CI 矩阵之外
|
||||
- 添加更多以任务为中心的示例和操作员文档
|
||||
- 继续加强 Rust 实现的功能覆盖并磨炼用户体验(UX)
|
||||
|
||||
Send one prompt, stream the answer, then exit:
|
||||
## 发行版本说明
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- prompt "Summarize the architecture of this repository"
|
||||
```
|
||||
- 0.1.0 发行说明草案:[`docs/releases/0.1.0.md`](docs/releases/0.1.0.md)
|
||||
|
||||
Use a specific model:
|
||||
## 许可
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- --model claude-sonnet-4-20250514 prompt "List the key crates in this workspace"
|
||||
```
|
||||
|
||||
Restrict enabled tools in an interactive session:
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- --allowedTools read,glob
|
||||
```
|
||||
|
||||
Bootstrap Claude project files for the current repo:
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- init
|
||||
```
|
||||
|
||||
### 2) REPL mode
|
||||
|
||||
Start the interactive shell:
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli --
|
||||
```
|
||||
|
||||
Inside the REPL, useful commands include:
|
||||
|
||||
```text
|
||||
/help
|
||||
/status
|
||||
/model claude-sonnet-4-20250514
|
||||
/permissions workspace-write
|
||||
/cost
|
||||
/compact
|
||||
/memory
|
||||
/config
|
||||
/init
|
||||
/diff
|
||||
/version
|
||||
/export notes.txt
|
||||
/sessions
|
||||
/session list
|
||||
/exit
|
||||
```
|
||||
|
||||
### 3) Resume an existing session
|
||||
|
||||
Inspect or maintain a saved session file without entering the REPL:
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- --resume session-123456 /status /compact /cost
|
||||
```
|
||||
|
||||
You can also inspect memory/config state for a restored session:
|
||||
|
||||
```bash
|
||||
cd rust
|
||||
cargo run -p rusty-claude-cli -- --resume ~/.claude/sessions/session-123456.json /memory /config
|
||||
```
|
||||
|
||||
## Available commands
|
||||
|
||||
### Top-level CLI commands
|
||||
|
||||
- `prompt <text...>` — run one prompt non-interactively
|
||||
- `--resume <session-id-or-path> [/commands...]` — inspect or maintain a saved session stored under `~/.claude/sessions/`
|
||||
- `dump-manifests` — print extracted upstream manifest counts
|
||||
- `bootstrap-plan` — print the current bootstrap skeleton
|
||||
- `system-prompt [--cwd PATH] [--date YYYY-MM-DD]` — render the synthesized system prompt
|
||||
- `self-update` — update the installed binary from the latest GitHub release when a matching asset is available
|
||||
- `--help` / `-h` — show CLI help
|
||||
- `--version` / `-V` — print the CLI version and build info locally (no API call)
|
||||
- `--output-format text|json` — choose non-interactive prompt output rendering
|
||||
- `--allowedTools <tool[,tool...]>` — restrict enabled tools for interactive sessions and prompt-mode tool use
|
||||
|
||||
### Interactive slash commands
|
||||
|
||||
- `/help` — show command help
|
||||
- `/status` — show current session status
|
||||
- `/compact` — compact local session history
|
||||
- `/model [model]` — inspect or switch the active model
|
||||
- `/permissions [read-only|workspace-write|danger-full-access]` — inspect or switch permissions
|
||||
- `/clear [--confirm]` — clear the current local session
|
||||
- `/cost` — show token usage totals
|
||||
- `/resume <session-id-or-path>` — load a saved session into the REPL
|
||||
- `/config [env|hooks|model]` — inspect discovered Claude config
|
||||
- `/memory` — inspect loaded instruction memory files
|
||||
- `/init` — bootstrap `.claude.json`, `.claude/`, `CLAUDE.md`, and local ignore rules
|
||||
- `/diff` — show the current git diff for the workspace
|
||||
- `/version` — print version and build metadata locally
|
||||
- `/export [file]` — export the current conversation transcript
|
||||
- `/sessions` — list recent managed local sessions from `~/.claude/sessions/`
|
||||
- `/session [list|switch <session-id>]` — inspect or switch managed local sessions
|
||||
- `/exit` — leave the REPL
|
||||
|
||||
## Environment variables
|
||||
|
||||
### Anthropic/API
|
||||
|
||||
- `ANTHROPIC_API_KEY` — highest-precedence API credential
|
||||
- `ANTHROPIC_AUTH_TOKEN` — bearer-token override used when no API key is set
|
||||
- Persisted OAuth credentials in `~/.claude/credentials.json` — used when neither env var is set
|
||||
- `ANTHROPIC_BASE_URL` — override the Anthropic API base URL
|
||||
- `ANTHROPIC_MODEL` — default model used by selected live integration tests
|
||||
|
||||
### CLI/runtime
|
||||
|
||||
- `RUSTY_CLAUDE_PERMISSION_MODE` — default REPL permission mode (`read-only`, `workspace-write`, or `danger-full-access`)
|
||||
- `CLAUDE_CONFIG_HOME` — override Claude config discovery root
|
||||
- `CLAUDE_CODE_REMOTE` — enable remote-session bootstrap handling when supported
|
||||
- `CLAUDE_CODE_REMOTE_SESSION_ID` — remote session identifier when using remote mode
|
||||
- `CLAUDE_CODE_UPSTREAM` — override the upstream TS source path for compat-harness extraction
|
||||
- `CLAWD_WEB_SEARCH_BASE_URL` — override the built-in web search service endpoint used by tooling
|
||||
|
||||
## Notes
|
||||
|
||||
- `compat-harness` exists to compare the Rust port against the upstream TypeScript codebase and is intentionally excluded from the requested release test run.
|
||||
- The CLI currently focuses on a practical integrated workflow: prompt execution, REPL operation, session inspection/resume, config discovery, and tool/runtime plumbing.
|
||||
有关许可详情,请参阅仓库根目录。
|
||||
|
||||
@@ -9,7 +9,8 @@ publish.workspace = true
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
|
||||
runtime = { path = "../runtime" }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
serde_json.workspace = true
|
||||
telemetry = { path = "../telemetry" }
|
||||
tokio = { version = "1", features = ["io-util", "macros", "net", "rt-multi-thread", "time"] }
|
||||
|
||||
[lints]
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
# API 模块 (api)
|
||||
|
||||
本模块提供了与大型语言模型 (LLM) 服务提供商(主要是 Anthropic 的 Claude 和兼容 OpenAI 的服务)进行交互的高层抽象和客户端。
|
||||
|
||||
## 概览
|
||||
|
||||
`api` 模块负责以下职责:
|
||||
- 标准化与不同 AI 提供商的通信。
|
||||
- 通过服务器发送事件 (SSE) 处理流式响应。
|
||||
- 管理身份验证源(API 密钥、OAuth 令牌)。
|
||||
- 提供消息、工具和使用情况跟踪的共享数据结构。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **提供商抽象 (Provider Abstraction)**:支持多种 AI 后端,包括:
|
||||
- `ClawApiClient`: Claude 模型的主要提供商。
|
||||
- `OpenAiCompatClient`: 支持兼容 OpenAI 的 API(如本地模型、专门的提供商)。
|
||||
- **流式支持 (Streaming Support)**:健壮的 SSE 解析实现 (`SseParser`),用于处理实时的内容生成。
|
||||
- **工具集成 (Tool Integration)**:为 `ToolDefinition`、`ToolChoice` 和 `ToolResultContentBlock` 提供强类型定义,支持智能代理 (Agentic) 工作流。
|
||||
- **身份验证管理 (Auth Management)**:用于解析启动身份验证源和管理 OAuth 令牌的实用工具。
|
||||
- **模型智能 (Model Intelligence)**:解析模型别名和计算最大标记 (Token) 限制的元数据及辅助函数。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`client.rs`**: 定义了 `ProviderClient` 特性 (Trait) 和基础客户端逻辑。它使用 `reqwest` 处理 HTTP 请求,并管理消息流的生命周期。
|
||||
- **`types.rs`**: 包含 API 的核心数据模型,如 `InputMessage`、`OutputContentBlock` 以及 `MessageRequest`/`MessageResponse`。
|
||||
- **`sse.rs`**: 实现了一个状态化的 SSE 解析器,能够处理分段的数据块并发出类型化的 `StreamEvent`。
|
||||
- **`providers/`**: 包含针对不同 LLM 端点的特定逻辑,将它们的独特格式映射到本模块使用的共享类型。
|
||||
|
||||
### 数据流
|
||||
|
||||
1. 构建包含模型详情、消息和工具定义的 `MessageRequest`。
|
||||
2. `ApiClient` 将此请求转换为提供商特定的 HTTP 请求。
|
||||
3. 如果启用了流式传输,客户端返回一个 `MessageStream`,该流使用 `SseParser` 来产生 `StreamEvent`。
|
||||
4. 最终响应包含用于跟踪 Token 消耗的 `Usage` 信息。
|
||||
|
||||
## 使用示例
|
||||
|
||||
```rust
|
||||
use api::{ApiClient, MessageRequest, InputMessage};
|
||||
|
||||
// 示例初始化(已简化)
|
||||
let client = ApiClient::new(auth_source);
|
||||
let request = MessageRequest {
|
||||
model: "claude-3-5-sonnet-20241022".to_string(),
|
||||
messages: vec![InputMessage::user("你好,世界!")],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let stream = client.create_message_stream(request).await?;
|
||||
```
|
||||
+161
-927
File diff suppressed because it is too large
Load Diff
+463
-25
@@ -2,19 +2,55 @@ 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 {
|
||||
MissingApiKey,
|
||||
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(serde_json::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,
|
||||
},
|
||||
@@ -30,18 +66,162 @@ pub enum ApiError {
|
||||
}
|
||||
|
||||
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::MissingApiKey
|
||||
Self::MissingCredentials { .. }
|
||||
| Self::ContextWindowExceeded { .. }
|
||||
| Self::ExpiredOAuthToken
|
||||
| Self::Auth(_)
|
||||
| Self::InvalidApiKeyEnv(_)
|
||||
| Self::Io(_)
|
||||
| Self::Json(_)
|
||||
| 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,
|
||||
}
|
||||
@@ -51,12 +231,44 @@ impl ApiError {
|
||||
impl Display for ApiError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::MissingApiKey => {
|
||||
Self::MissingCredentials {
|
||||
provider,
|
||||
env_vars,
|
||||
hint,
|
||||
} => {
|
||||
write!(
|
||||
f,
|
||||
"ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY is not set; export one before calling the Anthropic API"
|
||||
)
|
||||
"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,
|
||||
@@ -65,36 +277,45 @@ impl Display for ApiError {
|
||||
}
|
||||
Self::Auth(message) => write!(f, "auth error: {message}"),
|
||||
Self::InvalidApiKeyEnv(error) => {
|
||||
write!(
|
||||
f,
|
||||
"failed to read ANTHROPIC_AUTH_TOKEN / ANTHROPIC_API_KEY: {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(error) => write!(f, "json 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,
|
||||
..
|
||||
} => match (error_type, message) {
|
||||
(Some(error_type), Some(message)) => {
|
||||
write!(
|
||||
f,
|
||||
"anthropic api returned {status} ({error_type}): {message}"
|
||||
)
|
||||
} => {
|
||||
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}")
|
||||
}
|
||||
_ => write!(f, "anthropic api returned {status}: {body}"),
|
||||
},
|
||||
}
|
||||
Self::RetriesExhausted {
|
||||
attempts,
|
||||
last_error,
|
||||
} => write!(
|
||||
f,
|
||||
"anthropic api failed after {attempts} attempts: {last_error}"
|
||||
),
|
||||
} => write!(f, "api failed after {attempts} attempts: {last_error}"),
|
||||
Self::InvalidSseFrame(message) => write!(f, "invalid sse frame: {message}"),
|
||||
Self::BackoffOverflow {
|
||||
attempt,
|
||||
@@ -123,7 +344,12 @@ impl From<std::io::Error> for ApiError {
|
||||
|
||||
impl From<serde_json::Error> for ApiError {
|
||||
fn from(value: serde_json::Error) -> Self {
|
||||
Self::Json(value)
|
||||
Self::Json {
|
||||
provider: "unknown".to_string(),
|
||||
model: "unknown".to_string(),
|
||||
body_snippet: String::new(),
|
||||
source: value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -132,3 +358,215 @@ impl From<VarError> for ApiError {
|
||||
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:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
+26
-4
@@ -1,17 +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, resolve_saved_oauth_token, resolve_startup_auth_source,
|
||||
AnthropicClient, AuthSource, MessageStream, OAuthTokenSet,
|
||||
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,
|
||||
ImageSource, InputContentBlock, InputMessage, MessageDelta, MessageDeltaEvent, MessageRequest,
|
||||
InputContentBlock, InputMessage, MessageDelta, MessageDeltaEvent, MessageRequest,
|
||||
MessageResponse, MessageStartEvent, MessageStopEvent, OutputContentBlock, StreamEvent,
|
||||
ThinkingConfig, ToolChoice, ToolDefinition, ToolResultContentBlock, Usage,
|
||||
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(), ¤t);
|
||||
|
||||
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), ¤t)
|
||||
.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(
|
||||
¤t_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), ¤t)
|
||||
.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
+114
-3
@@ -4,6 +4,8 @@ use crate::types::StreamEvent;
|
||||
#[derive(Debug, Default)]
|
||||
pub struct SseParser {
|
||||
buffer: Vec<u8>,
|
||||
provider: Option<String>,
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
impl SseParser {
|
||||
@@ -12,12 +14,23 @@ impl SseParser {
|
||||
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) = parse_frame(&frame)? {
|
||||
if let Some(event) = self.parse_frame_with_context(&frame)? {
|
||||
events.push(event);
|
||||
}
|
||||
}
|
||||
@@ -31,12 +44,18 @@ impl SseParser {
|
||||
}
|
||||
|
||||
let trailing = std::mem::take(&mut self.buffer);
|
||||
match parse_frame(&String::from_utf8_lossy(&trailing))? {
|
||||
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
|
||||
@@ -61,6 +80,14 @@ impl SseParser {
|
||||
}
|
||||
|
||||
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);
|
||||
@@ -97,7 +124,7 @@ pub fn parse_frame(frame: &str) -> Result<Option<StreamEvent>, ApiError> {
|
||||
|
||||
serde_json::from_str::<StreamEvent>(&payload)
|
||||
.map(Some)
|
||||
.map_err(ApiError::from)
|
||||
.map_err(|error| ApiError::json_deserialize(provider, model, &payload, error))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -216,4 +243,88 @@ mod tests {
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[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(),
|
||||
}))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+100
-39
@@ -1,7 +1,8 @@
|
||||
use runtime::{pricing_for_model, TokenUsage, UsageCostEstimate};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
|
||||
pub struct MessageRequest {
|
||||
pub model: String,
|
||||
pub max_tokens: u32,
|
||||
@@ -12,10 +13,24 @@ pub struct MessageRequest {
|
||||
pub tools: Option<Vec<ToolDefinition>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_choice: Option<ToolChoice>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thinking: Option<ThinkingConfig>,
|
||||
#[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 {
|
||||
@@ -26,23 +41,6 @@ impl MessageRequest {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ThinkingConfig {
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String,
|
||||
pub budget_tokens: u32,
|
||||
}
|
||||
|
||||
impl ThinkingConfig {
|
||||
#[must_use]
|
||||
pub fn enabled(budget_tokens: u32) -> Self {
|
||||
Self {
|
||||
kind: "enabled".to_string(),
|
||||
budget_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct InputMessage {
|
||||
pub role: String,
|
||||
@@ -83,9 +81,6 @@ pub enum InputContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
Image {
|
||||
source: ImageSource,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
@@ -99,14 +94,6 @@ pub enum InputContentBlock {
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ImageSource {
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String,
|
||||
pub media_type: String,
|
||||
pub data: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ToolResultContentBlock {
|
||||
@@ -142,6 +129,7 @@ pub struct MessageResponse {
|
||||
pub stop_reason: Option<String>,
|
||||
#[serde(default)]
|
||||
pub stop_sequence: Option<String>,
|
||||
#[serde(default)]
|
||||
pub usage: Usage,
|
||||
#[serde(default)]
|
||||
pub request_id: Option<String>,
|
||||
@@ -160,32 +148,60 @@ pub enum OutputContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
Thinking {
|
||||
thinking: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
signature: Option<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, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[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.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),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -197,6 +213,7 @@ pub struct MessageStartEvent {
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct MessageDeltaEvent {
|
||||
pub delta: MessageDelta,
|
||||
#[serde(default)]
|
||||
pub usage: Usage,
|
||||
}
|
||||
|
||||
@@ -224,9 +241,9 @@ pub struct ContentBlockDeltaEvent {
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ContentBlockDelta {
|
||||
TextDelta { text: String },
|
||||
InputJsonDelta { partial_json: String },
|
||||
ThinkingDelta { thinking: String },
|
||||
SignatureDelta { signature: String },
|
||||
InputJsonDelta { partial_json: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
@@ -247,3 +264,47 @@ pub enum StreamEvent {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,17 +1,27 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Mutex as StdMutex, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use api::{
|
||||
AnthropicClient, ApiError, ContentBlockDelta, ContentBlockDeltaEvent, ContentBlockStartEvent,
|
||||
ImageSource, InputContentBlock, InputMessage, MessageDeltaEvent, MessageRequest,
|
||||
OutputContentBlock, StreamEvent, ToolChoice, ToolDefinition,
|
||||
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()));
|
||||
@@ -34,7 +44,7 @@ async fn send_message_posts_json_and_parses_response() {
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = AnthropicClient::new("test-key")
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_auth_token(Some("proxy-token".to_string()))
|
||||
.with_base_url(server.base_url());
|
||||
let response = client
|
||||
@@ -45,6 +55,8 @@ async fn send_message_posts_json_and_parses_response() {
|
||||
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 {
|
||||
@@ -64,6 +76,18 @@ async fn send_message_posts_json_and_parses_response() {
|
||||
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!(
|
||||
@@ -73,47 +97,237 @@ async fn send_message_posts_json_and_parses_response() {
|
||||
assert!(body.get("stream").is_none());
|
||||
assert_eq!(body["tools"][0]["name"], json!("get_weather"));
|
||||
assert_eq!(body["tool_choice"]["type"], json!("auto"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_content_blocks_serialize_with_base64_source() {
|
||||
let request = MessageRequest {
|
||||
model: "claude-3-7-sonnet-latest".to_string(),
|
||||
max_tokens: 64,
|
||||
messages: vec![InputMessage {
|
||||
role: "user".to_string(),
|
||||
content: vec![InputContentBlock::Image {
|
||||
source: ImageSource {
|
||||
kind: "base64".to_string(),
|
||||
media_type: "image/png".to_string(),
|
||||
data: "AQID".to_string(),
|
||||
},
|
||||
}],
|
||||
}],
|
||||
system: None,
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
stream: false,
|
||||
};
|
||||
|
||||
let json = serde_json::to_value(request).expect("request should serialize");
|
||||
assert_eq!(json["messages"][0]["content"][0]["type"], json!("image"));
|
||||
assert_eq!(
|
||||
json["messages"][0]["content"][0]["source"],
|
||||
json!({
|
||||
"type": "base64",
|
||||
"media_type": "image/png",
|
||||
"data": "AQID"
|
||||
})
|
||||
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,\"output_tokens\":0}}}\n\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",
|
||||
@@ -121,7 +335,7 @@ async fn stream_message_parses_sse_events_with_tool_use() {
|
||||
"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,\"output_tokens\":1}}\n\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"
|
||||
@@ -137,9 +351,10 @@ async fn stream_message_parses_sse_events_with_tool_use() {
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = AnthropicClient::new("test-key")
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_auth_token(Some("proxy-token".to_string()))
|
||||
.with_base_url(server.base_url());
|
||||
.with_base_url(server.base_url())
|
||||
.with_prompt_cache(PromptCache::new("stream-session"));
|
||||
let mut stream = client
|
||||
.stream_message(&sample_request(false))
|
||||
.await
|
||||
@@ -193,6 +408,20 @@ async fn stream_message_parses_sse_events_with_tool_use() {
|
||||
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]
|
||||
@@ -215,7 +444,7 @@ async fn retries_retryable_failures_before_succeeding() {
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = AnthropicClient::new("test-key")
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_base_url(server.base_url())
|
||||
.with_retry_policy(2, Duration::from_millis(1), Duration::from_millis(2));
|
||||
|
||||
@@ -228,6 +457,47 @@ async fn retries_retryable_failures_before_succeeding() {
|
||||
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()));
|
||||
@@ -248,7 +518,7 @@ async fn surfaces_retry_exhaustion_for_persistent_retryable_errors() {
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = AnthropicClient::new("test-key")
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_base_url(server.base_url())
|
||||
.with_retry_policy(1, Duration::from_millis(1), Duration::from_millis(2));
|
||||
|
||||
@@ -276,10 +546,190 @@ async fn surfaces_retry_exhaustion_for_persistent_retryable_errors() {
|
||||
}
|
||||
}
|
||||
|
||||
#[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 = AnthropicClient::from_env().expect("ANTHROPIC_API_KEY must be set");
|
||||
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")
|
||||
@@ -291,8 +741,8 @@ async fn live_stream_smoke_test() {
|
||||
system: None,
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
thinking: None,
|
||||
stream: false,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.expect("live stream should start");
|
||||
@@ -472,7 +922,7 @@ fn sample_request(stream: bool) -> MessageRequest {
|
||||
}),
|
||||
}]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
thinking: None,
|
||||
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"));
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
[package]
|
||||
name = "claw-cli"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "claw"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
api = { path = "../api" }
|
||||
commands = { path = "../commands" }
|
||||
compat-harness = { path = "../compat-harness" }
|
||||
crossterm = "0.28"
|
||||
pulldown-cmark = "0.13"
|
||||
rustyline = "15"
|
||||
runtime = { path = "../runtime" }
|
||||
plugins = { path = "../plugins" }
|
||||
serde_json.workspace = true
|
||||
syntect = "5"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "time"] }
|
||||
tools = { path = "../tools" }
|
||||
dotenvy = "0.15"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,58 @@
|
||||
# Claw CLI 模块 (claw-cli)
|
||||
|
||||
本模块实现了 Claw 应用程序的主要命令行界面 (CLI)。它提供了交互式的 REPL 环境和非交互式的命令执行功能。
|
||||
|
||||
## 概览
|
||||
|
||||
`claw-cli` 模块是整个项目的“胶水”,负责编排 `runtime`、`api`、`tools` 和 `plugins` 模块之间的交互。它捕获用户输入,管理应用程序状态,并以用户友好的格式渲染 AI 响应。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **交互式 REPL**:一个功能齐全的 Read-Eval-Print Loop,用于与 AI 对话,支持:
|
||||
- 通过 `rustyline` 实现多行输入和命令历史。
|
||||
- 实时的 Markdown 流式显示和代码语法高亮。
|
||||
- 为耗时的工具调用显示动态的加载动画 (Spinner) 和进度指示器。
|
||||
- **子命令**:
|
||||
- `prompt`:运行单次 Prompt 并退出(单次模式)。
|
||||
- `login`/`logout`:处理与 Claw 平台的 OAuth 身份验证。
|
||||
- `init`:初始化新项目/仓库以配合 Claw 使用。
|
||||
- `resume`:恢复并继续之前的对话会话。
|
||||
- `agents`/`skills`:管理并发现可用的智能体和技能。
|
||||
- **权限管理**:对工具执行权限的细粒度控制:
|
||||
- `read-only`:安全的分析模式。
|
||||
- `workspace-write`:允许在当前工作区内进行修改。
|
||||
- `danger-full-access`:高级任务的无限制访问。
|
||||
- **OAuth 流程**:集成了本地 HTTP 服务器,无缝处理用户机器上的 OAuth 回调重定向。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`main.rs`**: 主要入口点。负责解析命令行参数、初始化环境,并调度到适当的操作(REPL 或子命令)。
|
||||
- **`render/`**: 包含 `TerminalRenderer` 和 Markdown 流渲染逻辑。使用 `syntect` 进行语法高亮,并使用 `crossterm` 进行终端操作。
|
||||
- **`input/`**: 处理用户输入捕获,包括对多行 Prompt 和斜杠命令 (Slash Command) 的特殊处理。
|
||||
- **`init.rs`**: 处理项目级初始化和仓库设置。
|
||||
|
||||
### 交互循环流程
|
||||
|
||||
1. CLI 初始化 `ConversationRuntime` 并加载项目上下文。
|
||||
2. 使用 `rustyline` 进入捕获用户输入的循环。
|
||||
3. 检查用户输入是否为“斜杠命令”(如 `/compact`、`/model`)。
|
||||
4. 普通 Prompt 通过 `runtime` 发送给 AI。
|
||||
5. AI 事件(文本增量、工具调用)逐步渲染到终端。
|
||||
6. 会话定期保存,以便将来可以恢复。
|
||||
|
||||
## 使用方法
|
||||
|
||||
主二进制程序名为 `claw`。
|
||||
|
||||
```bash
|
||||
# 启动交互式 REPL
|
||||
claw
|
||||
|
||||
# 运行单次 Prompt
|
||||
claw prompt "解释一下这个项目的架构"
|
||||
|
||||
# 登录服务
|
||||
claw login
|
||||
```
|
||||
@@ -112,7 +112,7 @@ impl CliApp {
|
||||
|
||||
pub fn run_repl(&mut self) -> io::Result<()> {
|
||||
let mut editor = LineEditor::new("› ", Vec::new());
|
||||
println!("Rusty Claude CLI interactive mode");
|
||||
println!("Claw Code interactive mode");
|
||||
println!("Type /help for commands. Shift+Enter or Ctrl+J inserts a newline.");
|
||||
|
||||
loop {
|
||||
@@ -162,6 +162,10 @@ impl CliApp {
|
||||
writeln!(out, "Unknown slash command: /{name}")?;
|
||||
Ok(CommandResult::Continue)
|
||||
}
|
||||
_ => {
|
||||
writeln!(out, "Slash command unavailable in this mode")?;
|
||||
Ok(CommandResult::Continue)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -172,7 +176,7 @@ impl CliApp {
|
||||
SlashCommand::Help => "/help",
|
||||
SlashCommand::Status => "/status",
|
||||
SlashCommand::Compact => "/compact",
|
||||
SlashCommand::Unknown(_) => continue,
|
||||
_ => continue,
|
||||
};
|
||||
writeln!(out, " {name:<9} {}", handler.summary)?;
|
||||
}
|
||||
@@ -385,14 +389,14 @@ mod tests {
|
||||
#[test]
|
||||
fn session_state_tracks_config_values() {
|
||||
let config = SessionConfig {
|
||||
model: "claude".into(),
|
||||
permission_mode: PermissionMode::WorkspaceWrite,
|
||||
model: "sonnet".into(),
|
||||
permission_mode: PermissionMode::DangerFullAccess,
|
||||
config: Some(PathBuf::from("settings.toml")),
|
||||
output_format: OutputFormat::Text,
|
||||
};
|
||||
|
||||
assert_eq!(config.model, "claude");
|
||||
assert_eq!(config.permission_mode, PermissionMode::WorkspaceWrite);
|
||||
assert_eq!(config.model, "sonnet");
|
||||
assert_eq!(config.permission_mode, PermissionMode::DangerFullAccess);
|
||||
assert_eq!(config.config, Some(PathBuf::from("settings.toml")));
|
||||
}
|
||||
}
|
||||
@@ -3,16 +3,12 @@ use std::path::PathBuf;
|
||||
use clap::{Parser, Subcommand, ValueEnum};
|
||||
|
||||
#[derive(Debug, Clone, Parser, PartialEq, Eq)]
|
||||
#[command(
|
||||
name = "rusty-claude-cli",
|
||||
version,
|
||||
about = "Rust Claude CLI prototype"
|
||||
)]
|
||||
#[command(name = "claw-cli", version, about = "Claw Code CLI")]
|
||||
pub struct Cli {
|
||||
#[arg(long, default_value = "claude-3-7-sonnet")]
|
||||
#[arg(long, default_value = "claude-opus-4-6")]
|
||||
pub model: String,
|
||||
|
||||
#[arg(long, value_enum, default_value_t = PermissionMode::WorkspaceWrite)]
|
||||
#[arg(long, value_enum, default_value_t = PermissionMode::DangerFullAccess)]
|
||||
pub permission_mode: PermissionMode,
|
||||
|
||||
#[arg(long)]
|
||||
@@ -62,9 +58,9 @@ mod tests {
|
||||
#[test]
|
||||
fn parses_requested_flags() {
|
||||
let cli = Cli::parse_from([
|
||||
"rusty-claude-cli",
|
||||
"claw-cli",
|
||||
"--model",
|
||||
"claude-3-5-haiku",
|
||||
"claude-haiku-4-5-20251213",
|
||||
"--permission-mode",
|
||||
"read-only",
|
||||
"--config",
|
||||
@@ -76,7 +72,7 @@ mod tests {
|
||||
"world",
|
||||
]);
|
||||
|
||||
assert_eq!(cli.model, "claude-3-5-haiku");
|
||||
assert_eq!(cli.model, "claude-haiku-4-5-20251213");
|
||||
assert_eq!(cli.permission_mode, PermissionMode::ReadOnly);
|
||||
assert_eq!(
|
||||
cli.config.as_deref(),
|
||||
@@ -93,10 +89,16 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn parses_login_and_logout_commands() {
|
||||
let login = Cli::parse_from(["rusty-claude-cli", "login"]);
|
||||
let login = Cli::parse_from(["claw-cli", "login"]);
|
||||
assert_eq!(login.command, Some(Command::Login));
|
||||
|
||||
let logout = Cli::parse_from(["rusty-claude-cli", "logout"]);
|
||||
let logout = Cli::parse_from(["claw-cli", "logout"]);
|
||||
assert_eq!(logout.command, Some(Command::Logout));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defaults_to_danger_full_access_permission_mode() {
|
||||
let cli = Cli::parse_from(["claw-cli"]);
|
||||
assert_eq!(cli.permission_mode, PermissionMode::DangerFullAccess);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,432 @@
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const STARTER_CLAW_JSON: &str = concat!(
|
||||
"{\n",
|
||||
" \"permissions\": {\n",
|
||||
" \"defaultMode\": \"dontAsk\"\n",
|
||||
" }\n",
|
||||
"}\n",
|
||||
);
|
||||
const GITIGNORE_COMMENT: &str = "# Claw Code local artifacts";
|
||||
const GITIGNORE_ENTRIES: [&str; 2] = [".claw/settings.local.json", ".claw/sessions/"];
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum InitStatus {
|
||||
Created,
|
||||
Updated,
|
||||
Skipped,
|
||||
}
|
||||
|
||||
impl InitStatus {
|
||||
#[must_use]
|
||||
pub(crate) fn label(self) -> &'static str {
|
||||
match self {
|
||||
Self::Created => "created",
|
||||
Self::Updated => "updated",
|
||||
Self::Skipped => "skipped (already exists)",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct InitArtifact {
|
||||
pub(crate) name: &'static str,
|
||||
pub(crate) status: InitStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct InitReport {
|
||||
pub(crate) project_root: PathBuf,
|
||||
pub(crate) artifacts: Vec<InitArtifact>,
|
||||
}
|
||||
|
||||
impl InitReport {
|
||||
#[must_use]
|
||||
pub(crate) fn render(&self) -> String {
|
||||
let mut lines = vec![
|
||||
"Init".to_string(),
|
||||
format!(" Project {}", self.project_root.display()),
|
||||
];
|
||||
for artifact in &self.artifacts {
|
||||
lines.push(format!(
|
||||
" {:<16} {}",
|
||||
artifact.name,
|
||||
artifact.status.label()
|
||||
));
|
||||
}
|
||||
lines.push(" Next step Review and tailor the generated guidance".to_string());
|
||||
lines.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
#[allow(clippy::struct_excessive_bools)]
|
||||
struct RepoDetection {
|
||||
rust_workspace: bool,
|
||||
rust_root: bool,
|
||||
python: bool,
|
||||
package_json: bool,
|
||||
typescript: bool,
|
||||
nextjs: bool,
|
||||
react: bool,
|
||||
vite: bool,
|
||||
nest: bool,
|
||||
src_dir: bool,
|
||||
tests_dir: bool,
|
||||
rust_dir: bool,
|
||||
}
|
||||
|
||||
pub(crate) fn initialize_repo(cwd: &Path) -> Result<InitReport, Box<dyn std::error::Error>> {
|
||||
let mut artifacts = Vec::new();
|
||||
|
||||
let claw_dir = cwd.join(".claw");
|
||||
artifacts.push(InitArtifact {
|
||||
name: ".claw/",
|
||||
status: ensure_dir(&claw_dir)?,
|
||||
});
|
||||
|
||||
let claw_json = cwd.join(".claw.json");
|
||||
artifacts.push(InitArtifact {
|
||||
name: ".claw.json",
|
||||
status: write_file_if_missing(&claw_json, STARTER_CLAW_JSON)?,
|
||||
});
|
||||
|
||||
let gitignore = cwd.join(".gitignore");
|
||||
artifacts.push(InitArtifact {
|
||||
name: ".gitignore",
|
||||
status: ensure_gitignore_entries(&gitignore)?,
|
||||
});
|
||||
|
||||
let claw_md = cwd.join("CLAW.md");
|
||||
let content = render_init_claw_md(cwd);
|
||||
artifacts.push(InitArtifact {
|
||||
name: "CLAW.md",
|
||||
status: write_file_if_missing(&claw_md, &content)?,
|
||||
});
|
||||
|
||||
Ok(InitReport {
|
||||
project_root: cwd.to_path_buf(),
|
||||
artifacts,
|
||||
})
|
||||
}
|
||||
|
||||
fn ensure_dir(path: &Path) -> Result<InitStatus, std::io::Error> {
|
||||
if path.is_dir() {
|
||||
return Ok(InitStatus::Skipped);
|
||||
}
|
||||
fs::create_dir_all(path)?;
|
||||
Ok(InitStatus::Created)
|
||||
}
|
||||
|
||||
fn write_file_if_missing(path: &Path, content: &str) -> Result<InitStatus, std::io::Error> {
|
||||
if path.exists() {
|
||||
return Ok(InitStatus::Skipped);
|
||||
}
|
||||
fs::write(path, content)?;
|
||||
Ok(InitStatus::Created)
|
||||
}
|
||||
|
||||
fn ensure_gitignore_entries(path: &Path) -> Result<InitStatus, std::io::Error> {
|
||||
if !path.exists() {
|
||||
let mut lines = vec![GITIGNORE_COMMENT.to_string()];
|
||||
lines.extend(GITIGNORE_ENTRIES.iter().map(|entry| (*entry).to_string()));
|
||||
fs::write(path, format!("{}\n", lines.join("\n")))?;
|
||||
return Ok(InitStatus::Created);
|
||||
}
|
||||
|
||||
let existing = fs::read_to_string(path)?;
|
||||
let mut lines = existing.lines().map(ToOwned::to_owned).collect::<Vec<_>>();
|
||||
let mut changed = false;
|
||||
|
||||
if !lines.iter().any(|line| line == GITIGNORE_COMMENT) {
|
||||
lines.push(GITIGNORE_COMMENT.to_string());
|
||||
changed = true;
|
||||
}
|
||||
|
||||
for entry in GITIGNORE_ENTRIES {
|
||||
if !lines.iter().any(|line| line == entry) {
|
||||
lines.push(entry.to_string());
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
if !changed {
|
||||
return Ok(InitStatus::Skipped);
|
||||
}
|
||||
|
||||
fs::write(path, format!("{}\n", lines.join("\n")))?;
|
||||
Ok(InitStatus::Updated)
|
||||
}
|
||||
|
||||
pub(crate) fn render_init_claw_md(cwd: &Path) -> String {
|
||||
let detection = detect_repo(cwd);
|
||||
let mut lines = vec![
|
||||
"# CLAW.md".to_string(),
|
||||
String::new(),
|
||||
"This file provides guidance to Claw Code (clawcode.dev) when working with code in this repository.".to_string(),
|
||||
String::new(),
|
||||
];
|
||||
|
||||
let detected_languages = detected_languages(&detection);
|
||||
let detected_frameworks = detected_frameworks(&detection);
|
||||
lines.push("## Detected stack".to_string());
|
||||
if detected_languages.is_empty() {
|
||||
lines.push("- No specific language markers were detected yet; document the primary language and verification commands once the project structure settles.".to_string());
|
||||
} else {
|
||||
lines.push(format!("- Languages: {}.", detected_languages.join(", ")));
|
||||
}
|
||||
if detected_frameworks.is_empty() {
|
||||
lines.push("- Frameworks: none detected from the supported starter markers.".to_string());
|
||||
} else {
|
||||
lines.push(format!(
|
||||
"- Frameworks/tooling markers: {}.",
|
||||
detected_frameworks.join(", ")
|
||||
));
|
||||
}
|
||||
lines.push(String::new());
|
||||
|
||||
let verification_lines = verification_lines(cwd, &detection);
|
||||
if !verification_lines.is_empty() {
|
||||
lines.push("## Verification".to_string());
|
||||
lines.extend(verification_lines);
|
||||
lines.push(String::new());
|
||||
}
|
||||
|
||||
let structure_lines = repository_shape_lines(&detection);
|
||||
if !structure_lines.is_empty() {
|
||||
lines.push("## Repository shape".to_string());
|
||||
lines.extend(structure_lines);
|
||||
lines.push(String::new());
|
||||
}
|
||||
|
||||
let framework_lines = framework_notes(&detection);
|
||||
if !framework_lines.is_empty() {
|
||||
lines.push("## Framework notes".to_string());
|
||||
lines.extend(framework_lines);
|
||||
lines.push(String::new());
|
||||
}
|
||||
|
||||
lines.push("## Working agreement".to_string());
|
||||
lines.push("- Prefer small, reviewable changes and keep generated bootstrap files aligned with actual repo workflows.".to_string());
|
||||
lines.push("- Keep shared defaults in `.claw.json`; reserve `.claw/settings.local.json` for machine-local overrides.".to_string());
|
||||
lines.push("- Do not overwrite existing `CLAW.md` content automatically; update it intentionally when repo workflows change.".to_string());
|
||||
lines.push(String::new());
|
||||
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
fn detect_repo(cwd: &Path) -> RepoDetection {
|
||||
let package_json_contents = fs::read_to_string(cwd.join("package.json"))
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
RepoDetection {
|
||||
rust_workspace: cwd.join("rust").join("Cargo.toml").is_file(),
|
||||
rust_root: cwd.join("Cargo.toml").is_file(),
|
||||
python: cwd.join("pyproject.toml").is_file()
|
||||
|| cwd.join("requirements.txt").is_file()
|
||||
|| cwd.join("setup.py").is_file(),
|
||||
package_json: cwd.join("package.json").is_file(),
|
||||
typescript: cwd.join("tsconfig.json").is_file()
|
||||
|| package_json_contents.contains("typescript"),
|
||||
nextjs: package_json_contents.contains("\"next\""),
|
||||
react: package_json_contents.contains("\"react\""),
|
||||
vite: package_json_contents.contains("\"vite\""),
|
||||
nest: package_json_contents.contains("@nestjs"),
|
||||
src_dir: cwd.join("src").is_dir(),
|
||||
tests_dir: cwd.join("tests").is_dir(),
|
||||
rust_dir: cwd.join("rust").is_dir(),
|
||||
}
|
||||
}
|
||||
|
||||
fn detected_languages(detection: &RepoDetection) -> Vec<&'static str> {
|
||||
let mut languages = Vec::new();
|
||||
if detection.rust_workspace || detection.rust_root {
|
||||
languages.push("Rust");
|
||||
}
|
||||
if detection.python {
|
||||
languages.push("Python");
|
||||
}
|
||||
if detection.typescript {
|
||||
languages.push("TypeScript");
|
||||
} else if detection.package_json {
|
||||
languages.push("JavaScript/Node.js");
|
||||
}
|
||||
languages
|
||||
}
|
||||
|
||||
fn detected_frameworks(detection: &RepoDetection) -> Vec<&'static str> {
|
||||
let mut frameworks = Vec::new();
|
||||
if detection.nextjs {
|
||||
frameworks.push("Next.js");
|
||||
}
|
||||
if detection.react {
|
||||
frameworks.push("React");
|
||||
}
|
||||
if detection.vite {
|
||||
frameworks.push("Vite");
|
||||
}
|
||||
if detection.nest {
|
||||
frameworks.push("NestJS");
|
||||
}
|
||||
frameworks
|
||||
}
|
||||
|
||||
fn verification_lines(cwd: &Path, detection: &RepoDetection) -> Vec<String> {
|
||||
let mut lines = Vec::new();
|
||||
if detection.rust_workspace {
|
||||
lines.push("- Run Rust verification from `rust/`: `cargo fmt`, `cargo clippy --workspace --all-targets -- -D warnings`, `cargo test --workspace`".to_string());
|
||||
} else if detection.rust_root {
|
||||
lines.push("- Run Rust verification from the repo root: `cargo fmt`, `cargo clippy --workspace --all-targets -- -D warnings`, `cargo test --workspace`".to_string());
|
||||
}
|
||||
if detection.python {
|
||||
if cwd.join("pyproject.toml").is_file() {
|
||||
lines.push("- Run the Python project checks declared in `pyproject.toml` (for example: `pytest`, `ruff check`, and `mypy` when configured).".to_string());
|
||||
} else {
|
||||
lines.push(
|
||||
"- Run the repo's Python test/lint commands before shipping changes.".to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
if detection.package_json {
|
||||
lines.push("- Run the JavaScript/TypeScript checks from `package.json` before shipping changes (`npm test`, `npm run lint`, `npm run build`, or the repo equivalent).".to_string());
|
||||
}
|
||||
if detection.tests_dir && detection.src_dir {
|
||||
lines.push("- `src/` and `tests/` are both present; update both surfaces together when behavior changes.".to_string());
|
||||
}
|
||||
lines
|
||||
}
|
||||
|
||||
fn repository_shape_lines(detection: &RepoDetection) -> Vec<String> {
|
||||
let mut lines = Vec::new();
|
||||
if detection.rust_dir {
|
||||
lines.push(
|
||||
"- `rust/` contains the Rust workspace and active CLI/runtime implementation."
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
if detection.src_dir {
|
||||
lines.push("- `src/` contains source files that should stay consistent with generated guidance and tests.".to_string());
|
||||
}
|
||||
if detection.tests_dir {
|
||||
lines.push("- `tests/` contains validation surfaces that should be reviewed alongside code changes.".to_string());
|
||||
}
|
||||
lines
|
||||
}
|
||||
|
||||
fn framework_notes(detection: &RepoDetection) -> Vec<String> {
|
||||
let mut lines = Vec::new();
|
||||
if detection.nextjs {
|
||||
lines.push("- Next.js detected: preserve routing/data-fetching conventions and verify production builds after changing app structure.".to_string());
|
||||
}
|
||||
if detection.react && !detection.nextjs {
|
||||
lines.push("- React detected: keep component behavior covered with focused tests and avoid unnecessary prop/API churn.".to_string());
|
||||
}
|
||||
if detection.vite {
|
||||
lines.push("- Vite detected: validate the production bundle after changing build-sensitive configuration or imports.".to_string());
|
||||
}
|
||||
if detection.nest {
|
||||
lines.push("- NestJS detected: keep module/provider boundaries explicit and verify controller/service wiring after refactors.".to_string());
|
||||
}
|
||||
lines
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{initialize_repo, render_init_claw_md};
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_dir() -> 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!("claw-init-{nanos}"))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn initialize_repo_creates_expected_files_and_gitignore_entries() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(root.join("rust")).expect("create rust dir");
|
||||
fs::write(root.join("rust").join("Cargo.toml"), "[workspace]\n").expect("write cargo");
|
||||
|
||||
let report = initialize_repo(&root).expect("init should succeed");
|
||||
let rendered = report.render();
|
||||
assert!(rendered.contains(".claw/ created"));
|
||||
assert!(rendered.contains(".claw.json created"));
|
||||
assert!(rendered.contains(".gitignore created"));
|
||||
assert!(rendered.contains("CLAW.md created"));
|
||||
assert!(root.join(".claw").is_dir());
|
||||
assert!(root.join(".claw.json").is_file());
|
||||
assert!(root.join("CLAW.md").is_file());
|
||||
assert_eq!(
|
||||
fs::read_to_string(root.join(".claw.json")).expect("read claw json"),
|
||||
concat!(
|
||||
"{\n",
|
||||
" \"permissions\": {\n",
|
||||
" \"defaultMode\": \"dontAsk\"\n",
|
||||
" }\n",
|
||||
"}\n",
|
||||
)
|
||||
);
|
||||
let gitignore = fs::read_to_string(root.join(".gitignore")).expect("read gitignore");
|
||||
assert!(gitignore.contains(".claw/settings.local.json"));
|
||||
assert!(gitignore.contains(".claw/sessions/"));
|
||||
let claw_md = fs::read_to_string(root.join("CLAW.md")).expect("read claw md");
|
||||
assert!(claw_md.contains("Languages: Rust."));
|
||||
assert!(claw_md.contains("cargo clippy --workspace --all-targets -- -D warnings"));
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn initialize_repo_is_idempotent_and_preserves_existing_files() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create root");
|
||||
fs::write(root.join("CLAW.md"), "custom guidance\n").expect("write existing claw md");
|
||||
fs::write(root.join(".gitignore"), ".claw/settings.local.json\n").expect("write gitignore");
|
||||
|
||||
let first = initialize_repo(&root).expect("first init should succeed");
|
||||
assert!(first
|
||||
.render()
|
||||
.contains("CLAW.md skipped (already exists)"));
|
||||
let second = initialize_repo(&root).expect("second init should succeed");
|
||||
let second_rendered = second.render();
|
||||
assert!(second_rendered.contains(".claw/ skipped (already exists)"));
|
||||
assert!(second_rendered.contains(".claw.json skipped (already exists)"));
|
||||
assert!(second_rendered.contains(".gitignore skipped (already exists)"));
|
||||
assert!(second_rendered.contains("CLAW.md skipped (already exists)"));
|
||||
assert_eq!(
|
||||
fs::read_to_string(root.join("CLAW.md")).expect("read existing claw md"),
|
||||
"custom guidance\n"
|
||||
);
|
||||
let gitignore = fs::read_to_string(root.join(".gitignore")).expect("read gitignore");
|
||||
assert_eq!(gitignore.matches(".claw/settings.local.json").count(), 1);
|
||||
assert_eq!(gitignore.matches(".claw/sessions/").count(), 1);
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_init_template_mentions_detected_python_and_nextjs_markers() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create root");
|
||||
fs::write(root.join("pyproject.toml"), "[project]\nname = \"demo\"\n")
|
||||
.expect("write pyproject");
|
||||
fs::write(
|
||||
root.join("package.json"),
|
||||
r#"{"dependencies":{"next":"14.0.0","react":"18.0.0"},"devDependencies":{"typescript":"5.0.0"}}"#,
|
||||
)
|
||||
.expect("write package json");
|
||||
|
||||
let rendered = render_init_claw_md(Path::new(&root));
|
||||
assert!(rendered.contains("Languages: Python, TypeScript."));
|
||||
assert!(rendered.contains("Frameworks/tooling markers: Next.js, React."));
|
||||
assert!(rendered.contains("pyproject.toml"));
|
||||
assert!(rendered.contains("Next.js detected"));
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,797 @@
|
||||
use std::fmt::Write as FmtWrite;
|
||||
use std::io::{self, Write};
|
||||
|
||||
use crossterm::cursor::{MoveToColumn, RestorePosition, SavePosition};
|
||||
use crossterm::style::{Color, Print, ResetColor, SetForegroundColor, Stylize};
|
||||
use crossterm::terminal::{Clear, ClearType};
|
||||
use crossterm::{execute, queue};
|
||||
use pulldown_cmark::{CodeBlockKind, Event, Options, Parser, Tag, TagEnd};
|
||||
use syntect::easy::HighlightLines;
|
||||
use syntect::highlighting::{Theme, ThemeSet};
|
||||
use syntect::parsing::SyntaxSet;
|
||||
use syntect::util::{as_24_bit_terminal_escaped, LinesWithEndings};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ColorTheme {
|
||||
heading: Color,
|
||||
emphasis: Color,
|
||||
strong: Color,
|
||||
inline_code: Color,
|
||||
link: Color,
|
||||
quote: Color,
|
||||
table_border: Color,
|
||||
code_block_border: Color,
|
||||
spinner_active: Color,
|
||||
spinner_done: Color,
|
||||
spinner_failed: Color,
|
||||
}
|
||||
|
||||
impl Default for ColorTheme {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
heading: Color::Cyan,
|
||||
emphasis: Color::Magenta,
|
||||
strong: Color::Yellow,
|
||||
inline_code: Color::Green,
|
||||
link: Color::Blue,
|
||||
quote: Color::DarkGrey,
|
||||
table_border: Color::DarkCyan,
|
||||
code_block_border: Color::DarkGrey,
|
||||
spinner_active: Color::Blue,
|
||||
spinner_done: Color::Green,
|
||||
spinner_failed: Color::Red,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
pub struct Spinner {
|
||||
frame_index: usize,
|
||||
}
|
||||
|
||||
impl Spinner {
|
||||
const FRAMES: [&str; 10] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
|
||||
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn tick(
|
||||
&mut self,
|
||||
label: &str,
|
||||
theme: &ColorTheme,
|
||||
out: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
let frame = Self::FRAMES[self.frame_index % Self::FRAMES.len()];
|
||||
self.frame_index += 1;
|
||||
queue!(
|
||||
out,
|
||||
SavePosition,
|
||||
MoveToColumn(0),
|
||||
Clear(ClearType::CurrentLine),
|
||||
SetForegroundColor(theme.spinner_active),
|
||||
Print(format!("{frame} {label}")),
|
||||
ResetColor,
|
||||
RestorePosition
|
||||
)?;
|
||||
out.flush()
|
||||
}
|
||||
|
||||
pub fn finish(
|
||||
&mut self,
|
||||
label: &str,
|
||||
theme: &ColorTheme,
|
||||
out: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
self.frame_index = 0;
|
||||
execute!(
|
||||
out,
|
||||
MoveToColumn(0),
|
||||
Clear(ClearType::CurrentLine),
|
||||
SetForegroundColor(theme.spinner_done),
|
||||
Print(format!("✔ {label}\n")),
|
||||
ResetColor
|
||||
)?;
|
||||
out.flush()
|
||||
}
|
||||
|
||||
pub fn fail(
|
||||
&mut self,
|
||||
label: &str,
|
||||
theme: &ColorTheme,
|
||||
out: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
self.frame_index = 0;
|
||||
execute!(
|
||||
out,
|
||||
MoveToColumn(0),
|
||||
Clear(ClearType::CurrentLine),
|
||||
SetForegroundColor(theme.spinner_failed),
|
||||
Print(format!("✘ {label}\n")),
|
||||
ResetColor
|
||||
)?;
|
||||
out.flush()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum ListKind {
|
||||
Unordered,
|
||||
Ordered { next_index: u64 },
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
struct TableState {
|
||||
headers: Vec<String>,
|
||||
rows: Vec<Vec<String>>,
|
||||
current_row: Vec<String>,
|
||||
current_cell: String,
|
||||
in_head: bool,
|
||||
}
|
||||
|
||||
impl TableState {
|
||||
fn push_cell(&mut self) {
|
||||
let cell = self.current_cell.trim().to_string();
|
||||
self.current_row.push(cell);
|
||||
self.current_cell.clear();
|
||||
}
|
||||
|
||||
fn finish_row(&mut self) {
|
||||
if self.current_row.is_empty() {
|
||||
return;
|
||||
}
|
||||
let row = std::mem::take(&mut self.current_row);
|
||||
if self.in_head {
|
||||
self.headers = row;
|
||||
} else {
|
||||
self.rows.push(row);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
struct RenderState {
|
||||
emphasis: usize,
|
||||
strong: usize,
|
||||
heading_level: Option<u8>,
|
||||
quote: usize,
|
||||
list_stack: Vec<ListKind>,
|
||||
link_stack: Vec<LinkState>,
|
||||
table: Option<TableState>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct LinkState {
|
||||
destination: String,
|
||||
text: String,
|
||||
}
|
||||
|
||||
impl RenderState {
|
||||
fn style_text(&self, text: &str, theme: &ColorTheme) -> String {
|
||||
let mut style = text.stylize();
|
||||
|
||||
if matches!(self.heading_level, Some(1 | 2)) || self.strong > 0 {
|
||||
style = style.bold();
|
||||
}
|
||||
if self.emphasis > 0 {
|
||||
style = style.italic();
|
||||
}
|
||||
|
||||
if let Some(level) = self.heading_level {
|
||||
style = match level {
|
||||
1 => style.with(theme.heading),
|
||||
2 => style.white(),
|
||||
3 => style.with(Color::Blue),
|
||||
_ => style.with(Color::Grey),
|
||||
};
|
||||
} else if self.strong > 0 {
|
||||
style = style.with(theme.strong);
|
||||
} else if self.emphasis > 0 {
|
||||
style = style.with(theme.emphasis);
|
||||
}
|
||||
|
||||
if self.quote > 0 {
|
||||
style = style.with(theme.quote);
|
||||
}
|
||||
|
||||
format!("{style}")
|
||||
}
|
||||
|
||||
fn append_raw(&mut self, output: &mut String, text: &str) {
|
||||
if let Some(link) = self.link_stack.last_mut() {
|
||||
link.text.push_str(text);
|
||||
} else if let Some(table) = self.table.as_mut() {
|
||||
table.current_cell.push_str(text);
|
||||
} else {
|
||||
output.push_str(text);
|
||||
}
|
||||
}
|
||||
|
||||
fn append_styled(&mut self, output: &mut String, text: &str, theme: &ColorTheme) {
|
||||
let styled = self.style_text(text, theme);
|
||||
self.append_raw(output, &styled);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TerminalRenderer {
|
||||
syntax_set: SyntaxSet,
|
||||
syntax_theme: Theme,
|
||||
color_theme: ColorTheme,
|
||||
}
|
||||
|
||||
impl Default for TerminalRenderer {
|
||||
fn default() -> Self {
|
||||
let syntax_set = SyntaxSet::load_defaults_newlines();
|
||||
let syntax_theme = ThemeSet::load_defaults()
|
||||
.themes
|
||||
.remove("base16-ocean.dark")
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
syntax_set,
|
||||
syntax_theme,
|
||||
color_theme: ColorTheme::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TerminalRenderer {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn color_theme(&self) -> &ColorTheme {
|
||||
&self.color_theme
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn render_markdown(&self, markdown: &str) -> String {
|
||||
let mut output = String::new();
|
||||
let mut state = RenderState::default();
|
||||
let mut code_language = String::new();
|
||||
let mut code_buffer = String::new();
|
||||
let mut in_code_block = false;
|
||||
|
||||
for event in Parser::new_ext(markdown, Options::all()) {
|
||||
self.render_event(
|
||||
event,
|
||||
&mut state,
|
||||
&mut output,
|
||||
&mut code_buffer,
|
||||
&mut code_language,
|
||||
&mut in_code_block,
|
||||
);
|
||||
}
|
||||
|
||||
output.trim_end().to_string()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn markdown_to_ansi(&self, markdown: &str) -> String {
|
||||
self.render_markdown(markdown)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)]
|
||||
fn render_event(
|
||||
&self,
|
||||
event: Event<'_>,
|
||||
state: &mut RenderState,
|
||||
output: &mut String,
|
||||
code_buffer: &mut String,
|
||||
code_language: &mut String,
|
||||
in_code_block: &mut bool,
|
||||
) {
|
||||
match event {
|
||||
Event::Start(Tag::Heading { level, .. }) => {
|
||||
self.start_heading(state, level as u8, output);
|
||||
}
|
||||
Event::End(TagEnd::Paragraph) => output.push_str("\n\n"),
|
||||
Event::Start(Tag::BlockQuote(..)) => self.start_quote(state, output),
|
||||
Event::End(TagEnd::BlockQuote(..)) => {
|
||||
state.quote = state.quote.saturating_sub(1);
|
||||
output.push('\n');
|
||||
}
|
||||
Event::End(TagEnd::Heading(..)) => {
|
||||
state.heading_level = None;
|
||||
output.push_str("\n\n");
|
||||
}
|
||||
Event::End(TagEnd::Item) | Event::SoftBreak | Event::HardBreak => {
|
||||
state.append_raw(output, "\n");
|
||||
}
|
||||
Event::Start(Tag::List(first_item)) => {
|
||||
let kind = match first_item {
|
||||
Some(index) => ListKind::Ordered { next_index: index },
|
||||
None => ListKind::Unordered,
|
||||
};
|
||||
state.list_stack.push(kind);
|
||||
}
|
||||
Event::End(TagEnd::List(..)) => {
|
||||
state.list_stack.pop();
|
||||
output.push('\n');
|
||||
}
|
||||
Event::Start(Tag::Item) => Self::start_item(state, output),
|
||||
Event::Start(Tag::CodeBlock(kind)) => {
|
||||
*in_code_block = true;
|
||||
*code_language = match kind {
|
||||
CodeBlockKind::Indented => String::from("text"),
|
||||
CodeBlockKind::Fenced(lang) => lang.to_string(),
|
||||
};
|
||||
code_buffer.clear();
|
||||
self.start_code_block(code_language, output);
|
||||
}
|
||||
Event::End(TagEnd::CodeBlock) => {
|
||||
self.finish_code_block(code_buffer, code_language, output);
|
||||
*in_code_block = false;
|
||||
code_language.clear();
|
||||
code_buffer.clear();
|
||||
}
|
||||
Event::Start(Tag::Emphasis) => state.emphasis += 1,
|
||||
Event::End(TagEnd::Emphasis) => state.emphasis = state.emphasis.saturating_sub(1),
|
||||
Event::Start(Tag::Strong) => state.strong += 1,
|
||||
Event::End(TagEnd::Strong) => state.strong = state.strong.saturating_sub(1),
|
||||
Event::Code(code) => {
|
||||
let rendered =
|
||||
format!("{}", format!("`{code}`").with(self.color_theme.inline_code));
|
||||
state.append_raw(output, &rendered);
|
||||
}
|
||||
Event::Rule => output.push_str("---\n"),
|
||||
Event::Text(text) => {
|
||||
self.push_text(text.as_ref(), state, output, code_buffer, *in_code_block);
|
||||
}
|
||||
Event::Html(html) | Event::InlineHtml(html) => {
|
||||
state.append_raw(output, &html);
|
||||
}
|
||||
Event::FootnoteReference(reference) => {
|
||||
state.append_raw(output, &format!("[{reference}]"));
|
||||
}
|
||||
Event::TaskListMarker(done) => {
|
||||
state.append_raw(output, if done { "[x] " } else { "[ ] " });
|
||||
}
|
||||
Event::InlineMath(math) | Event::DisplayMath(math) => {
|
||||
state.append_raw(output, &math);
|
||||
}
|
||||
Event::Start(Tag::Link { dest_url, .. }) => {
|
||||
state.link_stack.push(LinkState {
|
||||
destination: dest_url.to_string(),
|
||||
text: String::new(),
|
||||
});
|
||||
}
|
||||
Event::End(TagEnd::Link) => {
|
||||
if let Some(link) = state.link_stack.pop() {
|
||||
let label = if link.text.is_empty() {
|
||||
link.destination.clone()
|
||||
} else {
|
||||
link.text
|
||||
};
|
||||
let rendered = format!(
|
||||
"{}",
|
||||
format!("[{label}]({})", link.destination)
|
||||
.underlined()
|
||||
.with(self.color_theme.link)
|
||||
);
|
||||
state.append_raw(output, &rendered);
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::Image { dest_url, .. }) => {
|
||||
let rendered = format!(
|
||||
"{}",
|
||||
format!("[image:{dest_url}]").with(self.color_theme.link)
|
||||
);
|
||||
state.append_raw(output, &rendered);
|
||||
}
|
||||
Event::Start(Tag::Table(..)) => state.table = Some(TableState::default()),
|
||||
Event::End(TagEnd::Table) => {
|
||||
if let Some(table) = state.table.take() {
|
||||
output.push_str(&self.render_table(&table));
|
||||
output.push_str("\n\n");
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::TableHead) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.in_head = true;
|
||||
}
|
||||
}
|
||||
Event::End(TagEnd::TableHead) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.finish_row();
|
||||
table.in_head = false;
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::TableRow) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.current_row.clear();
|
||||
table.current_cell.clear();
|
||||
}
|
||||
}
|
||||
Event::End(TagEnd::TableRow) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.finish_row();
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::TableCell) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.current_cell.clear();
|
||||
}
|
||||
}
|
||||
Event::End(TagEnd::TableCell) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.push_cell();
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::Paragraph | Tag::MetadataBlock(..) | _)
|
||||
| Event::End(TagEnd::Image | TagEnd::MetadataBlock(..) | _) => {}
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::unused_self)]
|
||||
fn start_heading(&self, state: &mut RenderState, level: u8, output: &mut String) {
|
||||
state.heading_level = Some(level);
|
||||
if !output.is_empty() {
|
||||
output.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
fn start_quote(&self, state: &mut RenderState, output: &mut String) {
|
||||
state.quote += 1;
|
||||
let _ = write!(output, "{}", "│ ".with(self.color_theme.quote));
|
||||
}
|
||||
|
||||
fn start_item(state: &mut RenderState, output: &mut String) {
|
||||
let depth = state.list_stack.len().saturating_sub(1);
|
||||
output.push_str(&" ".repeat(depth));
|
||||
|
||||
let marker = match state.list_stack.last_mut() {
|
||||
Some(ListKind::Ordered { next_index }) => {
|
||||
let value = *next_index;
|
||||
*next_index += 1;
|
||||
format!("{value}. ")
|
||||
}
|
||||
_ => "• ".to_string(),
|
||||
};
|
||||
output.push_str(&marker);
|
||||
}
|
||||
|
||||
fn start_code_block(&self, code_language: &str, output: &mut String) {
|
||||
let label = if code_language.is_empty() {
|
||||
"code".to_string()
|
||||
} else {
|
||||
code_language.to_string()
|
||||
};
|
||||
let _ = writeln!(
|
||||
output,
|
||||
"{}",
|
||||
format!("╭─ {label}")
|
||||
.bold()
|
||||
.with(self.color_theme.code_block_border)
|
||||
);
|
||||
}
|
||||
|
||||
fn finish_code_block(&self, code_buffer: &str, code_language: &str, output: &mut String) {
|
||||
output.push_str(&self.highlight_code(code_buffer, code_language));
|
||||
let _ = write!(
|
||||
output,
|
||||
"{}",
|
||||
"╰─".bold().with(self.color_theme.code_block_border)
|
||||
);
|
||||
output.push_str("\n\n");
|
||||
}
|
||||
|
||||
fn push_text(
|
||||
&self,
|
||||
text: &str,
|
||||
state: &mut RenderState,
|
||||
output: &mut String,
|
||||
code_buffer: &mut String,
|
||||
in_code_block: bool,
|
||||
) {
|
||||
if in_code_block {
|
||||
code_buffer.push_str(text);
|
||||
} else {
|
||||
state.append_styled(output, text, &self.color_theme);
|
||||
}
|
||||
}
|
||||
|
||||
fn render_table(&self, table: &TableState) -> String {
|
||||
let mut rows = Vec::new();
|
||||
if !table.headers.is_empty() {
|
||||
rows.push(table.headers.clone());
|
||||
}
|
||||
rows.extend(table.rows.iter().cloned());
|
||||
|
||||
if rows.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
|
||||
let column_count = rows.iter().map(Vec::len).max().unwrap_or(0);
|
||||
let widths = (0..column_count)
|
||||
.map(|column| {
|
||||
rows.iter()
|
||||
.filter_map(|row| row.get(column))
|
||||
.map(|cell| visible_width(cell))
|
||||
.max()
|
||||
.unwrap_or(0)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let border = format!("{}", "│".with(self.color_theme.table_border));
|
||||
let separator = widths
|
||||
.iter()
|
||||
.map(|width| "─".repeat(*width + 2))
|
||||
.collect::<Vec<_>>()
|
||||
.join(&format!("{}", "┼".with(self.color_theme.table_border)));
|
||||
let separator = format!("{border}{separator}{border}");
|
||||
|
||||
let mut output = String::new();
|
||||
if !table.headers.is_empty() {
|
||||
output.push_str(&self.render_table_row(&table.headers, &widths, true));
|
||||
output.push('\n');
|
||||
output.push_str(&separator);
|
||||
if !table.rows.is_empty() {
|
||||
output.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
for (index, row) in table.rows.iter().enumerate() {
|
||||
output.push_str(&self.render_table_row(row, &widths, false));
|
||||
if index + 1 < table.rows.len() {
|
||||
output.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
fn render_table_row(&self, row: &[String], widths: &[usize], is_header: bool) -> String {
|
||||
let border = format!("{}", "│".with(self.color_theme.table_border));
|
||||
let mut line = String::new();
|
||||
line.push_str(&border);
|
||||
|
||||
for (index, width) in widths.iter().enumerate() {
|
||||
let cell = row.get(index).map_or("", String::as_str);
|
||||
line.push(' ');
|
||||
if is_header {
|
||||
let _ = write!(line, "{}", cell.bold().with(self.color_theme.heading));
|
||||
} else {
|
||||
line.push_str(cell);
|
||||
}
|
||||
let padding = width.saturating_sub(visible_width(cell));
|
||||
line.push_str(&" ".repeat(padding + 1));
|
||||
line.push_str(&border);
|
||||
}
|
||||
|
||||
line
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn highlight_code(&self, code: &str, language: &str) -> String {
|
||||
let syntax = self
|
||||
.syntax_set
|
||||
.find_syntax_by_token(language)
|
||||
.unwrap_or_else(|| self.syntax_set.find_syntax_plain_text());
|
||||
let mut syntax_highlighter = HighlightLines::new(syntax, &self.syntax_theme);
|
||||
let mut colored_output = String::new();
|
||||
|
||||
for line in LinesWithEndings::from(code) {
|
||||
match syntax_highlighter.highlight_line(line, &self.syntax_set) {
|
||||
Ok(ranges) => {
|
||||
let escaped = as_24_bit_terminal_escaped(&ranges[..], false);
|
||||
colored_output.push_str(&apply_code_block_background(&escaped));
|
||||
}
|
||||
Err(_) => colored_output.push_str(&apply_code_block_background(line)),
|
||||
}
|
||||
}
|
||||
|
||||
colored_output
|
||||
}
|
||||
|
||||
pub fn stream_markdown(&self, markdown: &str, out: &mut impl Write) -> io::Result<()> {
|
||||
let rendered_markdown = self.markdown_to_ansi(markdown);
|
||||
write!(out, "{rendered_markdown}")?;
|
||||
if !rendered_markdown.ends_with('\n') {
|
||||
writeln!(out)?;
|
||||
}
|
||||
out.flush()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
pub struct MarkdownStreamState {
|
||||
pending: String,
|
||||
}
|
||||
|
||||
impl MarkdownStreamState {
|
||||
#[must_use]
|
||||
pub fn push(&mut self, renderer: &TerminalRenderer, delta: &str) -> Option<String> {
|
||||
self.pending.push_str(delta);
|
||||
let split = find_stream_safe_boundary(&self.pending)?;
|
||||
let ready = self.pending[..split].to_string();
|
||||
self.pending.drain(..split);
|
||||
Some(renderer.markdown_to_ansi(&ready))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn flush(&mut self, renderer: &TerminalRenderer) -> Option<String> {
|
||||
if self.pending.trim().is_empty() {
|
||||
self.pending.clear();
|
||||
None
|
||||
} else {
|
||||
let pending = std::mem::take(&mut self.pending);
|
||||
Some(renderer.markdown_to_ansi(&pending))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_code_block_background(line: &str) -> String {
|
||||
let trimmed = line.trim_end_matches('\n');
|
||||
let trailing_newline = if trimmed.len() == line.len() {
|
||||
""
|
||||
} else {
|
||||
"\n"
|
||||
};
|
||||
let with_background = trimmed.replace("\u{1b}[0m", "\u{1b}[0;48;5;236m");
|
||||
format!("\u{1b}[48;5;236m{with_background}\u{1b}[0m{trailing_newline}")
|
||||
}
|
||||
|
||||
fn find_stream_safe_boundary(markdown: &str) -> Option<usize> {
|
||||
let mut in_fence = false;
|
||||
let mut last_boundary = None;
|
||||
|
||||
for (offset, line) in markdown.split_inclusive('\n').scan(0usize, |cursor, line| {
|
||||
let start = *cursor;
|
||||
*cursor += line.len();
|
||||
Some((start, line))
|
||||
}) {
|
||||
let trimmed = line.trim_start();
|
||||
if trimmed.starts_with("```") || trimmed.starts_with("~~~") {
|
||||
in_fence = !in_fence;
|
||||
if !in_fence {
|
||||
last_boundary = Some(offset + line.len());
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if in_fence {
|
||||
continue;
|
||||
}
|
||||
|
||||
if trimmed.is_empty() {
|
||||
last_boundary = Some(offset + line.len());
|
||||
}
|
||||
}
|
||||
|
||||
last_boundary
|
||||
}
|
||||
|
||||
fn visible_width(input: &str) -> usize {
|
||||
strip_ansi(input).chars().count()
|
||||
}
|
||||
|
||||
fn strip_ansi(input: &str) -> String {
|
||||
let mut output = String::new();
|
||||
let mut chars = input.chars().peekable();
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
if ch == '\u{1b}' {
|
||||
if chars.peek() == Some(&'[') {
|
||||
chars.next();
|
||||
for next in chars.by_ref() {
|
||||
if next.is_ascii_alphabetic() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
output.push(ch);
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{strip_ansi, MarkdownStreamState, Spinner, TerminalRenderer};
|
||||
|
||||
#[test]
|
||||
fn renders_markdown_with_styling_and_lists() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output = terminal_renderer
|
||||
.render_markdown("# Heading\n\nThis is **bold** and *italic*.\n\n- item\n\n`code`");
|
||||
|
||||
assert!(markdown_output.contains("Heading"));
|
||||
assert!(markdown_output.contains("• item"));
|
||||
assert!(markdown_output.contains("code"));
|
||||
assert!(markdown_output.contains('\u{1b}'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_links_as_colored_markdown_labels() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output =
|
||||
terminal_renderer.render_markdown("See [Claw](https://example.com/docs) now.");
|
||||
let plain_text = strip_ansi(&markdown_output);
|
||||
|
||||
assert!(plain_text.contains("[Claw](https://example.com/docs)"));
|
||||
assert!(markdown_output.contains('\u{1b}'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn highlights_fenced_code_blocks() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output =
|
||||
terminal_renderer.markdown_to_ansi("```rust\nfn hi() { println!(\"hi\"); }\n```");
|
||||
let plain_text = strip_ansi(&markdown_output);
|
||||
|
||||
assert!(plain_text.contains("╭─ rust"));
|
||||
assert!(plain_text.contains("fn hi"));
|
||||
assert!(markdown_output.contains('\u{1b}'));
|
||||
assert!(markdown_output.contains("[48;5;236m"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_ordered_and_nested_lists() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output =
|
||||
terminal_renderer.render_markdown("1. first\n2. second\n - nested\n - child");
|
||||
let plain_text = strip_ansi(&markdown_output);
|
||||
|
||||
assert!(plain_text.contains("1. first"));
|
||||
assert!(plain_text.contains("2. second"));
|
||||
assert!(plain_text.contains(" • nested"));
|
||||
assert!(plain_text.contains(" • child"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_tables_with_alignment() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output = terminal_renderer
|
||||
.render_markdown("| Name | Value |\n| ---- | ----- |\n| alpha | 1 |\n| beta | 22 |");
|
||||
let plain_text = strip_ansi(&markdown_output);
|
||||
let lines = plain_text.lines().collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(lines[0], "│ Name │ Value │");
|
||||
assert_eq!(lines[1], "│───────┼───────│");
|
||||
assert_eq!(lines[2], "│ alpha │ 1 │");
|
||||
assert_eq!(lines[3], "│ beta │ 22 │");
|
||||
assert!(markdown_output.contains('\u{1b}'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streaming_state_waits_for_complete_blocks() {
|
||||
let renderer = TerminalRenderer::new();
|
||||
let mut state = MarkdownStreamState::default();
|
||||
|
||||
assert_eq!(state.push(&renderer, "# Heading"), None);
|
||||
let flushed = state
|
||||
.push(&renderer, "\n\nParagraph\n\n")
|
||||
.expect("completed block");
|
||||
let plain_text = strip_ansi(&flushed);
|
||||
assert!(plain_text.contains("Heading"));
|
||||
assert!(plain_text.contains("Paragraph"));
|
||||
|
||||
assert_eq!(state.push(&renderer, "```rust\nfn main() {}\n"), None);
|
||||
let code = state
|
||||
.push(&renderer, "```\n")
|
||||
.expect("closed code fence flushes");
|
||||
assert!(strip_ansi(&code).contains("fn main()"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn spinner_advances_frames() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let mut spinner = Spinner::new();
|
||||
let mut out = Vec::new();
|
||||
spinner
|
||||
.tick("Working", terminal_renderer.color_theme(), &mut out)
|
||||
.expect("tick succeeds");
|
||||
spinner
|
||||
.tick("Working", terminal_renderer.color_theme(), &mut out)
|
||||
.expect("tick succeeds");
|
||||
|
||||
let output = String::from_utf8_lossy(&out);
|
||||
assert!(output.contains("Working"));
|
||||
}
|
||||
}
|
||||
@@ -9,4 +9,6 @@ publish.workspace = true
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
plugins = { path = "../plugins" }
|
||||
runtime = { path = "../runtime" }
|
||||
serde_json.workspace = true
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
# 命令模块 (commands)
|
||||
|
||||
本模块负责定义和管理 Claw 交互界面中使用的“斜杠命令”(Slash Commands),并提供相关的解析和执行逻辑。
|
||||
|
||||
## 概览
|
||||
|
||||
`commands` 模块的主要职责包括:
|
||||
- 定义所有可用的斜杠命令及其元数据(别名、说明、类别等)。
|
||||
- 提供命令注册表 (`CommandRegistry`),用于在 CLI 中发现和分发命令。
|
||||
- 实现复杂的管理命令,如插件管理 (`/plugins`)、智能体查看 (`/agents`) 和技能查看 (`/skills`)。
|
||||
- 提供命令建议功能,支持基于编辑距离 (Levenshtein distance) 的模糊匹配。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **斜杠命令规范 (SlashCommandSpec)**:每个命令都包含详尽的元数据,包括所属类别(核心、工作区、会话、Git、自动化)以及是否支持在恢复会话时执行。
|
||||
- **命令分类**:
|
||||
- **核心 (Core)**:`/help`, `/status`, `/model`, `/permissions`, `/cost` 等。
|
||||
- **工作区 (Workspace)**:`/config`, `/memory`, `/diff`, `/teleport` 等。
|
||||
- **会话 (Session)**:`/clear`, `/resume`, `/export`, `/session` 等。
|
||||
- **Git 交互**:`/branch`, `/commit`, `/pr`, `/issue` 等。
|
||||
- **自动化 (Automation)**:`/plugins`, `/agents`, `/skills`, `/ultraplan` 等。
|
||||
- **模糊匹配与建议**:当用户输入错误的命令时,系统会自动推荐最接近的合法命令。
|
||||
- **插件集成**:`/plugins` 命令允许用户动态安装、启用、禁用或卸载插件,并能通知运行时重新加载环境。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`lib.rs`**: 包含了绝大部分逻辑。
|
||||
- **`SlashCommand` 枚举**: 定义了所有命令的强类型表示。
|
||||
- **`SlashCommandSpec` 结构体**: 存储命令的静态配置信息。
|
||||
- **`handle_plugins_slash_command`**: 处理复杂的插件管理工作流。
|
||||
- **`suggest_slash_commands`**: 实现基于 Levenshtein 距离的建议算法。
|
||||
|
||||
### 工作流程
|
||||
|
||||
1. 用户在 REPL 中输入以 `/` 开头的字符串。
|
||||
2. `claw-cli` 调用 `SlashCommand::parse` 进行解析。
|
||||
3. 解析后的命令被分发到相应的处理器。
|
||||
4. 处理结果(通常包含要显示给用户的消息,以及可选的会话更新或运行时重新加载请求)返回给 CLI。
|
||||
|
||||
## 使用示例 (内部)
|
||||
|
||||
```rust
|
||||
use commands::{SlashCommand, suggest_slash_commands};
|
||||
|
||||
// 解析命令
|
||||
if let Some(cmd) = SlashCommand::parse("/model sonnet") {
|
||||
// 处理模型切换逻辑
|
||||
}
|
||||
|
||||
// 获取建议
|
||||
let suggestions = suggest_slash_commands("hpel", 3);
|
||||
// 返回 ["/help"]
|
||||
```
|
||||
+5174
-134
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,48 @@
|
||||
# 兼容性测试套件模块 (compat-harness)
|
||||
|
||||
本模块提供了一套工具,专门用于分析和提取上游引用实现(如原始的 `claude-code` TypeScript 源码)中的元数据,以确保 Rust 版本的实现与其保持功能兼容。
|
||||
|
||||
## 概览
|
||||
|
||||
`compat-harness` 的主要职责是:
|
||||
- 定位上游仓库的源码路径。
|
||||
- 从 TypeScript 源码文件中提取命令 (`commands`)、工具 (`tools`) 和启动阶段 (`bootstrap phases`) 的定义。
|
||||
- 自动生成功能清单 (`ExtractedManifest`),供运行时或测试使用,以验证 Rust 版本的覆盖率。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **上游路径解析 (UpstreamPaths)**:能够自动识别多种常见的上游仓库目录结构,并支持通过环境变量 `CLAUDE_CODE_UPSTREAM` 进行覆盖。
|
||||
- **静态代码分析**:通过解析 TypeScript 源码,识别特定的代码模式(如 `export const INTERNAL_ONLY_COMMANDS` 或基于 `feature()` 的功能开关)。
|
||||
- **清单提取 (Manifest Extraction)**:
|
||||
- **命令提取**:识别内置命令、仅限内部使用的命令以及受功能开关控制的命令。
|
||||
- **工具提取**:识别基础工具和条件加载的工具。
|
||||
- **启动计划提取**:分析 CLI 入口文件,重建启动时的各个阶段(如 `FastPathVersion`, `MainRuntime` 等)。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`lib.rs`**: 包含了核心的提取逻辑。
|
||||
- **`UpstreamPaths` 结构体**: 封装了寻找 `commands.ts`、`tools.ts` 和 `cli.tsx` 的逻辑。
|
||||
- **`extract_commands` & `extract_tools`**: 使用字符串解析技术,识别 TypeScript 的 `import` 和赋值操作,提取符号名称。
|
||||
- **`extract_bootstrap_plan`**: 搜索特定的标志性字符串(如 `--version` 或 `daemon-worker`),从而推断出上游程序的启动流程。
|
||||
|
||||
### 工作流程
|
||||
|
||||
1. 模块根据预设路径或环境变量寻找上游 `claude-code` 仓库。
|
||||
2. 读取关键的 `.ts` 或 `.tsx` 文件内容。
|
||||
3. 执行正则表达式风格的行解析,提取出所有定义的命令和工具名称。
|
||||
4. 将提取结果组织成 `ExtractedManifest` 对象。
|
||||
|
||||
## 使用示例 (内部测试)
|
||||
|
||||
```rust
|
||||
use compat_harness::{UpstreamPaths, extract_manifest};
|
||||
|
||||
// 指定工作区目录,自动寻找上游路径
|
||||
let paths = UpstreamPaths::from_workspace_dir("path/to/workspace");
|
||||
// 提取功能清单
|
||||
if let Ok(manifest) = extract_manifest(&paths) {
|
||||
println!("上游发现 {} 个工具", manifest.tools.entries().len());
|
||||
}
|
||||
```
|
||||
@@ -70,16 +70,12 @@ fn upstream_repo_candidates(primary_repo_root: &Path) -> Vec<PathBuf> {
|
||||
}
|
||||
|
||||
for ancestor in primary_repo_root.ancestors().take(4) {
|
||||
candidates.push(ancestor.join("claude-code"));
|
||||
candidates.push(ancestor.join("claw-code"));
|
||||
candidates.push(ancestor.join("clawd-code"));
|
||||
}
|
||||
|
||||
candidates.push(
|
||||
primary_repo_root
|
||||
.join("reference-source")
|
||||
.join("claude-code"),
|
||||
);
|
||||
candidates.push(primary_repo_root.join("vendor").join("claude-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 {
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
name = "lsp"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[dependencies]
|
||||
lsp-types.workspace = true
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json.workspace = true
|
||||
tokio = { version = "1", features = ["io-util", "macros", "process", "rt", "rt-multi-thread", "sync", "time"] }
|
||||
url = "2"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,56 @@
|
||||
# LSP 模块 (lsp)
|
||||
|
||||
本模块实现了语言服务协议 (Language Server Protocol, LSP) 的客户端功能,允许系统通过集成的编程语言服务器获取代码的语义信息、错误诊断和符号导航。
|
||||
|
||||
## 概览
|
||||
|
||||
`lsp` 模块的主要职责是:
|
||||
- 管理多个 LSP 服务器的生命周期(启动、初始化、关闭)。
|
||||
- 与服务器进行异步 JSON-RPC 通信。
|
||||
- 提供跨语言的代码智能功能,如:
|
||||
- **转到定义 (Go to Definition)**
|
||||
- **查找引用 (Find References)**
|
||||
- **工作区诊断 (Workspace Diagnostics)**
|
||||
- 为 AI 提示词 (Prompt) 提供上下文增强,将代码中的实时错误和符号关系反馈给 LLM。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **LspManager**: 核心管理类,负责协调不同语言的服务器配置和文档状态。
|
||||
- **上下文增强 (Context Enrichment)**:定义了 `LspContextEnrichment` 结构,能够将复杂的 LSP 响应(如诊断信息和定义)转换为易于 AI 理解的 Markdown 格式。
|
||||
- **多服务器支持**:支持根据文件扩展名将请求路由到不同的语言服务器(如 `rust-analyzer`, `pyright` 等)。
|
||||
- **同步机制**:处理文档的 `didOpen`、`didChange` 和 `didSave` 消息,确保服务器拥有最新的代码视图。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`manager.rs`**: 实现了 `LspManager`。它维护一个服务器池,并提供高层 API 来执行跨服务器的请求。
|
||||
- **`client.rs`**: 实现底层的 LSP 客户端逻辑,处理基于 `tokio` 的异步 I/O 和 JSON-RPC 消息的分帧与解析。
|
||||
- **`types.rs`**: 定义了本模块使用的专用数据类型,并对 `lsp-types` 库中的类型进行了简化和包装,以便于内部使用。
|
||||
- **`error.rs`**: 定义了 LSP 相关的错误处理。
|
||||
|
||||
### 工作流程
|
||||
|
||||
1. 系统根据配置初始化 `LspManager`。
|
||||
2. 当打开一个文件时,`LspManager` 启动相应的服务器并发送 `initialize` 请求。
|
||||
3. `LspManager` 跟踪文档的打开状态,并在内容变化时同步到服务器。
|
||||
4. 当需要对某个符号进行分析时,调用 `go_to_definition` 等方法,模块负责发送请求并解析返回的 `Location`。
|
||||
5. 诊断信息异步通过 `textDocument/publishDiagnostics` 通知到达,模块会缓存这些信息供后续查询。
|
||||
|
||||
## 使用示例 (内部)
|
||||
|
||||
```rust
|
||||
use lsp::{LspManager, LspServerConfig};
|
||||
|
||||
// 配置并初始化管理器
|
||||
let configs = vec![LspServerConfig {
|
||||
name: "rust-analyzer".to_string(),
|
||||
command: "rust-analyzer".to_string(),
|
||||
..Default::default()
|
||||
}];
|
||||
let manager = LspManager::new(configs)?;
|
||||
|
||||
// 获取某个位置的上下文增强信息
|
||||
let enrichment = manager.context_enrichment(&file_path, position).await?;
|
||||
println!("{}", enrichment.render_prompt_section());
|
||||
```
|
||||
@@ -0,0 +1,465 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicI64, Ordering};
|
||||
|
||||
use lsp_types::{
|
||||
Diagnostic, GotoDefinitionResponse, Location, LocationLink, Position, PublishDiagnosticsParams,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
|
||||
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
|
||||
use tokio::sync::{oneshot, Mutex};
|
||||
|
||||
use crate::error::LspError;
|
||||
use crate::types::{LspServerConfig, SymbolLocation};
|
||||
|
||||
type PendingRequestMap = BTreeMap<i64, oneshot::Sender<Result<Value, LspError>>>;
|
||||
|
||||
pub(crate) struct LspClient {
|
||||
config: LspServerConfig,
|
||||
writer: Mutex<BufWriter<ChildStdin>>,
|
||||
child: Mutex<Child>,
|
||||
pending_requests: Arc<Mutex<PendingRequestMap>>,
|
||||
diagnostics: Arc<Mutex<BTreeMap<String, Vec<Diagnostic>>>>,
|
||||
open_documents: Mutex<BTreeMap<PathBuf, i32>>,
|
||||
next_request_id: AtomicI64,
|
||||
}
|
||||
|
||||
impl LspClient {
|
||||
pub(crate) async fn connect(config: LspServerConfig) -> Result<Self, LspError> {
|
||||
let mut command = Command::new(&config.command);
|
||||
command
|
||||
.args(&config.args)
|
||||
.current_dir(&config.workspace_root)
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.envs(config.env.clone());
|
||||
|
||||
let mut child = command.spawn()?;
|
||||
let stdin = child
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| LspError::Protocol("missing LSP stdin pipe".to_string()))?;
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| LspError::Protocol("missing LSP stdout pipe".to_string()))?;
|
||||
let stderr = child.stderr.take();
|
||||
|
||||
let client = Self {
|
||||
config,
|
||||
writer: Mutex::new(BufWriter::new(stdin)),
|
||||
child: Mutex::new(child),
|
||||
pending_requests: Arc::new(Mutex::new(BTreeMap::new())),
|
||||
diagnostics: Arc::new(Mutex::new(BTreeMap::new())),
|
||||
open_documents: Mutex::new(BTreeMap::new()),
|
||||
next_request_id: AtomicI64::new(1),
|
||||
};
|
||||
|
||||
client.spawn_reader(stdout);
|
||||
if let Some(stderr) = stderr {
|
||||
Self::spawn_stderr_drain(stderr);
|
||||
}
|
||||
client.initialize().await?;
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
pub(crate) async fn ensure_document_open(&self, path: &Path) -> Result<(), LspError> {
|
||||
if self.is_document_open(path).await {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let contents = std::fs::read_to_string(path)?;
|
||||
self.open_document(path, &contents).await
|
||||
}
|
||||
|
||||
pub(crate) async fn open_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
|
||||
let uri = file_url(path)?;
|
||||
let language_id = self
|
||||
.config
|
||||
.language_id_for(path)
|
||||
.ok_or_else(|| LspError::UnsupportedDocument(path.to_path_buf()))?;
|
||||
|
||||
self.notify(
|
||||
"textDocument/didOpen",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": uri,
|
||||
"languageId": language_id,
|
||||
"version": 1,
|
||||
"text": text,
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
self.open_documents
|
||||
.lock()
|
||||
.await
|
||||
.insert(path.to_path_buf(), 1);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn change_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
|
||||
if !self.is_document_open(path).await {
|
||||
return self.open_document(path, text).await;
|
||||
}
|
||||
|
||||
let uri = file_url(path)?;
|
||||
let next_version = {
|
||||
let mut open_documents = self.open_documents.lock().await;
|
||||
let version = open_documents
|
||||
.entry(path.to_path_buf())
|
||||
.and_modify(|value| *value += 1)
|
||||
.or_insert(1);
|
||||
*version
|
||||
};
|
||||
|
||||
self.notify(
|
||||
"textDocument/didChange",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": uri,
|
||||
"version": next_version,
|
||||
},
|
||||
"contentChanges": [{
|
||||
"text": text,
|
||||
}],
|
||||
}),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn save_document(&self, path: &Path) -> Result<(), LspError> {
|
||||
if !self.is_document_open(path).await {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
self.notify(
|
||||
"textDocument/didSave",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": file_url(path)?,
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn close_document(&self, path: &Path) -> Result<(), LspError> {
|
||||
if !self.is_document_open(path).await {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
self.notify(
|
||||
"textDocument/didClose",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": file_url(path)?,
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
self.open_documents.lock().await.remove(path);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn is_document_open(&self, path: &Path) -> bool {
|
||||
self.open_documents.lock().await.contains_key(path)
|
||||
}
|
||||
|
||||
pub(crate) async fn go_to_definition(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
) -> Result<Vec<SymbolLocation>, LspError> {
|
||||
self.ensure_document_open(path).await?;
|
||||
let response = self
|
||||
.request::<Option<GotoDefinitionResponse>>(
|
||||
"textDocument/definition",
|
||||
json!({
|
||||
"textDocument": { "uri": file_url(path)? },
|
||||
"position": position,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(match response {
|
||||
Some(GotoDefinitionResponse::Scalar(location)) => {
|
||||
location_to_symbol_locations(vec![location])
|
||||
}
|
||||
Some(GotoDefinitionResponse::Array(locations)) => location_to_symbol_locations(locations),
|
||||
Some(GotoDefinitionResponse::Link(links)) => location_links_to_symbol_locations(links),
|
||||
None => Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn find_references(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
include_declaration: bool,
|
||||
) -> Result<Vec<SymbolLocation>, LspError> {
|
||||
self.ensure_document_open(path).await?;
|
||||
let response = self
|
||||
.request::<Option<Vec<Location>>>(
|
||||
"textDocument/references",
|
||||
json!({
|
||||
"textDocument": { "uri": file_url(path)? },
|
||||
"position": position,
|
||||
"context": {
|
||||
"includeDeclaration": include_declaration,
|
||||
},
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(location_to_symbol_locations(response.unwrap_or_default()))
|
||||
}
|
||||
|
||||
pub(crate) async fn diagnostics_snapshot(&self) -> BTreeMap<String, Vec<Diagnostic>> {
|
||||
self.diagnostics.lock().await.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn shutdown(&self) -> Result<(), LspError> {
|
||||
let _ = self.request::<Value>("shutdown", json!({})).await;
|
||||
let _ = self.notify("exit", Value::Null).await;
|
||||
|
||||
let mut child = self.child.lock().await;
|
||||
if child.kill().await.is_err() {
|
||||
let _ = child.wait().await;
|
||||
return Ok(());
|
||||
}
|
||||
let _ = child.wait().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn spawn_reader(&self, stdout: ChildStdout) {
|
||||
let diagnostics = &self.diagnostics;
|
||||
let pending_requests = &self.pending_requests;
|
||||
|
||||
let diagnostics = diagnostics.clone();
|
||||
let pending_requests = pending_requests.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut reader = BufReader::new(stdout);
|
||||
let result = async {
|
||||
while let Some(message) = read_message(&mut reader).await? {
|
||||
if let Some(id) = message.get("id").and_then(Value::as_i64) {
|
||||
let response = if let Some(error) = message.get("error") {
|
||||
Err(LspError::Protocol(error.to_string()))
|
||||
} else {
|
||||
Ok(message.get("result").cloned().unwrap_or(Value::Null))
|
||||
};
|
||||
|
||||
if let Some(sender) = pending_requests.lock().await.remove(&id) {
|
||||
let _ = sender.send(response);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(method) = message.get("method").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
if method != "textDocument/publishDiagnostics" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let params = message.get("params").cloned().unwrap_or(Value::Null);
|
||||
let notification = serde_json::from_value::<PublishDiagnosticsParams>(params)?;
|
||||
let mut diagnostics_map = diagnostics.lock().await;
|
||||
if notification.diagnostics.is_empty() {
|
||||
diagnostics_map.remove(¬ification.uri.to_string());
|
||||
} else {
|
||||
diagnostics_map.insert(notification.uri.to_string(), notification.diagnostics);
|
||||
}
|
||||
}
|
||||
Ok::<(), LspError>(())
|
||||
}
|
||||
.await;
|
||||
|
||||
if let Err(error) = result {
|
||||
let mut pending = pending_requests.lock().await;
|
||||
let drained = pending
|
||||
.keys()
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
for id in drained {
|
||||
if let Some(sender) = pending.remove(&id) {
|
||||
let _ = sender.send(Err(LspError::Protocol(error.to_string())));
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn spawn_stderr_drain<R>(stderr: R)
|
||||
where
|
||||
R: AsyncRead + Unpin + Send + 'static,
|
||||
{
|
||||
tokio::spawn(async move {
|
||||
let mut reader = BufReader::new(stderr);
|
||||
let mut sink = Vec::new();
|
||||
let _ = reader.read_to_end(&mut sink).await;
|
||||
});
|
||||
}
|
||||
|
||||
async fn initialize(&self) -> Result<(), LspError> {
|
||||
let workspace_uri = file_url(&self.config.workspace_root)?;
|
||||
let _ = self
|
||||
.request::<Value>(
|
||||
"initialize",
|
||||
json!({
|
||||
"processId": std::process::id(),
|
||||
"rootUri": workspace_uri,
|
||||
"rootPath": self.config.workspace_root,
|
||||
"workspaceFolders": [{
|
||||
"uri": workspace_uri,
|
||||
"name": self.config.name,
|
||||
}],
|
||||
"initializationOptions": self.config.initialization_options.clone().unwrap_or(Value::Null),
|
||||
"capabilities": {
|
||||
"textDocument": {
|
||||
"publishDiagnostics": {
|
||||
"relatedInformation": true,
|
||||
},
|
||||
"definition": {
|
||||
"linkSupport": true,
|
||||
},
|
||||
"references": {}
|
||||
},
|
||||
"workspace": {
|
||||
"configuration": false,
|
||||
"workspaceFolders": true,
|
||||
},
|
||||
"general": {
|
||||
"positionEncodings": ["utf-16"],
|
||||
}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
self.notify("initialized", json!({})).await
|
||||
}
|
||||
|
||||
async fn request<T>(&self, method: &str, params: Value) -> Result<T, LspError>
|
||||
where
|
||||
T: for<'de> serde::Deserialize<'de>,
|
||||
{
|
||||
let id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
|
||||
let (sender, receiver) = oneshot::channel();
|
||||
self.pending_requests.lock().await.insert(id, sender);
|
||||
|
||||
if let Err(error) = self
|
||||
.send_message(&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"method": method,
|
||||
"params": params,
|
||||
}))
|
||||
.await
|
||||
{
|
||||
self.pending_requests.lock().await.remove(&id);
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
let response = receiver
|
||||
.await
|
||||
.map_err(|_| LspError::Protocol(format!("request channel closed for {method}")))??;
|
||||
Ok(serde_json::from_value(response)?)
|
||||
}
|
||||
|
||||
async fn notify(&self, method: &str, params: Value) -> Result<(), LspError> {
|
||||
self.send_message(&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": method,
|
||||
"params": params,
|
||||
}))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_message(&self, payload: &Value) -> Result<(), LspError> {
|
||||
let body = serde_json::to_vec(payload)?;
|
||||
let mut writer = self.writer.lock().await;
|
||||
writer
|
||||
.write_all(format!("Content-Length: {}\r\n\r\n", body.len()).as_bytes())
|
||||
.await?;
|
||||
writer.write_all(&body).await?;
|
||||
writer.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_message<R>(reader: &mut BufReader<R>) -> Result<Option<Value>, LspError>
|
||||
where
|
||||
R: AsyncRead + Unpin,
|
||||
{
|
||||
let mut content_length = None;
|
||||
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
let read = reader.read_line(&mut line).await?;
|
||||
if read == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
|
||||
let trimmed = line.trim_end_matches(['\r', '\n']);
|
||||
if let Some((name, value)) = trimmed.split_once(':') {
|
||||
if name.eq_ignore_ascii_case("Content-Length") {
|
||||
let value = value.trim().to_string();
|
||||
content_length = Some(
|
||||
value
|
||||
.parse::<usize>()
|
||||
.map_err(|_| LspError::InvalidContentLength(value.clone()))?,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
return Err(LspError::InvalidHeader(trimmed.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
let content_length = content_length.ok_or(LspError::MissingContentLength)?;
|
||||
let mut body = vec![0_u8; content_length];
|
||||
reader.read_exact(&mut body).await?;
|
||||
Ok(Some(serde_json::from_slice(&body)?))
|
||||
}
|
||||
|
||||
fn file_url(path: &Path) -> Result<String, LspError> {
|
||||
url::Url::from_file_path(path)
|
||||
.map(|url| url.to_string())
|
||||
.map_err(|()| LspError::PathToUrl(path.to_path_buf()))
|
||||
}
|
||||
|
||||
fn location_to_symbol_locations(locations: Vec<Location>) -> Vec<SymbolLocation> {
|
||||
locations
|
||||
.into_iter()
|
||||
.filter_map(|location| {
|
||||
uri_to_path(&location.uri.to_string()).map(|path| SymbolLocation {
|
||||
path,
|
||||
range: location.range,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn location_links_to_symbol_locations(links: Vec<LocationLink>) -> Vec<SymbolLocation> {
|
||||
links.into_iter()
|
||||
.filter_map(|link| {
|
||||
uri_to_path(&link.target_uri.to_string()).map(|path| SymbolLocation {
|
||||
path,
|
||||
range: link.target_selection_range,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn uri_to_path(uri: &str) -> Option<PathBuf> {
|
||||
url::Url::parse(uri).ok()?.to_file_path().ok()
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum LspError {
|
||||
Io(std::io::Error),
|
||||
Json(serde_json::Error),
|
||||
InvalidHeader(String),
|
||||
MissingContentLength,
|
||||
InvalidContentLength(String),
|
||||
UnsupportedDocument(PathBuf),
|
||||
UnknownServer(String),
|
||||
DuplicateExtension {
|
||||
extension: String,
|
||||
existing_server: String,
|
||||
new_server: String,
|
||||
},
|
||||
PathToUrl(PathBuf),
|
||||
Protocol(String),
|
||||
}
|
||||
|
||||
impl Display for LspError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Io(error) => write!(f, "{error}"),
|
||||
Self::Json(error) => write!(f, "{error}"),
|
||||
Self::InvalidHeader(header) => write!(f, "invalid LSP header: {header}"),
|
||||
Self::MissingContentLength => write!(f, "missing LSP Content-Length header"),
|
||||
Self::InvalidContentLength(value) => {
|
||||
write!(f, "invalid LSP Content-Length value: {value}")
|
||||
}
|
||||
Self::UnsupportedDocument(path) => {
|
||||
write!(f, "no LSP server configured for {}", path.display())
|
||||
}
|
||||
Self::UnknownServer(name) => write!(f, "unknown LSP server: {name}"),
|
||||
Self::DuplicateExtension {
|
||||
extension,
|
||||
existing_server,
|
||||
new_server,
|
||||
} => write!(
|
||||
f,
|
||||
"duplicate LSP extension mapping for {extension}: {existing_server} and {new_server}"
|
||||
),
|
||||
Self::PathToUrl(path) => write!(f, "failed to convert path to file URL: {}", path.display()),
|
||||
Self::Protocol(message) => write!(f, "LSP protocol error: {message}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for LspError {}
|
||||
|
||||
impl From<std::io::Error> for LspError {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::Io(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for LspError {
|
||||
fn from(value: serde_json::Error) -> Self {
|
||||
Self::Json(value)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,284 @@
|
||||
mod client;
|
||||
mod error;
|
||||
mod manager;
|
||||
mod types;
|
||||
|
||||
pub use error::LspError;
|
||||
pub use manager::LspManager;
|
||||
pub use types::{
|
||||
FileDiagnostics, LspContextEnrichment, LspServerConfig, SymbolLocation, WorkspaceDiagnostics,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use lsp_types::{DiagnosticSeverity, Position};
|
||||
|
||||
use crate::{LspManager, LspServerConfig};
|
||||
|
||||
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!("lsp-{label}-{nanos}"))
|
||||
}
|
||||
|
||||
fn python3_path() -> Option<String> {
|
||||
let candidates = ["python3", "/usr/bin/python3"];
|
||||
candidates.iter().find_map(|candidate| {
|
||||
Command::new(candidate)
|
||||
.arg("--version")
|
||||
.output()
|
||||
.ok()
|
||||
.filter(|output| output.status.success())
|
||||
.map(|_| (*candidate).to_string())
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)]
|
||||
fn write_mock_server_script(root: &std::path::Path) -> PathBuf {
|
||||
let script_path = root.join("mock_lsp_server.py");
|
||||
fs::write(
|
||||
&script_path,
|
||||
r#"import json
|
||||
import sys
|
||||
|
||||
|
||||
def read_message():
|
||||
headers = {}
|
||||
while True:
|
||||
line = sys.stdin.buffer.readline()
|
||||
if not line:
|
||||
return None
|
||||
if line == b"\r\n":
|
||||
break
|
||||
key, value = line.decode("utf-8").split(":", 1)
|
||||
headers[key.lower()] = value.strip()
|
||||
length = int(headers["content-length"])
|
||||
body = sys.stdin.buffer.read(length)
|
||||
return json.loads(body)
|
||||
|
||||
|
||||
def write_message(payload):
|
||||
raw = json.dumps(payload).encode("utf-8")
|
||||
sys.stdout.buffer.write(f"Content-Length: {len(raw)}\r\n\r\n".encode("utf-8"))
|
||||
sys.stdout.buffer.write(raw)
|
||||
sys.stdout.buffer.flush()
|
||||
|
||||
|
||||
while True:
|
||||
message = read_message()
|
||||
if message is None:
|
||||
break
|
||||
|
||||
method = message.get("method")
|
||||
if method == "initialize":
|
||||
write_message({
|
||||
"jsonrpc": "2.0",
|
||||
"id": message["id"],
|
||||
"result": {
|
||||
"capabilities": {
|
||||
"definitionProvider": True,
|
||||
"referencesProvider": True,
|
||||
"textDocumentSync": 1,
|
||||
}
|
||||
},
|
||||
})
|
||||
elif method == "initialized":
|
||||
continue
|
||||
elif method == "textDocument/didOpen":
|
||||
document = message["params"]["textDocument"]
|
||||
write_message({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "textDocument/publishDiagnostics",
|
||||
"params": {
|
||||
"uri": document["uri"],
|
||||
"diagnostics": [
|
||||
{
|
||||
"range": {
|
||||
"start": {"line": 0, "character": 0},
|
||||
"end": {"line": 0, "character": 3},
|
||||
},
|
||||
"severity": 1,
|
||||
"source": "mock-server",
|
||||
"message": "mock error",
|
||||
}
|
||||
],
|
||||
},
|
||||
})
|
||||
elif method == "textDocument/didChange":
|
||||
continue
|
||||
elif method == "textDocument/didSave":
|
||||
continue
|
||||
elif method == "textDocument/definition":
|
||||
uri = message["params"]["textDocument"]["uri"]
|
||||
write_message({
|
||||
"jsonrpc": "2.0",
|
||||
"id": message["id"],
|
||||
"result": [
|
||||
{
|
||||
"uri": uri,
|
||||
"range": {
|
||||
"start": {"line": 0, "character": 0},
|
||||
"end": {"line": 0, "character": 3},
|
||||
},
|
||||
}
|
||||
],
|
||||
})
|
||||
elif method == "textDocument/references":
|
||||
uri = message["params"]["textDocument"]["uri"]
|
||||
write_message({
|
||||
"jsonrpc": "2.0",
|
||||
"id": message["id"],
|
||||
"result": [
|
||||
{
|
||||
"uri": uri,
|
||||
"range": {
|
||||
"start": {"line": 0, "character": 0},
|
||||
"end": {"line": 0, "character": 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"uri": uri,
|
||||
"range": {
|
||||
"start": {"line": 1, "character": 4},
|
||||
"end": {"line": 1, "character": 7},
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
elif method == "shutdown":
|
||||
write_message({"jsonrpc": "2.0", "id": message["id"], "result": None})
|
||||
elif method == "exit":
|
||||
break
|
||||
"#,
|
||||
)
|
||||
.expect("mock server should be written");
|
||||
script_path
|
||||
}
|
||||
|
||||
async fn wait_for_diagnostics(manager: &LspManager) {
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
loop {
|
||||
if manager
|
||||
.collect_workspace_diagnostics()
|
||||
.await
|
||||
.expect("diagnostics snapshot should load")
|
||||
.total_diagnostics()
|
||||
> 0
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("diagnostics should arrive from mock server");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn collects_diagnostics_and_symbol_navigation_from_mock_server() {
|
||||
let Some(python) = python3_path() else {
|
||||
return;
|
||||
};
|
||||
|
||||
// given
|
||||
let root = temp_dir("manager");
|
||||
fs::create_dir_all(root.join("src")).expect("workspace root should exist");
|
||||
let script_path = write_mock_server_script(&root);
|
||||
let source_path = root.join("src").join("main.rs");
|
||||
fs::write(&source_path, "fn main() {}\nlet value = 1;\n").expect("source file should exist");
|
||||
let manager = LspManager::new(vec![LspServerConfig {
|
||||
name: "rust-analyzer".to_string(),
|
||||
command: python,
|
||||
args: vec![script_path.display().to_string()],
|
||||
env: BTreeMap::new(),
|
||||
workspace_root: root.clone(),
|
||||
initialization_options: None,
|
||||
extension_to_language: BTreeMap::from([(".rs".to_string(), "rust".to_string())]),
|
||||
}])
|
||||
.expect("manager should build");
|
||||
manager
|
||||
.open_document(&source_path, &fs::read_to_string(&source_path).expect("source read should succeed"))
|
||||
.await
|
||||
.expect("document should open");
|
||||
wait_for_diagnostics(&manager).await;
|
||||
|
||||
// when
|
||||
let diagnostics = manager
|
||||
.collect_workspace_diagnostics()
|
||||
.await
|
||||
.expect("diagnostics should be available");
|
||||
let definitions = manager
|
||||
.go_to_definition(&source_path, Position::new(0, 0))
|
||||
.await
|
||||
.expect("definition request should succeed");
|
||||
let references = manager
|
||||
.find_references(&source_path, Position::new(0, 0), true)
|
||||
.await
|
||||
.expect("references request should succeed");
|
||||
|
||||
// then
|
||||
assert_eq!(diagnostics.files.len(), 1);
|
||||
assert_eq!(diagnostics.total_diagnostics(), 1);
|
||||
assert_eq!(diagnostics.files[0].diagnostics[0].severity, Some(DiagnosticSeverity::ERROR));
|
||||
assert_eq!(definitions.len(), 1);
|
||||
assert_eq!(definitions[0].start_line(), 1);
|
||||
assert_eq!(references.len(), 2);
|
||||
|
||||
manager.shutdown().await.expect("shutdown should succeed");
|
||||
fs::remove_dir_all(root).expect("temp workspace should be removed");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn renders_runtime_context_enrichment_for_prompt_usage() {
|
||||
let Some(python) = python3_path() else {
|
||||
return;
|
||||
};
|
||||
|
||||
// given
|
||||
let root = temp_dir("prompt");
|
||||
fs::create_dir_all(root.join("src")).expect("workspace root should exist");
|
||||
let script_path = write_mock_server_script(&root);
|
||||
let source_path = root.join("src").join("lib.rs");
|
||||
fs::write(&source_path, "pub fn answer() -> i32 { 42 }\n").expect("source file should exist");
|
||||
let manager = LspManager::new(vec![LspServerConfig {
|
||||
name: "rust-analyzer".to_string(),
|
||||
command: python,
|
||||
args: vec![script_path.display().to_string()],
|
||||
env: BTreeMap::new(),
|
||||
workspace_root: root.clone(),
|
||||
initialization_options: None,
|
||||
extension_to_language: BTreeMap::from([(".rs".to_string(), "rust".to_string())]),
|
||||
}])
|
||||
.expect("manager should build");
|
||||
manager
|
||||
.open_document(&source_path, &fs::read_to_string(&source_path).expect("source read should succeed"))
|
||||
.await
|
||||
.expect("document should open");
|
||||
wait_for_diagnostics(&manager).await;
|
||||
|
||||
// when
|
||||
let enrichment = manager
|
||||
.context_enrichment(&source_path, Position::new(0, 0))
|
||||
.await
|
||||
.expect("context enrichment should succeed");
|
||||
let rendered = enrichment.render_prompt_section();
|
||||
|
||||
// then
|
||||
assert!(rendered.contains("# LSP context"));
|
||||
assert!(rendered.contains("Workspace diagnostics: 1 across 1 file(s)"));
|
||||
assert!(rendered.contains("Definitions:"));
|
||||
assert!(rendered.contains("References:"));
|
||||
assert!(rendered.contains("mock error"));
|
||||
|
||||
manager.shutdown().await.expect("shutdown should succeed");
|
||||
fs::remove_dir_all(root).expect("temp workspace should be removed");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use lsp_types::Position;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::client::LspClient;
|
||||
use crate::error::LspError;
|
||||
use crate::types::{
|
||||
normalize_extension, FileDiagnostics, LspContextEnrichment, LspServerConfig, SymbolLocation,
|
||||
WorkspaceDiagnostics,
|
||||
};
|
||||
|
||||
pub struct LspManager {
|
||||
server_configs: BTreeMap<String, LspServerConfig>,
|
||||
extension_map: BTreeMap<String, String>,
|
||||
clients: Mutex<BTreeMap<String, Arc<LspClient>>>,
|
||||
}
|
||||
|
||||
impl LspManager {
|
||||
pub fn new(server_configs: Vec<LspServerConfig>) -> Result<Self, LspError> {
|
||||
let mut configs_by_name = BTreeMap::new();
|
||||
let mut extension_map = BTreeMap::new();
|
||||
|
||||
for config in server_configs {
|
||||
for extension in config.extension_to_language.keys() {
|
||||
let normalized = normalize_extension(extension);
|
||||
if let Some(existing_server) = extension_map.insert(normalized.clone(), config.name.clone()) {
|
||||
return Err(LspError::DuplicateExtension {
|
||||
extension: normalized,
|
||||
existing_server,
|
||||
new_server: config.name.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
configs_by_name.insert(config.name.clone(), config);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
server_configs: configs_by_name,
|
||||
extension_map,
|
||||
clients: Mutex::new(BTreeMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn supports_path(&self, path: &Path) -> bool {
|
||||
path.extension().is_some_and(|extension| {
|
||||
let normalized = normalize_extension(extension.to_string_lossy().as_ref());
|
||||
self.extension_map.contains_key(&normalized)
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn open_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
|
||||
self.client_for_path(path).await?.open_document(path, text).await
|
||||
}
|
||||
|
||||
pub async fn sync_document_from_disk(&self, path: &Path) -> Result<(), LspError> {
|
||||
let contents = std::fs::read_to_string(path)?;
|
||||
self.change_document(path, &contents).await?;
|
||||
self.save_document(path).await
|
||||
}
|
||||
|
||||
pub async fn change_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
|
||||
self.client_for_path(path).await?.change_document(path, text).await
|
||||
}
|
||||
|
||||
pub async fn save_document(&self, path: &Path) -> Result<(), LspError> {
|
||||
self.client_for_path(path).await?.save_document(path).await
|
||||
}
|
||||
|
||||
pub async fn close_document(&self, path: &Path) -> Result<(), LspError> {
|
||||
self.client_for_path(path).await?.close_document(path).await
|
||||
}
|
||||
|
||||
pub async fn go_to_definition(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
) -> Result<Vec<SymbolLocation>, LspError> {
|
||||
let mut locations = self.client_for_path(path).await?.go_to_definition(path, position).await?;
|
||||
dedupe_locations(&mut locations);
|
||||
Ok(locations)
|
||||
}
|
||||
|
||||
pub async fn find_references(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
include_declaration: bool,
|
||||
) -> Result<Vec<SymbolLocation>, LspError> {
|
||||
let mut locations = self
|
||||
.client_for_path(path)
|
||||
.await?
|
||||
.find_references(path, position, include_declaration)
|
||||
.await?;
|
||||
dedupe_locations(&mut locations);
|
||||
Ok(locations)
|
||||
}
|
||||
|
||||
pub async fn collect_workspace_diagnostics(&self) -> Result<WorkspaceDiagnostics, LspError> {
|
||||
let clients = self.clients.lock().await.values().cloned().collect::<Vec<_>>();
|
||||
let mut files = Vec::new();
|
||||
|
||||
for client in clients {
|
||||
for (uri, diagnostics) in client.diagnostics_snapshot().await {
|
||||
let Ok(path) = url::Url::parse(&uri)
|
||||
.and_then(|url| url.to_file_path().map_err(|()| url::ParseError::RelativeUrlWithoutBase))
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if diagnostics.is_empty() {
|
||||
continue;
|
||||
}
|
||||
files.push(FileDiagnostics {
|
||||
path,
|
||||
uri,
|
||||
diagnostics,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
files.sort_by(|left, right| left.path.cmp(&right.path));
|
||||
Ok(WorkspaceDiagnostics { files })
|
||||
}
|
||||
|
||||
pub async fn context_enrichment(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
) -> Result<LspContextEnrichment, LspError> {
|
||||
Ok(LspContextEnrichment {
|
||||
file_path: path.to_path_buf(),
|
||||
diagnostics: self.collect_workspace_diagnostics().await?,
|
||||
definitions: self.go_to_definition(path, position).await?,
|
||||
references: self.find_references(path, position, true).await?,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn shutdown(&self) -> Result<(), LspError> {
|
||||
let mut clients = self.clients.lock().await;
|
||||
let drained = clients.values().cloned().collect::<Vec<_>>();
|
||||
clients.clear();
|
||||
drop(clients);
|
||||
|
||||
for client in drained {
|
||||
client.shutdown().await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn client_for_path(&self, path: &Path) -> Result<Arc<LspClient>, LspError> {
|
||||
let extension = path
|
||||
.extension()
|
||||
.map(|extension| normalize_extension(extension.to_string_lossy().as_ref()))
|
||||
.ok_or_else(|| LspError::UnsupportedDocument(path.to_path_buf()))?;
|
||||
let server_name = self
|
||||
.extension_map
|
||||
.get(&extension)
|
||||
.cloned()
|
||||
.ok_or_else(|| LspError::UnsupportedDocument(path.to_path_buf()))?;
|
||||
|
||||
let mut clients = self.clients.lock().await;
|
||||
if let Some(client) = clients.get(&server_name) {
|
||||
return Ok(client.clone());
|
||||
}
|
||||
|
||||
let config = self
|
||||
.server_configs
|
||||
.get(&server_name)
|
||||
.cloned()
|
||||
.ok_or_else(|| LspError::UnknownServer(server_name.clone()))?;
|
||||
let client = Arc::new(LspClient::connect(config).await?);
|
||||
clients.insert(server_name, client.clone());
|
||||
Ok(client)
|
||||
}
|
||||
}
|
||||
|
||||
fn dedupe_locations(locations: &mut Vec<SymbolLocation>) {
|
||||
let mut seen = BTreeSet::new();
|
||||
locations.retain(|location| {
|
||||
seen.insert((
|
||||
location.path.clone(),
|
||||
location.range.start.line,
|
||||
location.range.start.character,
|
||||
location.range.end.line,
|
||||
location.range.end.character,
|
||||
))
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use lsp_types::{Diagnostic, Range};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LspServerConfig {
|
||||
pub name: String,
|
||||
pub command: String,
|
||||
pub args: Vec<String>,
|
||||
pub env: BTreeMap<String, String>,
|
||||
pub workspace_root: PathBuf,
|
||||
pub initialization_options: Option<Value>,
|
||||
pub extension_to_language: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
impl LspServerConfig {
|
||||
#[must_use]
|
||||
pub fn language_id_for(&self, path: &Path) -> Option<&str> {
|
||||
let extension = normalize_extension(path.extension()?.to_string_lossy().as_ref());
|
||||
self.extension_to_language
|
||||
.get(&extension)
|
||||
.map(String::as_str)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct FileDiagnostics {
|
||||
pub path: PathBuf,
|
||||
pub uri: String,
|
||||
pub diagnostics: Vec<Diagnostic>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct WorkspaceDiagnostics {
|
||||
pub files: Vec<FileDiagnostics>,
|
||||
}
|
||||
|
||||
impl WorkspaceDiagnostics {
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.files.is_empty()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn total_diagnostics(&self) -> usize {
|
||||
self.files.iter().map(|file| file.diagnostics.len()).sum()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SymbolLocation {
|
||||
pub path: PathBuf,
|
||||
pub range: Range,
|
||||
}
|
||||
|
||||
impl SymbolLocation {
|
||||
#[must_use]
|
||||
pub fn start_line(&self) -> u32 {
|
||||
self.range.start.line + 1
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn start_character(&self) -> u32 {
|
||||
self.range.start.character + 1
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for SymbolLocation {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"{}:{}:{}",
|
||||
self.path.display(),
|
||||
self.start_line(),
|
||||
self.start_character()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct LspContextEnrichment {
|
||||
pub file_path: PathBuf,
|
||||
pub diagnostics: WorkspaceDiagnostics,
|
||||
pub definitions: Vec<SymbolLocation>,
|
||||
pub references: Vec<SymbolLocation>,
|
||||
}
|
||||
|
||||
impl LspContextEnrichment {
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.diagnostics.is_empty() && self.definitions.is_empty() && self.references.is_empty()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn render_prompt_section(&self) -> String {
|
||||
const MAX_RENDERED_DIAGNOSTICS: usize = 12;
|
||||
const MAX_RENDERED_LOCATIONS: usize = 12;
|
||||
|
||||
let mut lines = vec!["# LSP context".to_string()];
|
||||
lines.push(format!(" - Focus file: {}", self.file_path.display()));
|
||||
lines.push(format!(
|
||||
" - Workspace diagnostics: {} across {} file(s)",
|
||||
self.diagnostics.total_diagnostics(),
|
||||
self.diagnostics.files.len()
|
||||
));
|
||||
|
||||
if !self.diagnostics.files.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("Diagnostics:".to_string());
|
||||
let mut rendered = 0usize;
|
||||
for file in &self.diagnostics.files {
|
||||
for diagnostic in &file.diagnostics {
|
||||
if rendered == MAX_RENDERED_DIAGNOSTICS {
|
||||
lines.push(" - Additional diagnostics omitted for brevity.".to_string());
|
||||
break;
|
||||
}
|
||||
let severity = diagnostic_severity_label(diagnostic.severity);
|
||||
lines.push(format!(
|
||||
" - {}:{}:{} [{}] {}",
|
||||
file.path.display(),
|
||||
diagnostic.range.start.line + 1,
|
||||
diagnostic.range.start.character + 1,
|
||||
severity,
|
||||
diagnostic.message.replace('\n', " ")
|
||||
));
|
||||
rendered += 1;
|
||||
}
|
||||
if rendered == MAX_RENDERED_DIAGNOSTICS {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !self.definitions.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("Definitions:".to_string());
|
||||
lines.extend(
|
||||
self.definitions
|
||||
.iter()
|
||||
.take(MAX_RENDERED_LOCATIONS)
|
||||
.map(|location| format!(" - {location}")),
|
||||
);
|
||||
if self.definitions.len() > MAX_RENDERED_LOCATIONS {
|
||||
lines.push(" - Additional definitions omitted for brevity.".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
if !self.references.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("References:".to_string());
|
||||
lines.extend(
|
||||
self.references
|
||||
.iter()
|
||||
.take(MAX_RENDERED_LOCATIONS)
|
||||
.map(|location| format!(" - {location}")),
|
||||
);
|
||||
if self.references.len() > MAX_RENDERED_LOCATIONS {
|
||||
lines.push(" - Additional references omitted for brevity.".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
lines.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub(crate) fn normalize_extension(extension: &str) -> String {
|
||||
if extension.starts_with('.') {
|
||||
extension.to_ascii_lowercase()
|
||||
} else {
|
||||
format!(".{}", extension.to_ascii_lowercase())
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnostic_severity_label(severity: Option<lsp_types::DiagnosticSeverity>) -> &'static str {
|
||||
match severity {
|
||||
Some(lsp_types::DiagnosticSeverity::ERROR) => "error",
|
||||
Some(lsp_types::DiagnosticSeverity::WARNING) => "warning",
|
||||
Some(lsp_types::DiagnosticSeverity::INFORMATION) => "info",
|
||||
Some(lsp_types::DiagnosticSeverity::HINT) => "hint",
|
||||
_ => "unknown",
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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,59 @@
|
||||
# 插件模块 (plugins)
|
||||
|
||||
本模块实现了 Claw 的插件系统,允许通过外部扩展来增强 AI 的功能,包括自定义工具、命令以及在工具执行前后运行的钩子 (Hooks)。
|
||||
|
||||
## 概览
|
||||
|
||||
`plugins` 模块的主要职责是:
|
||||
- 定义插件的清单格式 (`plugin.json`) 和元数据结构。
|
||||
- 管理插件的完整生命周期:安装、加载、初始化、启用/禁用、更新和卸载。
|
||||
- 提供插件类型的抽象:
|
||||
- **Builtin (内置)**:编译在程序内部的插件。
|
||||
- **Bundled (绑定)**:随应用程序分发但作为独立文件存在的插件。
|
||||
- **External (外部)**:用户自行安装或从远程仓库下载的插件。
|
||||
- 实现插件隔离与执行机制,支持插件定义的自定义工具。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **插件清单 (PluginManifest)**:每个插件必须包含一个 `plugin.json`,详细说明其名称、版本、所需权限、钩子、生命周期脚本以及它所暴露的工具。
|
||||
- **自定义工具 (PluginTool)**:插件可以定义全新的工具供 AI 调用。这些工具在执行时被作为独立的外部进程启动。
|
||||
- **钩子系统 (Hooks)**:支持 `PreToolUse` 和 `PostToolUse` 钩子,允许插件在 AI 调用任何工具之前或之后执行特定的逻辑。
|
||||
- **生命周期管理**:提供 `Init` 和 `Shutdown` 阶段,允许插件在加载时进行环境准备,在卸载或关闭时进行清理。
|
||||
- **权限模型**:强制要求插件声明权限(`read`, `write`, `execute`),并为定义的工具指定安全级别(`read-only`, `workspace-write`, `danger-full-access`)。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`lib.rs`**: 包含了插件定义的各种结构体(Manifest, Metadata, Tool, Permission 等)以及插件特性的定义。
|
||||
- **`manager.rs`**: 实现了 `PluginManager`,负责插件在磁盘上的组织、注册表的维护以及安装/更新逻辑。
|
||||
- **`hooks.rs`**: 实现了钩子执行器 (`HookRunner`),负责在正确的时机触发插件定义的脚本。
|
||||
|
||||
### 插件加载与执行流程
|
||||
|
||||
1. `PluginManager` 扫描指定的目录(内置、绑定及外部安装目录)。
|
||||
2. 读取并验证每个插件的 `plugin.json`。
|
||||
3. 如果插件被启用,则初始化该插件并将其定义的工具注册到系统的全局工具注册表中。
|
||||
4. 当 AI 调用插件工具时,系统根据清单中定义的命令行信息启动一个子进程,并通过标准输入/环境变量传递参数。
|
||||
5. 钩子逻辑会在工具执行的生命周期内被自动触发。
|
||||
|
||||
## 使用示例 (插件定义样例)
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "my-custom-plugin",
|
||||
"version": "1.0.0",
|
||||
"description": "一个演示插件",
|
||||
"permissions": ["read", "execute"],
|
||||
"tools": [
|
||||
{
|
||||
"name": "custom_search",
|
||||
"description": "执行自定义搜索",
|
||||
"inputSchema": { "type": "object", "properties": { "query": { "type": "string" } } },
|
||||
"command": "python3",
|
||||
"args": ["search_script.py"],
|
||||
"requiredPermission": "read-only"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
@@ -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'
|
||||
@@ -0,0 +1,563 @@
|
||||
use std::ffi::OsStr;
|
||||
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(®istry).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
@@ -7,11 +7,14 @@ publish.workspace = true
|
||||
|
||||
[dependencies]
|
||||
sha2 = "0.10"
|
||||
which = "7"
|
||||
glob = "0.3"
|
||||
plugins = { path = "../plugins" }
|
||||
regex = "1"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
tokio = { version = "1", features = ["io-util", "macros", "process", "rt", "rt-multi-thread", "time"] }
|
||||
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]
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
# 运行时模块 (runtime)
|
||||
|
||||
本模块是 Claw 的核心引擎,负责协调 AI 模型、工具执行、会话管理和权限控制之间的所有交互。
|
||||
|
||||
## 概览
|
||||
|
||||
`runtime` 模块是整个系统的“中枢神经”,其主要职责包括:
|
||||
- **对话驱动**:管理“用户-助手-工具”的循环迭代。
|
||||
- **会话持久化**:负责会话的加载、保存及历史记录的压缩 (Compaction)。
|
||||
- **MCP 客户端**:实现模型上下文协议 (Model Context Protocol),支持与外部 MCP 服务器通信。
|
||||
- **安全沙箱与权限**:实施基于策略的工具执行权限检查。
|
||||
- **上下文构建**:动态生成系统提示词 (System Prompt),集成工作区上下文。
|
||||
- **消耗统计**:精确跟踪 Token 使用情况和 Token 缓存状态。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **ConversationRuntime**:核心驱动类,支持流式响应处理和多轮工具调用迭代。
|
||||
- **权限引擎 (Permissions)**:提供多种模式(`ReadOnly`, `WorkspaceWrite`, `DangerFullAccess`),并支持交互式权限确认。
|
||||
- **会话压缩 (Compaction)**:当对话历史过长影响性能或成本时,自动将旧消息总结为摘要,保持上下文精简。
|
||||
- **钩子集成 (Hooks)**:在工具执行的前后触发插件定义的钩子,支持干预工具输入或处理执行结果。
|
||||
- **沙箱执行 (Sandbox)**:为 Bash 等敏感工具提供受限的执行环境。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心子模块
|
||||
|
||||
- **`conversation.rs`**: 定义了核心的 `ConversationRuntime` 和 `ApiClient`/`ToolExecutor` 特性。
|
||||
- **`mcp_stdio.rs` / `mcp_client.rs`**: 实现了完整的 MCP 规范,支持通过标准输入/输出与外部工具服务器交互。
|
||||
- **`session.rs`**: 定义了消息模型 (`ConversationMessage`)、内容块 (`ContentBlock`) 和会话序列化逻辑。
|
||||
- **`permissions.rs`**: 实现了权限审核逻辑和提示器接口。
|
||||
- **`compact.rs`**: 包含了基于 LLM 的会话摘要生成和历史裁剪算法。
|
||||
- **`config.rs`**: 负责加载和合并多层级的配置文件。
|
||||
|
||||
### 对话循环流程 (run_turn)
|
||||
|
||||
1. 将用户输入推入 `Session`。
|
||||
2. 调用 `ApiClient` 发起流式请求。
|
||||
3. 监听 `AssistantEvent`,解析文本内容和工具调用请求。
|
||||
4. **工具权限审核**:针对每个 `ToolUse`,根据 `PermissionPolicy` 决定是允许、拒绝还是询问用户。
|
||||
5. **执行工具**:若允许,则通过 `ToolExecutor`(或 MCP 客户端)执行工具,并运行相关的 `Pre/Post Hooks`。
|
||||
6. 将工具结果反馈给 AI,进入下一轮迭代,直到 AI 给出最终回复。
|
||||
|
||||
## 使用示例 (内部)
|
||||
|
||||
```rust
|
||||
use runtime::{ConversationRuntime, Session, PermissionPolicy, PermissionMode};
|
||||
|
||||
// 初始化运行时
|
||||
let mut runtime = ConversationRuntime::new(
|
||||
Session::new(),
|
||||
api_client,
|
||||
tool_executor,
|
||||
PermissionPolicy::new(PermissionMode::WorkspaceWrite),
|
||||
system_prompt,
|
||||
);
|
||||
|
||||
// 运行一轮对话
|
||||
let summary = runtime.run_turn("帮我重构 src/lib.rs", Some(&mut cli_prompter))?;
|
||||
println!("共迭代 {} 次,消耗 {} tokens", summary.iterations, summary.usage.total_tokens());
|
||||
```
|
||||
+210
-9
@@ -1,3 +1,4 @@
|
||||
use std::env;
|
||||
use std::io;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::time::Duration;
|
||||
@@ -7,6 +8,13 @@ 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,
|
||||
@@ -16,8 +24,17 @@ pub struct BashCommandInput {
|
||||
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,
|
||||
@@ -45,13 +62,18 @@ pub struct BashCommandOutput {
|
||||
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 child = Command::new("sh")
|
||||
.arg("-lc")
|
||||
.arg(&input.command)
|
||||
let mut child = prepare_command(&input.command, &cwd, &sandbox_status, false);
|
||||
let child = child
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
@@ -72,16 +94,20 @@ pub fn execute_bash(input: BashCommandInput) -> io::Result<BashCommandOutput> {
|
||||
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))
|
||||
runtime.block_on(execute_bash_async(input, sandbox_status, cwd))
|
||||
}
|
||||
|
||||
async fn execute_bash_async(input: BashCommandInput) -> io::Result<BashCommandOutput> {
|
||||
let mut command = TokioCommand::new("sh");
|
||||
command.arg("-lc").arg(&input.command);
|
||||
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 {
|
||||
@@ -102,6 +128,7 @@ async fn execute_bash_async(input: BashCommandInput) -> io::Result<BashCommandOu
|
||||
structured_content: None,
|
||||
persisted_output_path: None,
|
||||
persisted_output_size: None,
|
||||
sandbox_status: Some(sandbox_status),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -110,8 +137,8 @@ async fn execute_bash_async(input: BashCommandInput) -> io::Result<BashCommandOu
|
||||
};
|
||||
|
||||
let (output, interrupted) = output_result;
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).into_owned();
|
||||
let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
|
||||
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 {
|
||||
@@ -136,12 +163,113 @@ async fn execute_bash_async(input: BashCommandInput) -> io::Result<BashCommandOu
|
||||
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 prepared = if cfg!(target_os = "windows") && !sh_exists() {
|
||||
let mut p = Command::new("cmd");
|
||||
p.arg("/C").arg(command).current_dir(cwd);
|
||||
p
|
||||
} else {
|
||||
let mut p = Command::new("sh");
|
||||
p.arg("-lc").arg(command).current_dir(cwd);
|
||||
if sandbox_status.filesystem_active {
|
||||
p.env("HOME", cwd.join(".sandbox-home"));
|
||||
p.env("TMPDIR", cwd.join(".sandbox-tmp"));
|
||||
}
|
||||
p
|
||||
};
|
||||
prepared
|
||||
}
|
||||
|
||||
fn sh_exists() -> bool {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
which::which("sh").is_ok()
|
||||
}
|
||||
#[cfg(not(windows))]
|
||||
{
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
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 prepared = if cfg!(target_os = "windows") && !sh_exists() {
|
||||
let mut p = TokioCommand::new("cmd");
|
||||
p.arg("/C").arg(command).current_dir(cwd);
|
||||
p
|
||||
} else {
|
||||
let mut p = TokioCommand::new("sh");
|
||||
p.arg("-lc").arg(command).current_dir(cwd);
|
||||
if sandbox_status.filesystem_active {
|
||||
p.env("HOME", cwd.join(".sandbox-home"));
|
||||
p.env("TMPDIR", cwd.join(".sandbox-tmp"));
|
||||
}
|
||||
p
|
||||
};
|
||||
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() {
|
||||
@@ -151,10 +279,83 @@ mod tests {
|
||||
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
@@ -54,3 +54,58 @@ impl BootstrapPlan {
|
||||
&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());
|
||||
}
|
||||
}
|
||||
+382
-146
@@ -1,8 +1,11 @@
|
||||
use std::fs;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
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,
|
||||
@@ -18,6 +21,7 @@ impl Default for CompactionConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of compacting a session into a summary plus preserved tail messages.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CompactionResult {
|
||||
pub summary: String,
|
||||
@@ -26,17 +30,27 @@ pub struct CompactionResult {
|
||||
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 {
|
||||
session.messages.len() > config.preserve_recent_messages
|
||||
&& estimate_session_tokens(session) >= config.max_estimated_tokens
|
||||
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");
|
||||
@@ -52,6 +66,7 @@ pub fn format_compact_summary(summary: &str) -> String {
|
||||
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,
|
||||
@@ -59,21 +74,24 @@ pub fn get_compact_continuation_message(
|
||||
recent_messages_preserved: bool,
|
||||
) -> String {
|
||||
let mut base = format!(
|
||||
"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{}",
|
||||
"{COMPACT_CONTINUATION_PREAMBLE}{}",
|
||||
format_compact_summary(summary)
|
||||
);
|
||||
|
||||
if recent_messages_preserved {
|
||||
base.push_str("\n\nRecent messages are preserved verbatim.");
|
||||
base.push_str("\n\n");
|
||||
base.push_str(COMPACT_RECENT_MESSAGES_NOTE);
|
||||
}
|
||||
|
||||
if suppress_follow_up_questions {
|
||||
base.push_str("\nContinue 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.");
|
||||
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) {
|
||||
@@ -85,15 +103,65 @@ pub fn compact_session(session: &Session, config: CompactionConfig) -> Compactio
|
||||
};
|
||||
}
|
||||
|
||||
let keep_from = session
|
||||
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);
|
||||
let removed = &session.messages[..keep_from];
|
||||
// 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 = summarize_messages(removed);
|
||||
let summary =
|
||||
merge_compact_summaries(existing_summary.as_deref(), &summarize_messages(removed));
|
||||
let formatted_summary = format_compact_summary(&summary);
|
||||
persist_compact_summary(&formatted_summary);
|
||||
let continuation = get_compact_continuation_message(&summary, true, !preserved.is_empty());
|
||||
|
||||
let mut compacted_messages = vec![ConversationMessage {
|
||||
@@ -103,45 +171,26 @@ pub fn compact_session(session: &Session, config: CompactionConfig) -> Compactio
|
||||
}];
|
||||
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: Session {
|
||||
version: session.version,
|
||||
messages: compacted_messages,
|
||||
metadata: session.metadata.clone(),
|
||||
},
|
||||
compacted_session,
|
||||
removed_message_count: removed.len(),
|
||||
}
|
||||
}
|
||||
|
||||
fn persist_compact_summary(formatted_summary: &str) {
|
||||
if formatted_summary.trim().is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let Ok(cwd) = std::env::current_dir() else {
|
||||
return;
|
||||
};
|
||||
let memory_dir = cwd.join(".claude").join("memory");
|
||||
if fs::create_dir_all(&memory_dir).is_err() {
|
||||
return;
|
||||
}
|
||||
|
||||
let path = memory_dir.join(compact_summary_filename());
|
||||
let _ = fs::write(path, render_memory_file(formatted_summary));
|
||||
}
|
||||
|
||||
fn compact_summary_filename() -> String {
|
||||
let timestamp = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
format!("summary-{timestamp}.md")
|
||||
}
|
||||
|
||||
fn render_memory_file(formatted_summary: &str) -> String {
|
||||
format!("# Project memory\n\n{}\n", formatted_summary.trim())
|
||||
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 {
|
||||
@@ -164,7 +213,7 @@ fn summarize_messages(messages: &[ConversationMessage]) -> String {
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::ToolUse { name, .. } => Some(name.as_str()),
|
||||
ContentBlock::ToolResult { tool_name, .. } => Some(tool_name.as_str()),
|
||||
ContentBlock::Text { .. } | ContentBlock::Thinking { .. } => None,
|
||||
ContentBlock::Text { .. } => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
tool_names.sort_unstable();
|
||||
@@ -231,10 +280,44 @@ fn summarize_messages(messages: &[ConversationMessage]) -> 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::Thinking { text, .. } => format!("thinking: {text}"),
|
||||
ContentBlock::ToolUse { name, input, .. } => format!("tool_use {name}({input})"),
|
||||
ContentBlock::ToolResult {
|
||||
tool_name,
|
||||
@@ -293,7 +376,7 @@ fn collect_key_files(messages: &[ConversationMessage]) -> Vec<String> {
|
||||
.iter()
|
||||
.flat_map(|message| message.blocks.iter())
|
||||
.map(|block| match block {
|
||||
ContentBlock::Text { text } | ContentBlock::Thinking { text, .. } => text.as_str(),
|
||||
ContentBlock::Text { text } => text.as_str(),
|
||||
ContentBlock::ToolUse { input, .. } => input.as_str(),
|
||||
ContentBlock::ToolResult { output, .. } => output.as_str(),
|
||||
})
|
||||
@@ -315,15 +398,10 @@ fn infer_current_work(messages: &[ConversationMessage]) -> Option<String> {
|
||||
|
||||
fn first_text_block(message: &ConversationMessage) -> Option<&str> {
|
||||
message.blocks.iter().find_map(|block| match block {
|
||||
ContentBlock::Text { text } | ContentBlock::Thinking { text, .. }
|
||||
if !text.trim().is_empty() =>
|
||||
{
|
||||
Some(text.as_str())
|
||||
}
|
||||
ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.as_str()),
|
||||
ContentBlock::ToolUse { .. }
|
||||
| ContentBlock::ToolResult { .. }
|
||||
| ContentBlock::Text { .. }
|
||||
| ContentBlock::Thinking { .. } => None,
|
||||
| ContentBlock::Text { .. } => None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -368,7 +446,7 @@ fn estimate_message_tokens(message: &ConversationMessage) -> usize {
|
||||
.blocks
|
||||
.iter()
|
||||
.map(|block| match block {
|
||||
ContentBlock::Text { text } | ContentBlock::Thinking { text, .. } => text.len() / 4 + 1,
|
||||
ContentBlock::Text { text } => text.len() / 4 + 1,
|
||||
ContentBlock::ToolUse { name, input, .. } => (name.len() + input.len()) / 4 + 1,
|
||||
ContentBlock::ToolResult {
|
||||
tool_name, output, ..
|
||||
@@ -414,34 +492,84 @@ fn collapse_blank_lines(content: &str) -> String {
|
||||
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, estimate_session_tokens, format_compact_summary,
|
||||
infer_pending_work, render_memory_file, should_compact, CompactionConfig,
|
||||
collect_key_files, compact_session, format_compact_summary,
|
||||
get_compact_continuation_message, infer_pending_work, should_compact, CompactionConfig,
|
||||
};
|
||||
use std::fs;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
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");
|
||||
assert_eq!(
|
||||
render_memory_file("Summary:\nKept work"),
|
||||
"# Project memory\n\nSummary:\nKept work\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves_small_sessions_unchanged() {
|
||||
let session = Session {
|
||||
version: 1,
|
||||
messages: vec![ConversationMessage::user_text("hello")],
|
||||
metadata: None,
|
||||
};
|
||||
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);
|
||||
@@ -450,83 +578,23 @@ mod tests {
|
||||
assert!(result.formatted_summary.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn persists_compacted_summaries_under_dot_claude_memory() {
|
||||
let _guard = crate::test_env_lock();
|
||||
let temp = std::env::temp_dir().join(format!(
|
||||
"runtime-compact-memory-{}",
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time after epoch")
|
||||
.as_nanos()
|
||||
));
|
||||
fs::create_dir_all(&temp).expect("temp dir");
|
||||
let previous = std::env::current_dir().expect("cwd");
|
||||
std::env::set_current_dir(&temp).expect("set cwd");
|
||||
|
||||
let session = Session {
|
||||
version: 1,
|
||||
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,
|
||||
},
|
||||
);
|
||||
let memory_dir = temp.join(".claude").join("memory");
|
||||
let files = fs::read_dir(&memory_dir)
|
||||
.expect("memory dir exists")
|
||||
.flatten()
|
||||
.map(|entry| entry.path())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(result.removed_message_count, 2);
|
||||
assert_eq!(files.len(), 1);
|
||||
let persisted = fs::read_to_string(&files[0]).expect("memory file readable");
|
||||
|
||||
std::env::set_current_dir(previous).expect("restore cwd");
|
||||
fs::remove_dir_all(temp).expect("cleanup temp dir");
|
||||
|
||||
assert!(persisted.contains("# Project memory"));
|
||||
assert!(persisted.contains("Summary:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compacts_older_messages_into_a_system_summary() {
|
||||
let session = Session {
|
||||
version: 1,
|
||||
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,
|
||||
},
|
||||
],
|
||||
metadata: None,
|
||||
};
|
||||
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,
|
||||
@@ -536,7 +604,14 @@ mod tests {
|
||||
},
|
||||
);
|
||||
|
||||
assert_eq!(result.removed_message_count, 2);
|
||||
// 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
|
||||
@@ -554,11 +629,98 @@ mod tests {
|
||||
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!(
|
||||
estimate_session_tokens(&result.compacted_session) < estimate_session_tokens(&session)
|
||||
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 {
|
||||
@@ -577,6 +739,80 @@ mod tests {
|
||||
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(&[
|
||||
|
||||
+1252
-49
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(¤t);
|
||||
}
|
||||
|
||||
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"#
|
||||
);
|
||||
}
|
||||
}
|
||||
+1145
-80
File diff suppressed because it is too large
Load Diff
@@ -9,6 +9,41 @@ 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")]
|
||||
@@ -22,6 +57,7 @@ pub struct TextFilePayload {
|
||||
pub total_lines: usize,
|
||||
}
|
||||
|
||||
/// Output envelope for the `read_file` tool.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct ReadFileOutput {
|
||||
#[serde(rename = "type")]
|
||||
@@ -29,6 +65,7 @@ pub struct ReadFileOutput {
|
||||
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")]
|
||||
@@ -42,6 +79,7 @@ pub struct StructuredPatchHunk {
|
||||
pub lines: Vec<String>,
|
||||
}
|
||||
|
||||
/// Output envelope for full-file write operations.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct WriteFileOutput {
|
||||
#[serde(rename = "type")]
|
||||
@@ -57,6 +95,7 @@ pub struct WriteFileOutput {
|
||||
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")]
|
||||
@@ -77,6 +116,7 @@ pub struct EditFileOutput {
|
||||
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")]
|
||||
@@ -87,6 +127,7 @@ pub struct GlobSearchOutput {
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
/// Parameters accepted by the grep-style search tool.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct GrepSearchInput {
|
||||
pub pattern: String,
|
||||
@@ -112,6 +153,7 @@ pub struct GrepSearchInput {
|
||||
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>,
|
||||
@@ -129,12 +171,35 @@ pub struct GrepSearchOutput {
|
||||
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());
|
||||
@@ -155,7 +220,19 @@ pub fn read_file(
|
||||
})
|
||||
}
|
||||
|
||||
/// 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() {
|
||||
@@ -177,6 +254,7 @@ pub fn write_file(path: &str, content: &str) -> io::Result<WriteFileOutput> {
|
||||
})
|
||||
}
|
||||
|
||||
/// Performs an in-file string replacement and returns patch metadata.
|
||||
pub fn edit_file(
|
||||
path: &str,
|
||||
old_string: &str,
|
||||
@@ -217,6 +295,7 @@ pub fn edit_file(
|
||||
})
|
||||
}
|
||||
|
||||
/// 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
|
||||
@@ -229,12 +308,20 @@ pub fn glob_search(pattern: &str, path: Option<&str>) -> io::Result<GlobSearchOu
|
||||
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();
|
||||
let entries = glob::glob(&search_pattern)
|
||||
.map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error.to_string()))?;
|
||||
for entry in entries.flatten() {
|
||||
if entry.is_file() {
|
||||
matches.push(entry);
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -260,6 +347,7 @@ pub fn glob_search(pattern: &str, path: Option<&str>) -> io::Result<GlobSearchOu
|
||||
})
|
||||
}
|
||||
|
||||
/// 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
|
||||
@@ -477,11 +565,98 @@ fn normalize_path_allow_missing(path: &str) -> io::Result<PathBuf> {
|
||||
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, glob_search, grep_search, read_file, write_file, GrepSearchInput};
|
||||
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()
|
||||
@@ -513,6 +688,73 @@ mod tests {
|
||||
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");
|
||||
@@ -547,4 +789,51 @@ mod tests {
|
||||
.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", "tests@example.com"]);
|
||||
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", "tests@example.com"]);
|
||||
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", "tests@example.com"]);
|
||||
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,
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,988 @@
|
||||
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: Send {
|
||||
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)]
|
||||
#[allow(unused_mut)]
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
+106
-16
@@ -1,57 +1,112 @@
|
||||
//! 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, McpClaudeAiProxyServerConfig,
|
||||
McpConfigCollection, McpOAuthConfig, McpRemoteServerConfig, McpSdkServerConfig,
|
||||
ConfigEntry, ConfigError, ConfigLoader, ConfigSource, McpConfigCollection,
|
||||
McpManagedProxyServerConfig, McpOAuthConfig, McpRemoteServerConfig, McpSdkServerConfig,
|
||||
McpServerConfig, McpStdioServerConfig, McpTransport, McpWebSocketServerConfig, OAuthConfig,
|
||||
ResolvedPermissionMode, RuntimeConfig, RuntimeFeatureConfig, ScopedMcpServerConfig,
|
||||
CLAUDE_CODE_SETTINGS_SCHEMA_NAME,
|
||||
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::{
|
||||
ApiClient, ApiRequest, AssistantEvent, ConversationRuntime, RuntimeError, StaticToolExecutor,
|
||||
ToolError, ToolExecutor, TurnSummary,
|
||||
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::{
|
||||
McpClaudeAiProxyTransport, McpClientAuth, McpClientBootstrap, McpClientTransport,
|
||||
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, McpInitializeClientInfo, McpInitializeParams, McpInitializeResult,
|
||||
McpInitializeServerInfo, McpListResourcesParams, McpListResourcesResult, McpListToolsParams,
|
||||
McpListToolsResult, McpReadResourceParams, McpReadResourceResult, McpResource,
|
||||
McpResourceContents, McpServerManager, McpServerManagerError, McpStdioProcess, McpTool,
|
||||
McpToolCallContent, McpToolCallParams, McpToolCallResult, UnsupportedMcpServer,
|
||||
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,
|
||||
@@ -61,24 +116,59 @@ pub use oauth::{
|
||||
PkceChallengeMethod, PkceCodePair,
|
||||
};
|
||||
pub use permissions::{
|
||||
PermissionMode, PermissionOutcome, PermissionPolicy, PermissionPromptDecision,
|
||||
PermissionPrompter, PermissionRequest,
|
||||
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 session::{
|
||||
ContentBlock, ConversationMessage, MessageRole, Session, SessionError, SessionMetadata,
|
||||
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, ()> {
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -73,7 +73,7 @@ pub fn mcp_server_signature(config: &McpServerConfig) -> Option<String> {
|
||||
Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url)))
|
||||
}
|
||||
McpServerConfig::Ws(config) => Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url))),
|
||||
McpServerConfig::ClaudeAiProxy(config) => {
|
||||
McpServerConfig::ManagedProxy(config) => {
|
||||
Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url)))
|
||||
}
|
||||
McpServerConfig::Sdk(_) => None,
|
||||
@@ -84,10 +84,13 @@ pub fn mcp_server_signature(config: &McpServerConfig) -> Option<String> {
|
||||
pub fn scoped_mcp_config_hash(config: &ScopedMcpServerConfig) -> String {
|
||||
let rendered = match &config.config {
|
||||
McpServerConfig::Stdio(stdio) => format!(
|
||||
"stdio|{}|{}|{}",
|
||||
"stdio|{}|{}|{}|{}",
|
||||
stdio.command,
|
||||
render_command_signature(&stdio.args),
|
||||
render_env_signature(&stdio.env)
|
||||
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|{}|{}|{}|{}",
|
||||
@@ -110,7 +113,7 @@ pub fn scoped_mcp_config_hash(config: &ScopedMcpServerConfig) -> String {
|
||||
ws.headers_helper.as_deref().unwrap_or("")
|
||||
),
|
||||
McpServerConfig::Sdk(sdk) => format!("sdk|{}", sdk.name),
|
||||
McpServerConfig::ClaudeAiProxy(proxy) => {
|
||||
McpServerConfig::ManagedProxy(proxy) => {
|
||||
format!("claudeai-proxy|{}|{}", proxy.url, proxy.id)
|
||||
}
|
||||
};
|
||||
@@ -245,6 +248,7 @@ mod tests {
|
||||
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),
|
||||
|
||||
@@ -3,6 +3,8 @@ 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),
|
||||
@@ -10,7 +12,7 @@ pub enum McpClientTransport {
|
||||
Http(McpRemoteTransport),
|
||||
WebSocket(McpRemoteTransport),
|
||||
Sdk(McpSdkTransport),
|
||||
ClaudeAiProxy(McpClaudeAiProxyTransport),
|
||||
ManagedProxy(McpManagedProxyTransport),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -18,6 +20,7 @@ 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)]
|
||||
@@ -34,7 +37,7 @@ pub struct McpSdkTransport {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct McpClaudeAiProxyTransport {
|
||||
pub struct McpManagedProxyTransport {
|
||||
pub url: String,
|
||||
pub id: String,
|
||||
}
|
||||
@@ -75,6 +78,7 @@ impl McpClientTransport {
|
||||
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(),
|
||||
@@ -97,16 +101,22 @@ impl McpClientTransport {
|
||||
McpServerConfig::Sdk(config) => Self::Sdk(McpSdkTransport {
|
||||
name: config.name.clone(),
|
||||
}),
|
||||
McpServerConfig::ClaudeAiProxy(config) => {
|
||||
Self::ClaudeAiProxy(McpClaudeAiProxyTransport {
|
||||
url: config.url.clone(),
|
||||
id: config.id.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 {
|
||||
@@ -138,6 +148,7 @@ mod tests {
|
||||
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),
|
||||
}),
|
||||
};
|
||||
|
||||
@@ -156,6 +167,7 @@ mod tests {
|
||||
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:?}"),
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
+1354
-123
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(all(test, unix))]
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ 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,
|
||||
@@ -17,6 +18,7 @@ pub struct OAuthTokenSet {
|
||||
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,
|
||||
@@ -24,6 +26,7 @@ pub struct PkceCodePair {
|
||||
pub challenge_method: PkceChallengeMethod,
|
||||
}
|
||||
|
||||
/// Challenge algorithms supported by the local PKCE helpers.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PkceChallengeMethod {
|
||||
S256,
|
||||
@@ -38,6 +41,7 @@ impl PkceChallengeMethod {
|
||||
}
|
||||
}
|
||||
|
||||
/// Parameters needed to build an authorization URL for browser-based login.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct OAuthAuthorizationRequest {
|
||||
pub authorize_url: String,
|
||||
@@ -50,6 +54,7 @@ pub struct OAuthAuthorizationRequest {
|
||||
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,
|
||||
@@ -60,6 +65,7 @@ pub struct OAuthTokenExchangeRequest {
|
||||
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,
|
||||
@@ -68,6 +74,7 @@ pub struct OAuthRefreshRequest {
|
||||
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>,
|
||||
@@ -324,12 +331,19 @@ fn generate_random_token(bytes: usize) -> io::Result<String> {
|
||||
}
|
||||
|
||||
fn credentials_home_dir() -> io::Result<PathBuf> {
|
||||
if let Some(path) = std::env::var_os("CLAUDE_CONFIG_HOME") {
|
||||
if let Some(path) = std::env::var_os("CLAW_CONFIG_HOME") {
|
||||
return Ok(PathBuf::from(path));
|
||||
}
|
||||
let home = std::env::var_os("HOME")
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "HOME is not set"))?;
|
||||
Ok(PathBuf::from(home).join(".claude"))
|
||||
.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>> {
|
||||
@@ -442,7 +456,7 @@ fn decode_hex(byte: u8) -> Result<u8, String> {
|
||||
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-encoding byte: {byte}")),
|
||||
_ => Err(format!("invalid percent byte: {byte}")),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -541,7 +555,7 @@ mod tests {
|
||||
fn oauth_credentials_round_trip_and_clear_preserves_other_fields() {
|
||||
let _guard = env_lock();
|
||||
let config_home = temp_config_home();
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", &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");
|
||||
@@ -567,7 +581,7 @@ mod tests {
|
||||
assert!(cleared.contains("\"other\": \"value\""));
|
||||
assert!(!cleared.contains("\"oauth\""));
|
||||
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
std::env::remove_var("CLAW_CONFIG_HOME");
|
||||
std::fs::remove_dir_all(config_home).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
|
||||
@@ -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:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,10 +1,17 @@
|
||||
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 {
|
||||
@@ -14,38 +21,87 @@ impl PermissionMode {
|
||||
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 {
|
||||
@@ -54,6 +110,9 @@ impl PermissionPolicy {
|
||||
Self {
|
||||
active_mode,
|
||||
tool_requirements: BTreeMap::new(),
|
||||
allow_rules: Vec::new(),
|
||||
deny_rules: Vec::new(),
|
||||
ask_rules: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -68,6 +127,26 @@ impl PermissionPolicy {
|
||||
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
|
||||
@@ -86,37 +165,121 @@ impl PermissionPolicy {
|
||||
&self,
|
||||
tool_name: &str,
|
||||
input: &str,
|
||||
mut prompter: Option<&mut dyn PermissionPrompter>,
|
||||
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);
|
||||
if current_mode >= required_mode {
|
||||
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;
|
||||
}
|
||||
|
||||
let request = PermissionRequest {
|
||||
tool_name: tool_name.to_string(),
|
||||
input: input.to_string(),
|
||||
current_mode,
|
||||
required_mode,
|
||||
};
|
||||
|
||||
if current_mode == PermissionMode::WorkspaceWrite
|
||||
&& required_mode == PermissionMode::DangerFullAccess
|
||||
if current_mode == PermissionMode::Prompt
|
||||
|| (current_mode == PermissionMode::WorkspaceWrite
|
||||
&& required_mode == PermissionMode::DangerFullAccess)
|
||||
{
|
||||
return match prompter.as_mut() {
|
||||
Some(prompter) => match prompter.decide(&request) {
|
||||
PermissionPromptDecision::Allow => PermissionOutcome::Allow,
|
||||
PermissionPromptDecision::Deny { reason } => PermissionOutcome::Deny { reason },
|
||||
},
|
||||
None => PermissionOutcome::Deny {
|
||||
reason: format!(
|
||||
"tool '{tool_name}' requires approval to escalate from {} to {}",
|
||||
current_mode.as_str(),
|
||||
required_mode.as_str()
|
||||
),
|
||||
},
|
||||
};
|
||||
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 {
|
||||
@@ -127,14 +290,191 @@ impl PermissionPolicy {
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
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::{
|
||||
PermissionMode, PermissionOutcome, PermissionPolicy, PermissionPromptDecision,
|
||||
PermissionPrompter, PermissionRequest,
|
||||
PermissionContext, PermissionMode, PermissionOutcome, PermissionOverride, PermissionPolicy,
|
||||
PermissionPromptDecision, PermissionPrompter, PermissionRequest,
|
||||
};
|
||||
use crate::config::RuntimePermissionRuleConfig;
|
||||
|
||||
struct RecordingPrompter {
|
||||
seen: Vec<PermissionRequest>,
|
||||
@@ -224,4 +564,120 @@ mod tests {
|
||||
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")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,533 @@
|
||||
#![allow(clippy::redundant_closure_for_method_calls)]
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::config::RuntimePluginConfig;
|
||||
use crate::mcp_tool_bridge::{McpResourceInfo, McpToolInfo};
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
pub type ToolInfo = McpToolInfo;
|
||||
pub type ResourceInfo = McpResourceInfo;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ServerStatus {
|
||||
Healthy,
|
||||
Degraded,
|
||||
Failed,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ServerStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Healthy => write!(f, "healthy"),
|
||||
Self::Degraded => write!(f, "degraded"),
|
||||
Self::Failed => write!(f, "failed"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ServerHealth {
|
||||
pub server_name: String,
|
||||
pub status: ServerStatus,
|
||||
pub capabilities: Vec<String>,
|
||||
pub last_error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case", tag = "state")]
|
||||
pub enum PluginState {
|
||||
Unconfigured,
|
||||
Validated,
|
||||
Starting,
|
||||
Healthy,
|
||||
Degraded {
|
||||
healthy_servers: Vec<String>,
|
||||
failed_servers: Vec<ServerHealth>,
|
||||
},
|
||||
Failed {
|
||||
reason: String,
|
||||
},
|
||||
ShuttingDown,
|
||||
Stopped,
|
||||
}
|
||||
|
||||
impl PluginState {
|
||||
#[must_use]
|
||||
pub fn from_servers(servers: &[ServerHealth]) -> Self {
|
||||
if servers.is_empty() {
|
||||
return Self::Failed {
|
||||
reason: "no servers available".to_string(),
|
||||
};
|
||||
}
|
||||
|
||||
let healthy_servers = servers
|
||||
.iter()
|
||||
.filter(|server| server.status != ServerStatus::Failed)
|
||||
.map(|server| server.server_name.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let failed_servers = servers
|
||||
.iter()
|
||||
.filter(|server| server.status == ServerStatus::Failed)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let has_degraded_server = servers
|
||||
.iter()
|
||||
.any(|server| server.status == ServerStatus::Degraded);
|
||||
|
||||
if failed_servers.is_empty() && !has_degraded_server {
|
||||
Self::Healthy
|
||||
} else if healthy_servers.is_empty() {
|
||||
Self::Failed {
|
||||
reason: format!("all {} servers failed", failed_servers.len()),
|
||||
}
|
||||
} else {
|
||||
Self::Degraded {
|
||||
healthy_servers,
|
||||
failed_servers,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for PluginState {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Unconfigured => write!(f, "unconfigured"),
|
||||
Self::Validated => write!(f, "validated"),
|
||||
Self::Starting => write!(f, "starting"),
|
||||
Self::Healthy => write!(f, "healthy"),
|
||||
Self::Degraded { .. } => write!(f, "degraded"),
|
||||
Self::Failed { .. } => write!(f, "failed"),
|
||||
Self::ShuttingDown => write!(f, "shutting_down"),
|
||||
Self::Stopped => write!(f, "stopped"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct PluginHealthcheck {
|
||||
pub plugin_name: String,
|
||||
pub state: PluginState,
|
||||
pub servers: Vec<ServerHealth>,
|
||||
pub last_check: u64,
|
||||
}
|
||||
|
||||
impl PluginHealthcheck {
|
||||
#[must_use]
|
||||
pub fn new(plugin_name: impl Into<String>, servers: Vec<ServerHealth>) -> Self {
|
||||
let state = PluginState::from_servers(&servers);
|
||||
Self {
|
||||
plugin_name: plugin_name.into(),
|
||||
state,
|
||||
servers,
|
||||
last_check: now_secs(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn degraded_mode(&self, discovery: &DiscoveryResult) -> Option<DegradedMode> {
|
||||
match &self.state {
|
||||
PluginState::Degraded {
|
||||
healthy_servers,
|
||||
failed_servers,
|
||||
} => Some(DegradedMode {
|
||||
available_tools: discovery
|
||||
.tools
|
||||
.iter()
|
||||
.map(|tool| tool.name.clone())
|
||||
.collect(),
|
||||
unavailable_tools: failed_servers
|
||||
.iter()
|
||||
.flat_map(|server| server.capabilities.iter().cloned())
|
||||
.collect(),
|
||||
reason: format!(
|
||||
"{} servers healthy, {} servers failed",
|
||||
healthy_servers.len(),
|
||||
failed_servers.len()
|
||||
),
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DiscoveryResult {
|
||||
pub tools: Vec<ToolInfo>,
|
||||
pub resources: Vec<ResourceInfo>,
|
||||
pub partial: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct DegradedMode {
|
||||
pub available_tools: Vec<String>,
|
||||
pub unavailable_tools: Vec<String>,
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
impl DegradedMode {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
available_tools: Vec<String>,
|
||||
unavailable_tools: Vec<String>,
|
||||
reason: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
available_tools,
|
||||
unavailable_tools,
|
||||
reason: reason.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum PluginLifecycleEvent {
|
||||
ConfigValidated,
|
||||
StartupHealthy,
|
||||
StartupDegraded,
|
||||
StartupFailed,
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for PluginLifecycleEvent {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::ConfigValidated => write!(f, "config_validated"),
|
||||
Self::StartupHealthy => write!(f, "startup_healthy"),
|
||||
Self::StartupDegraded => write!(f, "startup_degraded"),
|
||||
Self::StartupFailed => write!(f, "startup_failed"),
|
||||
Self::Shutdown => write!(f, "shutdown"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait PluginLifecycle {
|
||||
fn validate_config(&self, config: &RuntimePluginConfig) -> Result<(), String>;
|
||||
fn healthcheck(&self) -> PluginHealthcheck;
|
||||
fn discover(&self) -> DiscoveryResult;
|
||||
fn shutdown(&mut self) -> Result<(), String>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MockPluginLifecycle {
|
||||
plugin_name: String,
|
||||
valid_config: bool,
|
||||
healthcheck: PluginHealthcheck,
|
||||
discovery: DiscoveryResult,
|
||||
shutdown_error: Option<String>,
|
||||
shutdown_called: bool,
|
||||
}
|
||||
|
||||
impl MockPluginLifecycle {
|
||||
fn new(
|
||||
plugin_name: &str,
|
||||
valid_config: bool,
|
||||
servers: Vec<ServerHealth>,
|
||||
discovery: DiscoveryResult,
|
||||
shutdown_error: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
plugin_name: plugin_name.to_string(),
|
||||
valid_config,
|
||||
healthcheck: PluginHealthcheck::new(plugin_name, servers),
|
||||
discovery,
|
||||
shutdown_error,
|
||||
shutdown_called: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PluginLifecycle for MockPluginLifecycle {
|
||||
fn validate_config(&self, _config: &RuntimePluginConfig) -> Result<(), String> {
|
||||
if self.valid_config {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!(
|
||||
"plugin `{}` failed configuration validation",
|
||||
self.plugin_name
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn healthcheck(&self) -> PluginHealthcheck {
|
||||
if self.shutdown_called {
|
||||
PluginHealthcheck {
|
||||
plugin_name: self.plugin_name.clone(),
|
||||
state: PluginState::Stopped,
|
||||
servers: self.healthcheck.servers.clone(),
|
||||
last_check: now_secs(),
|
||||
}
|
||||
} else {
|
||||
self.healthcheck.clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn discover(&self) -> DiscoveryResult {
|
||||
self.discovery.clone()
|
||||
}
|
||||
|
||||
fn shutdown(&mut self) -> Result<(), String> {
|
||||
if let Some(error) = &self.shutdown_error {
|
||||
return Err(error.clone());
|
||||
}
|
||||
|
||||
self.shutdown_called = true;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn healthy_server(name: &str, capabilities: &[&str]) -> ServerHealth {
|
||||
ServerHealth {
|
||||
server_name: name.to_string(),
|
||||
status: ServerStatus::Healthy,
|
||||
capabilities: capabilities
|
||||
.iter()
|
||||
.map(|capability| capability.to_string())
|
||||
.collect(),
|
||||
last_error: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn failed_server(name: &str, capabilities: &[&str], error: &str) -> ServerHealth {
|
||||
ServerHealth {
|
||||
server_name: name.to_string(),
|
||||
status: ServerStatus::Failed,
|
||||
capabilities: capabilities
|
||||
.iter()
|
||||
.map(|capability| capability.to_string())
|
||||
.collect(),
|
||||
last_error: Some(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn degraded_server(name: &str, capabilities: &[&str], error: &str) -> ServerHealth {
|
||||
ServerHealth {
|
||||
server_name: name.to_string(),
|
||||
status: ServerStatus::Degraded,
|
||||
capabilities: capabilities
|
||||
.iter()
|
||||
.map(|capability| capability.to_string())
|
||||
.collect(),
|
||||
last_error: Some(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn tool(name: &str) -> ToolInfo {
|
||||
ToolInfo {
|
||||
name: name.to_string(),
|
||||
description: Some(format!("{name} tool")),
|
||||
input_schema: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resource(name: &str, uri: &str) -> ResourceInfo {
|
||||
ResourceInfo {
|
||||
uri: uri.to_string(),
|
||||
name: name.to_string(),
|
||||
description: Some(format!("{name} resource")),
|
||||
mime_type: Some("application/json".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn full_lifecycle_happy_path() {
|
||||
// given
|
||||
let mut lifecycle = MockPluginLifecycle::new(
|
||||
"healthy-plugin",
|
||||
true,
|
||||
vec![
|
||||
healthy_server("alpha", &["search", "read"]),
|
||||
healthy_server("beta", &["write"]),
|
||||
],
|
||||
DiscoveryResult {
|
||||
tools: vec![tool("search"), tool("read"), tool("write")],
|
||||
resources: vec![resource("docs", "file:///docs")],
|
||||
partial: false,
|
||||
},
|
||||
None,
|
||||
);
|
||||
let config = RuntimePluginConfig::default();
|
||||
|
||||
// when
|
||||
let validation = lifecycle.validate_config(&config);
|
||||
let healthcheck = lifecycle.healthcheck();
|
||||
let discovery = lifecycle.discover();
|
||||
let shutdown = lifecycle.shutdown();
|
||||
let post_shutdown = lifecycle.healthcheck();
|
||||
|
||||
// then
|
||||
assert_eq!(validation, Ok(()));
|
||||
assert_eq!(healthcheck.state, PluginState::Healthy);
|
||||
assert_eq!(healthcheck.plugin_name, "healthy-plugin");
|
||||
assert_eq!(discovery.tools.len(), 3);
|
||||
assert_eq!(discovery.resources.len(), 1);
|
||||
assert!(!discovery.partial);
|
||||
assert_eq!(shutdown, Ok(()));
|
||||
assert_eq!(post_shutdown.state, PluginState::Stopped);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degraded_startup_when_one_of_three_servers_fails() {
|
||||
// given
|
||||
let lifecycle = MockPluginLifecycle::new(
|
||||
"degraded-plugin",
|
||||
true,
|
||||
vec![
|
||||
healthy_server("alpha", &["search"]),
|
||||
failed_server("beta", &["write"], "connection refused"),
|
||||
healthy_server("gamma", &["read"]),
|
||||
],
|
||||
DiscoveryResult {
|
||||
tools: vec![tool("search"), tool("read")],
|
||||
resources: vec![resource("alpha-docs", "file:///alpha")],
|
||||
partial: true,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let healthcheck = lifecycle.healthcheck();
|
||||
let discovery = lifecycle.discover();
|
||||
let degraded_mode = healthcheck
|
||||
.degraded_mode(&discovery)
|
||||
.expect("degraded startup should expose degraded mode");
|
||||
|
||||
// then
|
||||
match healthcheck.state {
|
||||
PluginState::Degraded {
|
||||
healthy_servers,
|
||||
failed_servers,
|
||||
} => {
|
||||
assert_eq!(
|
||||
healthy_servers,
|
||||
vec!["alpha".to_string(), "gamma".to_string()]
|
||||
);
|
||||
assert_eq!(failed_servers.len(), 1);
|
||||
assert_eq!(failed_servers[0].server_name, "beta");
|
||||
assert_eq!(
|
||||
failed_servers[0].last_error.as_deref(),
|
||||
Some("connection refused")
|
||||
);
|
||||
}
|
||||
other => panic!("expected degraded state, got {other:?}"),
|
||||
}
|
||||
assert!(discovery.partial);
|
||||
assert_eq!(
|
||||
degraded_mode.available_tools,
|
||||
vec!["search".to_string(), "read".to_string()]
|
||||
);
|
||||
assert_eq!(degraded_mode.unavailable_tools, vec!["write".to_string()]);
|
||||
assert_eq!(degraded_mode.reason, "2 servers healthy, 1 servers failed");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degraded_server_status_keeps_server_usable() {
|
||||
// given
|
||||
let lifecycle = MockPluginLifecycle::new(
|
||||
"soft-degraded-plugin",
|
||||
true,
|
||||
vec![
|
||||
healthy_server("alpha", &["search"]),
|
||||
degraded_server("beta", &["write"], "high latency"),
|
||||
],
|
||||
DiscoveryResult {
|
||||
tools: vec![tool("search"), tool("write")],
|
||||
resources: Vec::new(),
|
||||
partial: true,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let healthcheck = lifecycle.healthcheck();
|
||||
|
||||
// then
|
||||
match healthcheck.state {
|
||||
PluginState::Degraded {
|
||||
healthy_servers,
|
||||
failed_servers,
|
||||
} => {
|
||||
assert_eq!(
|
||||
healthy_servers,
|
||||
vec!["alpha".to_string(), "beta".to_string()]
|
||||
);
|
||||
assert!(failed_servers.is_empty());
|
||||
}
|
||||
other => panic!("expected degraded state, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_failure_when_all_servers_fail() {
|
||||
// given
|
||||
let lifecycle = MockPluginLifecycle::new(
|
||||
"failed-plugin",
|
||||
true,
|
||||
vec![
|
||||
failed_server("alpha", &["search"], "timeout"),
|
||||
failed_server("beta", &["read"], "handshake failed"),
|
||||
],
|
||||
DiscoveryResult {
|
||||
tools: Vec::new(),
|
||||
resources: Vec::new(),
|
||||
partial: false,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let healthcheck = lifecycle.healthcheck();
|
||||
let discovery = lifecycle.discover();
|
||||
|
||||
// then
|
||||
match &healthcheck.state {
|
||||
PluginState::Failed { reason } => {
|
||||
assert_eq!(reason, "all 2 servers failed");
|
||||
}
|
||||
other => panic!("expected failed state, got {other:?}"),
|
||||
}
|
||||
assert!(!discovery.partial);
|
||||
assert!(discovery.tools.is_empty());
|
||||
assert!(discovery.resources.is_empty());
|
||||
assert!(healthcheck.degraded_mode(&discovery).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graceful_shutdown() {
|
||||
// given
|
||||
let mut lifecycle = MockPluginLifecycle::new(
|
||||
"shutdown-plugin",
|
||||
true,
|
||||
vec![healthy_server("alpha", &["search"])],
|
||||
DiscoveryResult {
|
||||
tools: vec![tool("search")],
|
||||
resources: Vec::new(),
|
||||
partial: false,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let shutdown = lifecycle.shutdown();
|
||||
let post_shutdown = lifecycle.healthcheck();
|
||||
|
||||
// then
|
||||
assert_eq!(shutdown, Ok(()));
|
||||
assert_eq!(PluginLifecycleEvent::Shutdown.to_string(), "shutdown");
|
||||
assert_eq!(post_shutdown.state, PluginState::Stopped);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,581 @@
|
||||
use std::time::Duration;
|
||||
|
||||
pub type GreenLevel = u8;
|
||||
|
||||
const STALE_BRANCH_THRESHOLD: Duration = Duration::from_secs(60 * 60);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PolicyRule {
|
||||
pub name: String,
|
||||
pub condition: PolicyCondition,
|
||||
pub action: PolicyAction,
|
||||
pub priority: u32,
|
||||
}
|
||||
|
||||
impl PolicyRule {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
name: impl Into<String>,
|
||||
condition: PolicyCondition,
|
||||
action: PolicyAction,
|
||||
priority: u32,
|
||||
) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
condition,
|
||||
action,
|
||||
priority,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn matches(&self, context: &LaneContext) -> bool {
|
||||
self.condition.matches(context)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PolicyCondition {
|
||||
And(Vec<PolicyCondition>),
|
||||
Or(Vec<PolicyCondition>),
|
||||
GreenAt { level: GreenLevel },
|
||||
StaleBranch,
|
||||
StartupBlocked,
|
||||
LaneCompleted,
|
||||
LaneReconciled,
|
||||
ReviewPassed,
|
||||
ScopedDiff,
|
||||
TimedOut { duration: Duration },
|
||||
}
|
||||
|
||||
impl PolicyCondition {
|
||||
#[must_use]
|
||||
pub fn matches(&self, context: &LaneContext) -> bool {
|
||||
match self {
|
||||
Self::And(conditions) => conditions
|
||||
.iter()
|
||||
.all(|condition| condition.matches(context)),
|
||||
Self::Or(conditions) => conditions
|
||||
.iter()
|
||||
.any(|condition| condition.matches(context)),
|
||||
Self::GreenAt { level } => context.green_level >= *level,
|
||||
Self::StaleBranch => context.branch_freshness >= STALE_BRANCH_THRESHOLD,
|
||||
Self::StartupBlocked => context.blocker == LaneBlocker::Startup,
|
||||
Self::LaneCompleted => context.completed,
|
||||
Self::LaneReconciled => context.reconciled,
|
||||
Self::ReviewPassed => context.review_status == ReviewStatus::Approved,
|
||||
Self::ScopedDiff => context.diff_scope == DiffScope::Scoped,
|
||||
Self::TimedOut { duration } => context.branch_freshness >= *duration,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PolicyAction {
|
||||
MergeToDev,
|
||||
MergeForward,
|
||||
RecoverOnce,
|
||||
Escalate { reason: String },
|
||||
CloseoutLane,
|
||||
CleanupSession,
|
||||
Reconcile { reason: ReconcileReason },
|
||||
Notify { channel: String },
|
||||
Block { reason: String },
|
||||
Chain(Vec<PolicyAction>),
|
||||
}
|
||||
|
||||
/// Why a lane was reconciled without further action.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ReconcileReason {
|
||||
/// Branch already merged into main — no PR needed.
|
||||
AlreadyMerged,
|
||||
/// Work superseded by another lane or direct commit.
|
||||
Superseded,
|
||||
/// PR would be empty — all changes already landed.
|
||||
EmptyDiff,
|
||||
/// Lane manually closed by operator.
|
||||
ManualClose,
|
||||
}
|
||||
|
||||
impl PolicyAction {
|
||||
fn flatten_into(&self, actions: &mut Vec<PolicyAction>) {
|
||||
match self {
|
||||
Self::Chain(chained) => {
|
||||
for action in chained {
|
||||
action.flatten_into(actions);
|
||||
}
|
||||
}
|
||||
_ => actions.push(self.clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LaneBlocker {
|
||||
None,
|
||||
Startup,
|
||||
External,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ReviewStatus {
|
||||
Pending,
|
||||
Approved,
|
||||
Rejected,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum DiffScope {
|
||||
Full,
|
||||
Scoped,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LaneContext {
|
||||
pub lane_id: String,
|
||||
pub green_level: GreenLevel,
|
||||
pub branch_freshness: Duration,
|
||||
pub blocker: LaneBlocker,
|
||||
pub review_status: ReviewStatus,
|
||||
pub diff_scope: DiffScope,
|
||||
pub completed: bool,
|
||||
pub reconciled: bool,
|
||||
}
|
||||
|
||||
impl LaneContext {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
lane_id: impl Into<String>,
|
||||
green_level: GreenLevel,
|
||||
branch_freshness: Duration,
|
||||
blocker: LaneBlocker,
|
||||
review_status: ReviewStatus,
|
||||
diff_scope: DiffScope,
|
||||
completed: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
lane_id: lane_id.into(),
|
||||
green_level,
|
||||
branch_freshness,
|
||||
blocker,
|
||||
review_status,
|
||||
diff_scope,
|
||||
completed,
|
||||
reconciled: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a lane context that is already reconciled (no further action needed).
|
||||
#[must_use]
|
||||
pub fn reconciled(lane_id: impl Into<String>) -> Self {
|
||||
Self {
|
||||
lane_id: lane_id.into(),
|
||||
green_level: 0,
|
||||
branch_freshness: Duration::from_secs(0),
|
||||
blocker: LaneBlocker::None,
|
||||
review_status: ReviewStatus::Pending,
|
||||
diff_scope: DiffScope::Full,
|
||||
completed: true,
|
||||
reconciled: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PolicyEngine {
|
||||
rules: Vec<PolicyRule>,
|
||||
}
|
||||
|
||||
impl PolicyEngine {
|
||||
#[must_use]
|
||||
pub fn new(mut rules: Vec<PolicyRule>) -> Self {
|
||||
rules.sort_by_key(|rule| rule.priority);
|
||||
Self { rules }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn rules(&self) -> &[PolicyRule] {
|
||||
&self.rules
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn evaluate(&self, context: &LaneContext) -> Vec<PolicyAction> {
|
||||
evaluate(self, context)
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn evaluate(engine: &PolicyEngine, context: &LaneContext) -> Vec<PolicyAction> {
|
||||
let mut actions = Vec::new();
|
||||
for rule in &engine.rules {
|
||||
if rule.matches(context) {
|
||||
rule.action.flatten_into(&mut actions);
|
||||
}
|
||||
}
|
||||
actions
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{
|
||||
evaluate, DiffScope, LaneBlocker, LaneContext, PolicyAction, PolicyCondition, PolicyEngine,
|
||||
PolicyRule, ReconcileReason, ReviewStatus, STALE_BRANCH_THRESHOLD,
|
||||
};
|
||||
|
||||
fn default_context() -> LaneContext {
|
||||
LaneContext::new(
|
||||
"lane-7",
|
||||
0,
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_to_dev_rule_fires_for_green_scoped_reviewed_lane() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"merge-to-dev",
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::GreenAt { level: 2 },
|
||||
PolicyCondition::ScopedDiff,
|
||||
PolicyCondition::ReviewPassed,
|
||||
]),
|
||||
PolicyAction::MergeToDev,
|
||||
20,
|
||||
)]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
3,
|
||||
Duration::from_secs(5),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Approved,
|
||||
DiffScope::Scoped,
|
||||
false,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(actions, vec![PolicyAction::MergeToDev]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_rule_fires_at_threshold() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"merge-forward",
|
||||
PolicyCondition::StaleBranch,
|
||||
PolicyAction::MergeForward,
|
||||
10,
|
||||
)]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
1,
|
||||
STALE_BRANCH_THRESHOLD,
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(actions, vec![PolicyAction::MergeForward]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_blocked_rule_recovers_then_escalates() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"startup-recovery",
|
||||
PolicyCondition::StartupBlocked,
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::RecoverOnce,
|
||||
PolicyAction::Escalate {
|
||||
reason: "startup remained blocked".to_string(),
|
||||
},
|
||||
]),
|
||||
15,
|
||||
)]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
0,
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::Startup,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::RecoverOnce,
|
||||
PolicyAction::Escalate {
|
||||
reason: "startup remained blocked".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completed_lane_rule_closes_out_and_cleans_up() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"lane-closeout",
|
||||
PolicyCondition::LaneCompleted,
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::CloseoutLane,
|
||||
PolicyAction::CleanupSession,
|
||||
]),
|
||||
30,
|
||||
)]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
0,
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
true,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![PolicyAction::CloseoutLane, PolicyAction::CleanupSession]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matching_rules_are_returned_in_priority_order_with_stable_ties() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![
|
||||
PolicyRule::new(
|
||||
"late-cleanup",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::CleanupSession,
|
||||
30,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"first-notify",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::Notify {
|
||||
channel: "ops".to_string(),
|
||||
},
|
||||
10,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"second-notify",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::Notify {
|
||||
channel: "review".to_string(),
|
||||
},
|
||||
10,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"merge",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::MergeToDev,
|
||||
20,
|
||||
),
|
||||
]);
|
||||
let context = default_context();
|
||||
|
||||
// when
|
||||
let actions = evaluate(&engine, &context);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::Notify {
|
||||
channel: "ops".to_string(),
|
||||
},
|
||||
PolicyAction::Notify {
|
||||
channel: "review".to_string(),
|
||||
},
|
||||
PolicyAction::MergeToDev,
|
||||
PolicyAction::CleanupSession,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn combinators_handle_empty_cases_and_nested_chains() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![
|
||||
PolicyRule::new(
|
||||
"empty-and",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::Notify {
|
||||
channel: "orchestrator".to_string(),
|
||||
},
|
||||
5,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"empty-or",
|
||||
PolicyCondition::Or(vec![]),
|
||||
PolicyAction::Block {
|
||||
reason: "should not fire".to_string(),
|
||||
},
|
||||
10,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"nested",
|
||||
PolicyCondition::Or(vec![
|
||||
PolicyCondition::StartupBlocked,
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::GreenAt { level: 2 },
|
||||
PolicyCondition::TimedOut {
|
||||
duration: Duration::from_secs(5),
|
||||
},
|
||||
]),
|
||||
]),
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::Notify {
|
||||
channel: "alerts".to_string(),
|
||||
},
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::MergeForward,
|
||||
PolicyAction::CleanupSession,
|
||||
]),
|
||||
]),
|
||||
15,
|
||||
),
|
||||
]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
2,
|
||||
Duration::from_secs(10),
|
||||
LaneBlocker::External,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::Notify {
|
||||
channel: "orchestrator".to_string(),
|
||||
},
|
||||
PolicyAction::Notify {
|
||||
channel: "alerts".to_string(),
|
||||
},
|
||||
PolicyAction::MergeForward,
|
||||
PolicyAction::CleanupSession,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconciled_lane_emits_reconcile_and_cleanup() {
|
||||
// given — a lane where branch is already merged, no PR needed, session stale
|
||||
let engine = PolicyEngine::new(vec![
|
||||
PolicyRule::new(
|
||||
"reconcile-closeout",
|
||||
PolicyCondition::LaneReconciled,
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::Reconcile {
|
||||
reason: ReconcileReason::AlreadyMerged,
|
||||
},
|
||||
PolicyAction::CloseoutLane,
|
||||
PolicyAction::CleanupSession,
|
||||
]),
|
||||
5,
|
||||
),
|
||||
// This rule should NOT fire — reconciled lanes are completed but we want
|
||||
// the more specific reconcile rule to handle them
|
||||
PolicyRule::new(
|
||||
"generic-closeout",
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::LaneCompleted,
|
||||
// Only fire if NOT reconciled
|
||||
PolicyCondition::And(vec![]),
|
||||
]),
|
||||
PolicyAction::CloseoutLane,
|
||||
30,
|
||||
),
|
||||
]);
|
||||
let context = LaneContext::reconciled("lane-9411");
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then — reconcile rule fires first (priority 5), then generic closeout also fires
|
||||
// because reconciled context has completed=true
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::Reconcile {
|
||||
reason: ReconcileReason::AlreadyMerged,
|
||||
},
|
||||
PolicyAction::CloseoutLane,
|
||||
PolicyAction::CleanupSession,
|
||||
PolicyAction::CloseoutLane,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconciled_context_has_correct_defaults() {
|
||||
let ctx = LaneContext::reconciled("test-lane");
|
||||
assert_eq!(ctx.lane_id, "test-lane");
|
||||
assert!(ctx.completed);
|
||||
assert!(ctx.reconciled);
|
||||
assert_eq!(ctx.blocker, LaneBlocker::None);
|
||||
assert_eq!(ctx.green_level, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_reconciled_lane_does_not_trigger_reconcile_rule() {
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"reconcile-closeout",
|
||||
PolicyCondition::LaneReconciled,
|
||||
PolicyAction::Reconcile {
|
||||
reason: ReconcileReason::EmptyDiff,
|
||||
},
|
||||
5,
|
||||
)]);
|
||||
// Normal completed lane — not reconciled
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
0,
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
true,
|
||||
);
|
||||
|
||||
let actions = engine.evaluate(&context);
|
||||
assert!(actions.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconcile_reason_variants_are_distinct() {
|
||||
assert_ne!(ReconcileReason::AlreadyMerged, ReconcileReason::Superseded);
|
||||
assert_ne!(ReconcileReason::EmptyDiff, ReconcileReason::ManualClose);
|
||||
}
|
||||
}
|
||||
+230
-100
@@ -4,7 +4,9 @@ use std::path::{Path, PathBuf};
|
||||
use std::process::Command;
|
||||
|
||||
use crate::config::{ConfigError, ConfigLoader, RuntimeConfig};
|
||||
use crate::git_context::GitContext;
|
||||
|
||||
/// Errors raised while assembling the final system prompt.
|
||||
#[derive(Debug)]
|
||||
pub enum PromptBuildError {
|
||||
Io(std::io::Error),
|
||||
@@ -34,24 +36,29 @@ impl From<ConfigError> for PromptBuildError {
|
||||
}
|
||||
}
|
||||
|
||||
/// Marker separating static prompt scaffolding from dynamic runtime context.
|
||||
pub const SYSTEM_PROMPT_DYNAMIC_BOUNDARY: &str = "__SYSTEM_PROMPT_DYNAMIC_BOUNDARY__";
|
||||
/// Human-readable default frontier model name embedded into generated prompts.
|
||||
pub const FRONTIER_MODEL_NAME: &str = "Claude Opus 4.6";
|
||||
const MAX_INSTRUCTION_FILE_CHARS: usize = 4_000;
|
||||
const MAX_TOTAL_INSTRUCTION_CHARS: usize = 12_000;
|
||||
|
||||
/// Contents of an instruction file included in prompt construction.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ContextFile {
|
||||
pub path: PathBuf,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
/// Project-local context injected into the rendered system prompt.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct ProjectContext {
|
||||
pub cwd: PathBuf,
|
||||
pub current_date: String,
|
||||
pub git_status: Option<String>,
|
||||
pub git_diff: Option<String>,
|
||||
pub git_context: Option<GitContext>,
|
||||
pub instruction_files: Vec<ContextFile>,
|
||||
pub memory_files: Vec<ContextFile>,
|
||||
}
|
||||
|
||||
impl ProjectContext {
|
||||
@@ -61,13 +68,13 @@ impl ProjectContext {
|
||||
) -> std::io::Result<Self> {
|
||||
let cwd = cwd.into();
|
||||
let instruction_files = discover_instruction_files(&cwd)?;
|
||||
let memory_files = discover_memory_files(&cwd)?;
|
||||
Ok(Self {
|
||||
cwd,
|
||||
current_date: current_date.into(),
|
||||
git_status: None,
|
||||
git_diff: None,
|
||||
git_context: None,
|
||||
instruction_files,
|
||||
memory_files,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -77,10 +84,13 @@ impl ProjectContext {
|
||||
) -> std::io::Result<Self> {
|
||||
let mut context = Self::discover(cwd, current_date)?;
|
||||
context.git_status = read_git_status(&context.cwd);
|
||||
context.git_diff = read_git_diff(&context.cwd);
|
||||
context.git_context = GitContext::detect(&context.cwd);
|
||||
Ok(context)
|
||||
}
|
||||
}
|
||||
|
||||
/// Builder for the runtime system prompt and dynamic environment sections.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct SystemPromptBuilder {
|
||||
output_style_name: Option<String>,
|
||||
@@ -147,9 +157,6 @@ impl SystemPromptBuilder {
|
||||
if !project_context.instruction_files.is_empty() {
|
||||
sections.push(render_instruction_files(&project_context.instruction_files));
|
||||
}
|
||||
if !project_context.memory_files.is_empty() {
|
||||
sections.push(render_memory_files(&project_context.memory_files));
|
||||
}
|
||||
}
|
||||
if let Some(config) = &self.config {
|
||||
sections.push(render_config_section(config));
|
||||
@@ -187,12 +194,13 @@ impl SystemPromptBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
/// Formats each item as an indented bullet for prompt sections.
|
||||
#[must_use]
|
||||
pub fn prepend_bullets(items: Vec<String>) -> Vec<String> {
|
||||
items.into_iter().map(|item| format!(" - {item}")).collect()
|
||||
}
|
||||
|
||||
fn discover_context_directories(cwd: &Path) -> Vec<PathBuf> {
|
||||
fn discover_instruction_files(cwd: &Path) -> std::io::Result<Vec<ContextFile>> {
|
||||
let mut directories = Vec::new();
|
||||
let mut cursor = Some(cwd);
|
||||
while let Some(dir) = cursor {
|
||||
@@ -200,19 +208,14 @@ fn discover_context_directories(cwd: &Path) -> Vec<PathBuf> {
|
||||
cursor = dir.parent();
|
||||
}
|
||||
directories.reverse();
|
||||
directories
|
||||
}
|
||||
|
||||
fn discover_instruction_files(cwd: &Path) -> std::io::Result<Vec<ContextFile>> {
|
||||
let directories = discover_context_directories(cwd);
|
||||
|
||||
let mut files = Vec::new();
|
||||
for dir in directories {
|
||||
for candidate in [
|
||||
dir.join("CLAUDE.md"),
|
||||
dir.join("CLAUDE.local.md"),
|
||||
dir.join(".claude").join("CLAUDE.md"),
|
||||
dir.join(".claude").join("instructions.md"),
|
||||
dir.join(".claw").join("CLAUDE.md"),
|
||||
dir.join(".claw").join("instructions.md"),
|
||||
] {
|
||||
push_context_file(&mut files, candidate)?;
|
||||
}
|
||||
@@ -220,26 +223,6 @@ fn discover_instruction_files(cwd: &Path) -> std::io::Result<Vec<ContextFile>> {
|
||||
Ok(dedupe_instruction_files(files))
|
||||
}
|
||||
|
||||
fn discover_memory_files(cwd: &Path) -> std::io::Result<Vec<ContextFile>> {
|
||||
let mut files = Vec::new();
|
||||
for dir in discover_context_directories(cwd) {
|
||||
let memory_dir = dir.join(".claude").join("memory");
|
||||
let Ok(entries) = fs::read_dir(&memory_dir) else {
|
||||
continue;
|
||||
};
|
||||
let mut paths = entries
|
||||
.flatten()
|
||||
.map(|entry| entry.path())
|
||||
.filter(|path| path.is_file())
|
||||
.collect::<Vec<_>>();
|
||||
paths.sort();
|
||||
for path in paths {
|
||||
push_context_file(&mut files, path)?;
|
||||
}
|
||||
}
|
||||
Ok(dedupe_instruction_files(files))
|
||||
}
|
||||
|
||||
fn push_context_file(files: &mut Vec<ContextFile>, path: PathBuf) -> std::io::Result<()> {
|
||||
match fs::read_to_string(&path) {
|
||||
Ok(content) if !content.trim().is_empty() => {
|
||||
@@ -270,6 +253,38 @@ fn read_git_status(cwd: &Path) -> Option<String> {
|
||||
}
|
||||
}
|
||||
|
||||
fn read_git_diff(cwd: &Path) -> Option<String> {
|
||||
let mut sections = Vec::new();
|
||||
|
||||
let staged = read_git_output(cwd, &["diff", "--cached"])?;
|
||||
if !staged.trim().is_empty() {
|
||||
sections.push(format!("Staged changes:\n{}", staged.trim_end()));
|
||||
}
|
||||
|
||||
let unstaged = read_git_output(cwd, &["diff"])?;
|
||||
if !unstaged.trim().is_empty() {
|
||||
sections.push(format!("Unstaged changes:\n{}", unstaged.trim_end()));
|
||||
}
|
||||
|
||||
if sections.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(sections.join("\n\n"))
|
||||
}
|
||||
}
|
||||
|
||||
fn read_git_output(cwd: &Path, args: &[&str]) -> Option<String> {
|
||||
let output = Command::new("git")
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.ok()?;
|
||||
if !output.status.success() {
|
||||
return None;
|
||||
}
|
||||
String::from_utf8(output.stdout).ok()
|
||||
}
|
||||
|
||||
fn render_project_context(project_context: &ProjectContext) -> String {
|
||||
let mut lines = vec!["# Project context".to_string()];
|
||||
let mut bullets = vec![
|
||||
@@ -282,31 +297,38 @@ fn render_project_context(project_context: &ProjectContext) -> String {
|
||||
project_context.instruction_files.len()
|
||||
));
|
||||
}
|
||||
if !project_context.memory_files.is_empty() {
|
||||
bullets.push(format!(
|
||||
"Project memory files discovered: {}.",
|
||||
project_context.memory_files.len()
|
||||
));
|
||||
}
|
||||
lines.extend(prepend_bullets(bullets));
|
||||
if let Some(status) = &project_context.git_status {
|
||||
lines.push(String::new());
|
||||
lines.push("Git status snapshot:".to_string());
|
||||
lines.push(status.clone());
|
||||
}
|
||||
if let Some(ref gc) = project_context.git_context {
|
||||
if !gc.recent_commits.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("Recent commits (last 5):".to_string());
|
||||
for c in &gc.recent_commits {
|
||||
lines.push(format!(" {} {}", c.hash, c.subject));
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(diff) = &project_context.git_diff {
|
||||
lines.push(String::new());
|
||||
lines.push("Git diff snapshot:".to_string());
|
||||
lines.push(diff.clone());
|
||||
}
|
||||
if let Some(git_context) = &project_context.git_context {
|
||||
let rendered = git_context.render();
|
||||
if !rendered.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push(rendered);
|
||||
}
|
||||
}
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
fn render_instruction_files(files: &[ContextFile]) -> String {
|
||||
render_context_file_section("# Claude instructions", files)
|
||||
}
|
||||
|
||||
fn render_memory_files(files: &[ContextFile]) -> String {
|
||||
render_context_file_section("# Project memory", files)
|
||||
}
|
||||
|
||||
fn render_context_file_section(title: &str, files: &[ContextFile]) -> String {
|
||||
let mut sections = vec![title.to_string()];
|
||||
let mut sections = vec!["# Claude instructions".to_string()];
|
||||
let mut remaining_chars = MAX_TOTAL_INSTRUCTION_CHARS;
|
||||
for file in files {
|
||||
if remaining_chars == 0 {
|
||||
@@ -406,6 +428,7 @@ fn collapse_blank_lines(content: &str) -> String {
|
||||
result
|
||||
}
|
||||
|
||||
/// Loads config and project context, then renders the system prompt text.
|
||||
pub fn load_system_prompt(
|
||||
cwd: impl Into<PathBuf>,
|
||||
current_date: impl Into<String>,
|
||||
@@ -426,7 +449,7 @@ fn render_config_section(config: &RuntimeConfig) -> String {
|
||||
let mut lines = vec!["# Runtime config".to_string()];
|
||||
if config.loaded_entries().is_empty() {
|
||||
lines.extend(prepend_bullets(vec![
|
||||
"No Claude Code settings files loaded.".to_string(),
|
||||
"No Claw Code settings files loaded.".to_string()
|
||||
]));
|
||||
return lines.join("\n");
|
||||
}
|
||||
@@ -498,9 +521,8 @@ fn get_actions_section() -> String {
|
||||
mod tests {
|
||||
use super::{
|
||||
collapse_blank_lines, display_context_path, normalize_instruction_content,
|
||||
render_instruction_content, render_instruction_files, render_memory_files,
|
||||
truncate_instruction_content, ContextFile, ProjectContext, SystemPromptBuilder,
|
||||
SYSTEM_PROMPT_DYNAMIC_BOUNDARY,
|
||||
render_instruction_content, render_instruction_files, truncate_instruction_content,
|
||||
ContextFile, ProjectContext, SystemPromptBuilder, SYSTEM_PROMPT_DYNAMIC_BOUNDARY,
|
||||
};
|
||||
use crate::config::ConfigLoader;
|
||||
use std::fs;
|
||||
@@ -519,27 +541,34 @@ mod tests {
|
||||
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 discovers_instruction_files_from_ancestor_chain() {
|
||||
let root = temp_dir();
|
||||
let nested = root.join("apps").join("api");
|
||||
fs::create_dir_all(nested.join(".claude")).expect("nested claude dir");
|
||||
fs::create_dir_all(nested.join(".claw")).expect("nested claw dir");
|
||||
fs::write(root.join("CLAUDE.md"), "root instructions").expect("write root instructions");
|
||||
fs::write(root.join("CLAUDE.local.md"), "local instructions")
|
||||
.expect("write local instructions");
|
||||
fs::create_dir_all(root.join("apps")).expect("apps dir");
|
||||
fs::create_dir_all(root.join("apps").join(".claude")).expect("apps claude dir");
|
||||
fs::create_dir_all(root.join("apps").join(".claw")).expect("apps claw dir");
|
||||
fs::write(root.join("apps").join("CLAUDE.md"), "apps instructions")
|
||||
.expect("write apps instructions");
|
||||
fs::write(
|
||||
root.join("apps").join(".claude").join("instructions.md"),
|
||||
root.join("apps").join(".claw").join("instructions.md"),
|
||||
"apps dot claude instructions",
|
||||
)
|
||||
.expect("write apps dot claude instructions");
|
||||
fs::write(nested.join(".claude").join("CLAUDE.md"), "nested rules")
|
||||
fs::write(nested.join(".claw").join("CLAUDE.md"), "nested rules")
|
||||
.expect("write nested rules");
|
||||
fs::write(
|
||||
nested.join(".claude").join("instructions.md"),
|
||||
nested.join(".claw").join("instructions.md"),
|
||||
"nested instructions",
|
||||
)
|
||||
.expect("write nested instructions");
|
||||
@@ -565,35 +594,6 @@ mod tests {
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discovers_project_memory_files_from_ancestor_chain() {
|
||||
let root = temp_dir();
|
||||
let nested = root.join("apps").join("api");
|
||||
fs::create_dir_all(root.join(".claude").join("memory")).expect("root memory dir");
|
||||
fs::create_dir_all(nested.join(".claude").join("memory")).expect("nested memory dir");
|
||||
fs::write(
|
||||
root.join(".claude").join("memory").join("2026-03-30.md"),
|
||||
"root memory",
|
||||
)
|
||||
.expect("write root memory");
|
||||
fs::write(
|
||||
nested.join(".claude").join("memory").join("2026-03-31.md"),
|
||||
"nested memory",
|
||||
)
|
||||
.expect("write nested memory");
|
||||
|
||||
let context = ProjectContext::discover(&nested, "2026-03-31").expect("context should load");
|
||||
let contents = context
|
||||
.memory_files
|
||||
.iter()
|
||||
.map(|file| file.content.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(contents, vec!["root memory", "nested memory"]);
|
||||
assert!(render_memory_files(&context.memory_files).contains("# Project memory"));
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedupes_identical_instruction_content_across_scopes() {
|
||||
let root = temp_dir();
|
||||
@@ -628,13 +628,15 @@ mod tests {
|
||||
#[test]
|
||||
fn displays_context_paths_compactly() {
|
||||
assert_eq!(
|
||||
display_context_path(Path::new("/tmp/project/.claude/CLAUDE.md")),
|
||||
display_context_path(Path::new("/tmp/project/.claw/CLAUDE.md")),
|
||||
"CLAUDE.md"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_with_git_includes_status_snapshot() {
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir");
|
||||
std::process::Command::new("git")
|
||||
@@ -652,6 +654,133 @@ mod tests {
|
||||
assert!(status.contains("## No commits yet on") || status.contains("## "));
|
||||
assert!(status.contains("?? CLAUDE.md"));
|
||||
assert!(status.contains("?? tracked.txt"));
|
||||
assert!(context.git_diff.is_none());
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_with_git_includes_recent_commits_and_renders_them() {
|
||||
// given: a git repo with three commits and a current branch
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir");
|
||||
std::process::Command::new("git")
|
||||
.args(["init", "--quiet", "-b", "main"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git init should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["config", "user.email", "tests@example.com"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git config email should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["config", "user.name", "Runtime Prompt Tests"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git config name should run");
|
||||
for (file, message) in [
|
||||
("a.txt", "first commit"),
|
||||
("b.txt", "second commit"),
|
||||
("c.txt", "third commit"),
|
||||
] {
|
||||
fs::write(root.join(file), "x\n").expect("write commit file");
|
||||
std::process::Command::new("git")
|
||||
.args(["add", file])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git add should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["commit", "-m", message, "--quiet"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git commit should run");
|
||||
}
|
||||
fs::write(root.join("d.txt"), "staged\n").expect("write staged file");
|
||||
std::process::Command::new("git")
|
||||
.args(["add", "d.txt"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git add staged should run");
|
||||
|
||||
// when: discovering project context with git auto-include
|
||||
let context =
|
||||
ProjectContext::discover_with_git(&root, "2026-03-31").expect("context should load");
|
||||
let rendered = SystemPromptBuilder::new()
|
||||
.with_os("linux", "6.8")
|
||||
.with_project_context(context.clone())
|
||||
.render();
|
||||
|
||||
// then: branch, recent commits and staged files are present in context
|
||||
let gc = context
|
||||
.git_context
|
||||
.as_ref()
|
||||
.expect("git context should be present");
|
||||
let commits: String = gc
|
||||
.recent_commits
|
||||
.iter()
|
||||
.map(|c| c.subject.clone())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
assert!(commits.contains("first commit"));
|
||||
assert!(commits.contains("second commit"));
|
||||
assert!(commits.contains("third commit"));
|
||||
assert_eq!(gc.recent_commits.len(), 3);
|
||||
|
||||
let status = context.git_status.as_deref().expect("status snapshot");
|
||||
assert!(status.contains("## main"));
|
||||
assert!(status.contains("A d.txt"));
|
||||
|
||||
assert!(rendered.contains("Recent commits (last 5):"));
|
||||
assert!(rendered.contains("first commit"));
|
||||
assert!(rendered.contains("Git status snapshot:"));
|
||||
assert!(rendered.contains("## main"));
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_with_git_includes_diff_snapshot_for_tracked_changes() {
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir");
|
||||
std::process::Command::new("git")
|
||||
.args(["init", "--quiet"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git init should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["config", "user.email", "tests@example.com"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git config email should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["config", "user.name", "Runtime Prompt Tests"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git config name should run");
|
||||
fs::write(root.join("tracked.txt"), "hello\n").expect("write tracked file");
|
||||
std::process::Command::new("git")
|
||||
.args(["add", "tracked.txt"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git add should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["commit", "-m", "init", "--quiet"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git commit should run");
|
||||
fs::write(root.join("tracked.txt"), "hello\nworld\n").expect("rewrite tracked file");
|
||||
|
||||
let context =
|
||||
ProjectContext::discover_with_git(&root, "2026-03-31").expect("context should load");
|
||||
|
||||
let diff = context.git_diff.expect("git diff should be present");
|
||||
assert!(diff.contains("Unstaged changes:"));
|
||||
assert!(diff.contains("tracked.txt"));
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
@@ -659,20 +788,21 @@ mod tests {
|
||||
#[test]
|
||||
fn load_system_prompt_reads_claude_files_and_config() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(root.join(".claude")).expect("claude dir");
|
||||
fs::create_dir_all(root.join(".claw")).expect("claw dir");
|
||||
fs::write(root.join("CLAUDE.md"), "Project rules").expect("write instructions");
|
||||
fs::write(
|
||||
root.join(".claude").join("settings.json"),
|
||||
root.join(".claw").join("settings.json"),
|
||||
r#"{"permissionMode":"acceptEdits"}"#,
|
||||
)
|
||||
.expect("write settings");
|
||||
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let previous = std::env::current_dir().expect("cwd");
|
||||
let original_home = std::env::var("HOME").ok();
|
||||
let original_claude_home = std::env::var("CLAUDE_CONFIG_HOME").ok();
|
||||
let original_claw_home = std::env::var("CLAW_CONFIG_HOME").ok();
|
||||
std::env::set_var("HOME", &root);
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", root.join("missing-home"));
|
||||
std::env::set_var("CLAW_CONFIG_HOME", root.join("missing-home"));
|
||||
std::env::set_current_dir(&root).expect("change cwd");
|
||||
let prompt = super::load_system_prompt(&root, "2026-03-31", "linux", "6.8")
|
||||
.expect("system prompt should load")
|
||||
@@ -687,10 +817,10 @@ mod tests {
|
||||
} else {
|
||||
std::env::remove_var("HOME");
|
||||
}
|
||||
if let Some(value) = original_claude_home {
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", value);
|
||||
if let Some(value) = original_claw_home {
|
||||
std::env::set_var("CLAW_CONFIG_HOME", value);
|
||||
} else {
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
std::env::remove_var("CLAW_CONFIG_HOME");
|
||||
}
|
||||
|
||||
assert!(prompt.contains("Project rules"));
|
||||
@@ -701,10 +831,10 @@ mod tests {
|
||||
#[test]
|
||||
fn renders_claude_code_style_sections_with_project_context() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(root.join(".claude")).expect("claude dir");
|
||||
fs::create_dir_all(root.join(".claw")).expect("claw dir");
|
||||
fs::write(root.join("CLAUDE.md"), "Project rules").expect("write CLAUDE.md");
|
||||
fs::write(
|
||||
root.join(".claude").join("settings.json"),
|
||||
root.join(".claw").join("settings.json"),
|
||||
r#"{"permissionMode":"acceptEdits"}"#,
|
||||
)
|
||||
.expect("write settings");
|
||||
@@ -743,9 +873,9 @@ mod tests {
|
||||
fn discovers_dot_claude_instructions_markdown() {
|
||||
let root = temp_dir();
|
||||
let nested = root.join("apps").join("api");
|
||||
fs::create_dir_all(nested.join(".claude")).expect("nested claude dir");
|
||||
fs::create_dir_all(nested.join(".claw")).expect("nested claw dir");
|
||||
fs::write(
|
||||
nested.join(".claude").join("instructions.md"),
|
||||
nested.join(".claw").join("instructions.md"),
|
||||
"instruction markdown",
|
||||
)
|
||||
.expect("write instructions.md");
|
||||
@@ -754,7 +884,7 @@ mod tests {
|
||||
assert!(context
|
||||
.instruction_files
|
||||
.iter()
|
||||
.any(|file| file.path.ends_with(".claude/instructions.md")));
|
||||
.any(|file| file.path.ends_with(".claw/instructions.md")));
|
||||
assert!(
|
||||
render_instruction_files(&context.instruction_files).contains("instruction markdown")
|
||||
);
|
||||
|
||||
@@ -0,0 +1,631 @@
|
||||
#![allow(clippy::cast_possible_truncation, clippy::uninlined_format_args)]
|
||||
//! Recovery recipes for common failure scenarios.
|
||||
//!
|
||||
//! Encodes known automatic recoveries for the six failure scenarios
|
||||
//! listed in ROADMAP item 8, and enforces one automatic recovery
|
||||
//! attempt before escalation. Each attempt is emitted as a structured
|
||||
//! recovery event.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::worker_boot::WorkerFailureKind;
|
||||
|
||||
/// The six failure scenarios that have known recovery recipes.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum FailureScenario {
|
||||
TrustPromptUnresolved,
|
||||
PromptMisdelivery,
|
||||
StaleBranch,
|
||||
CompileRedCrossCrate,
|
||||
McpHandshakeFailure,
|
||||
PartialPluginStartup,
|
||||
ProviderFailure,
|
||||
}
|
||||
|
||||
impl FailureScenario {
|
||||
/// Returns all known failure scenarios.
|
||||
#[must_use]
|
||||
pub fn all() -> &'static [FailureScenario] {
|
||||
&[
|
||||
Self::TrustPromptUnresolved,
|
||||
Self::PromptMisdelivery,
|
||||
Self::StaleBranch,
|
||||
Self::CompileRedCrossCrate,
|
||||
Self::McpHandshakeFailure,
|
||||
Self::PartialPluginStartup,
|
||||
Self::ProviderFailure,
|
||||
]
|
||||
}
|
||||
|
||||
/// Map a `WorkerFailureKind` to the corresponding `FailureScenario`.
|
||||
/// This is the bridge that lets recovery policy consume worker boot events.
|
||||
#[must_use]
|
||||
pub fn from_worker_failure_kind(kind: WorkerFailureKind) -> Self {
|
||||
match kind {
|
||||
WorkerFailureKind::TrustGate => Self::TrustPromptUnresolved,
|
||||
WorkerFailureKind::PromptDelivery => Self::PromptMisdelivery,
|
||||
WorkerFailureKind::Protocol => Self::McpHandshakeFailure,
|
||||
WorkerFailureKind::Provider => Self::ProviderFailure,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for FailureScenario {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::TrustPromptUnresolved => write!(f, "trust_prompt_unresolved"),
|
||||
Self::PromptMisdelivery => write!(f, "prompt_misdelivery"),
|
||||
Self::StaleBranch => write!(f, "stale_branch"),
|
||||
Self::CompileRedCrossCrate => write!(f, "compile_red_cross_crate"),
|
||||
Self::McpHandshakeFailure => write!(f, "mcp_handshake_failure"),
|
||||
Self::PartialPluginStartup => write!(f, "partial_plugin_startup"),
|
||||
Self::ProviderFailure => write!(f, "provider_failure"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Individual step that can be executed as part of a recovery recipe.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RecoveryStep {
|
||||
AcceptTrustPrompt,
|
||||
RedirectPromptToAgent,
|
||||
RebaseBranch,
|
||||
CleanBuild,
|
||||
RetryMcpHandshake { timeout: u64 },
|
||||
RestartPlugin { name: String },
|
||||
RestartWorker,
|
||||
EscalateToHuman { reason: String },
|
||||
}
|
||||
|
||||
/// Policy governing what happens when automatic recovery is exhausted.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum EscalationPolicy {
|
||||
AlertHuman,
|
||||
LogAndContinue,
|
||||
Abort,
|
||||
}
|
||||
|
||||
/// A recovery recipe encodes the sequence of steps to attempt for a
|
||||
/// given failure scenario, along with the maximum number of automatic
|
||||
/// attempts and the escalation policy.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RecoveryRecipe {
|
||||
pub scenario: FailureScenario,
|
||||
pub steps: Vec<RecoveryStep>,
|
||||
pub max_attempts: u32,
|
||||
pub escalation_policy: EscalationPolicy,
|
||||
}
|
||||
|
||||
/// Outcome of a recovery attempt.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RecoveryResult {
|
||||
Recovered {
|
||||
steps_taken: u32,
|
||||
},
|
||||
PartialRecovery {
|
||||
recovered: Vec<RecoveryStep>,
|
||||
remaining: Vec<RecoveryStep>,
|
||||
},
|
||||
EscalationRequired {
|
||||
reason: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Structured event emitted during recovery.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RecoveryEvent {
|
||||
RecoveryAttempted {
|
||||
scenario: FailureScenario,
|
||||
recipe: RecoveryRecipe,
|
||||
result: RecoveryResult,
|
||||
},
|
||||
RecoverySucceeded,
|
||||
RecoveryFailed,
|
||||
Escalated,
|
||||
}
|
||||
|
||||
/// Minimal context for tracking recovery state and emitting events.
|
||||
///
|
||||
/// Holds per-scenario attempt counts, a structured event log, and an
|
||||
/// optional simulation knob for controlling step outcomes during tests.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct RecoveryContext {
|
||||
attempts: HashMap<FailureScenario, u32>,
|
||||
events: Vec<RecoveryEvent>,
|
||||
/// Optional step index at which simulated execution fails.
|
||||
/// `None` means all steps succeed.
|
||||
fail_at_step: Option<usize>,
|
||||
}
|
||||
|
||||
impl RecoveryContext {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Configure a step index at which simulated execution will fail.
|
||||
#[must_use]
|
||||
pub fn with_fail_at_step(mut self, index: usize) -> Self {
|
||||
self.fail_at_step = Some(index);
|
||||
self
|
||||
}
|
||||
|
||||
/// Returns the structured event log populated during recovery.
|
||||
#[must_use]
|
||||
pub fn events(&self) -> &[RecoveryEvent] {
|
||||
&self.events
|
||||
}
|
||||
|
||||
/// Returns the number of recovery attempts made for a scenario.
|
||||
#[must_use]
|
||||
pub fn attempt_count(&self, scenario: &FailureScenario) -> u32 {
|
||||
self.attempts.get(scenario).copied().unwrap_or(0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the known recovery recipe for the given failure scenario.
|
||||
#[must_use]
|
||||
pub fn recipe_for(scenario: &FailureScenario) -> RecoveryRecipe {
|
||||
match scenario {
|
||||
FailureScenario::TrustPromptUnresolved => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::AcceptTrustPrompt],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
FailureScenario::PromptMisdelivery => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::RedirectPromptToAgent],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
FailureScenario::StaleBranch => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::RebaseBranch, RecoveryStep::CleanBuild],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
FailureScenario::CompileRedCrossCrate => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::CleanBuild],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
FailureScenario::McpHandshakeFailure => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::RetryMcpHandshake { timeout: 5000 }],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::Abort,
|
||||
},
|
||||
FailureScenario::PartialPluginStartup => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![
|
||||
RecoveryStep::RestartPlugin {
|
||||
name: "stalled".to_string(),
|
||||
},
|
||||
RecoveryStep::RetryMcpHandshake { timeout: 3000 },
|
||||
],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::LogAndContinue,
|
||||
},
|
||||
FailureScenario::ProviderFailure => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::RestartWorker],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Attempts automatic recovery for the given failure scenario.
|
||||
///
|
||||
/// Looks up the recipe, enforces the one-attempt-before-escalation
|
||||
/// policy, simulates step execution (controlled by the context), and
|
||||
/// emits structured [`RecoveryEvent`]s for every attempt.
|
||||
pub fn attempt_recovery(scenario: &FailureScenario, ctx: &mut RecoveryContext) -> RecoveryResult {
|
||||
let recipe = recipe_for(scenario);
|
||||
let attempt_count = ctx.attempts.entry(*scenario).or_insert(0);
|
||||
|
||||
// Enforce one automatic recovery attempt before escalation.
|
||||
if *attempt_count >= recipe.max_attempts {
|
||||
let result = RecoveryResult::EscalationRequired {
|
||||
reason: format!(
|
||||
"max recovery attempts ({}) exceeded for {}",
|
||||
recipe.max_attempts, scenario
|
||||
),
|
||||
};
|
||||
ctx.events.push(RecoveryEvent::RecoveryAttempted {
|
||||
scenario: *scenario,
|
||||
recipe,
|
||||
result: result.clone(),
|
||||
});
|
||||
ctx.events.push(RecoveryEvent::Escalated);
|
||||
return result;
|
||||
}
|
||||
|
||||
*attempt_count += 1;
|
||||
|
||||
// Execute steps, honoring the optional fail_at_step simulation.
|
||||
let fail_index = ctx.fail_at_step;
|
||||
let mut executed = Vec::new();
|
||||
let mut failed = false;
|
||||
|
||||
for (i, step) in recipe.steps.iter().enumerate() {
|
||||
if fail_index == Some(i) {
|
||||
failed = true;
|
||||
break;
|
||||
}
|
||||
executed.push(step.clone());
|
||||
}
|
||||
|
||||
let result = if failed {
|
||||
let remaining: Vec<RecoveryStep> = recipe.steps[executed.len()..].to_vec();
|
||||
if executed.is_empty() {
|
||||
RecoveryResult::EscalationRequired {
|
||||
reason: format!("recovery failed at first step for {}", scenario),
|
||||
}
|
||||
} else {
|
||||
RecoveryResult::PartialRecovery {
|
||||
recovered: executed,
|
||||
remaining,
|
||||
}
|
||||
}
|
||||
} else {
|
||||
RecoveryResult::Recovered {
|
||||
steps_taken: recipe.steps.len() as u32,
|
||||
}
|
||||
};
|
||||
|
||||
// Emit the attempt as structured event data.
|
||||
ctx.events.push(RecoveryEvent::RecoveryAttempted {
|
||||
scenario: *scenario,
|
||||
recipe,
|
||||
result: result.clone(),
|
||||
});
|
||||
|
||||
match &result {
|
||||
RecoveryResult::Recovered { .. } => {
|
||||
ctx.events.push(RecoveryEvent::RecoverySucceeded);
|
||||
}
|
||||
RecoveryResult::PartialRecovery { .. } => {
|
||||
ctx.events.push(RecoveryEvent::RecoveryFailed);
|
||||
}
|
||||
RecoveryResult::EscalationRequired { .. } => {
|
||||
ctx.events.push(RecoveryEvent::Escalated);
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn each_scenario_has_a_matching_recipe() {
|
||||
// given
|
||||
let scenarios = FailureScenario::all();
|
||||
|
||||
// when / then
|
||||
for scenario in scenarios {
|
||||
let recipe = recipe_for(scenario);
|
||||
assert_eq!(
|
||||
recipe.scenario, *scenario,
|
||||
"recipe scenario should match requested scenario"
|
||||
);
|
||||
assert!(
|
||||
!recipe.steps.is_empty(),
|
||||
"recipe for {} should have at least one step",
|
||||
scenario
|
||||
);
|
||||
assert!(
|
||||
recipe.max_attempts >= 1,
|
||||
"recipe for {} should allow at least one attempt",
|
||||
scenario
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn successful_recovery_returns_recovered_and_emits_events() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::TrustPromptUnresolved;
|
||||
|
||||
// when
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
assert_eq!(result, RecoveryResult::Recovered { steps_taken: 1 });
|
||||
assert_eq!(ctx.events().len(), 2);
|
||||
assert!(matches!(
|
||||
&ctx.events()[0],
|
||||
RecoveryEvent::RecoveryAttempted {
|
||||
scenario: s,
|
||||
result: r,
|
||||
..
|
||||
} if *s == FailureScenario::TrustPromptUnresolved
|
||||
&& matches!(r, RecoveryResult::Recovered { steps_taken: 1 })
|
||||
));
|
||||
assert_eq!(ctx.events()[1], RecoveryEvent::RecoverySucceeded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escalation_after_max_attempts_exceeded() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::PromptMisdelivery;
|
||||
|
||||
// when — first attempt succeeds
|
||||
let first = attempt_recovery(&scenario, &mut ctx);
|
||||
assert!(matches!(first, RecoveryResult::Recovered { .. }));
|
||||
|
||||
// when — second attempt should escalate
|
||||
let second = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
assert!(
|
||||
matches!(
|
||||
&second,
|
||||
RecoveryResult::EscalationRequired { reason }
|
||||
if reason.contains("max recovery attempts")
|
||||
),
|
||||
"second attempt should require escalation, got: {second:?}"
|
||||
);
|
||||
assert_eq!(ctx.attempt_count(&scenario), 1);
|
||||
assert!(ctx
|
||||
.events()
|
||||
.iter()
|
||||
.any(|e| matches!(e, RecoveryEvent::Escalated)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_recovery_when_step_fails_midway() {
|
||||
// given — PartialPluginStartup has two steps; fail at step index 1
|
||||
let mut ctx = RecoveryContext::new().with_fail_at_step(1);
|
||||
let scenario = FailureScenario::PartialPluginStartup;
|
||||
|
||||
// when
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
match &result {
|
||||
RecoveryResult::PartialRecovery {
|
||||
recovered,
|
||||
remaining,
|
||||
} => {
|
||||
assert_eq!(recovered.len(), 1, "one step should have succeeded");
|
||||
assert_eq!(remaining.len(), 1, "one step should remain");
|
||||
assert!(matches!(recovered[0], RecoveryStep::RestartPlugin { .. }));
|
||||
assert!(matches!(
|
||||
remaining[0],
|
||||
RecoveryStep::RetryMcpHandshake { .. }
|
||||
));
|
||||
}
|
||||
other => panic!("expected PartialRecovery, got {other:?}"),
|
||||
}
|
||||
assert!(ctx
|
||||
.events()
|
||||
.iter()
|
||||
.any(|e| matches!(e, RecoveryEvent::RecoveryFailed)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_step_failure_escalates_immediately() {
|
||||
// given — fail at step index 0
|
||||
let mut ctx = RecoveryContext::new().with_fail_at_step(0);
|
||||
let scenario = FailureScenario::CompileRedCrossCrate;
|
||||
|
||||
// when
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
assert!(
|
||||
matches!(
|
||||
&result,
|
||||
RecoveryResult::EscalationRequired { reason }
|
||||
if reason.contains("failed at first step")
|
||||
),
|
||||
"zero-step failure should escalate, got: {result:?}"
|
||||
);
|
||||
assert!(ctx
|
||||
.events()
|
||||
.iter()
|
||||
.any(|e| matches!(e, RecoveryEvent::Escalated)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn emitted_events_include_structured_attempt_data() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::McpHandshakeFailure;
|
||||
|
||||
// when
|
||||
let _ = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then — verify the RecoveryAttempted event carries full context
|
||||
let attempted = ctx
|
||||
.events()
|
||||
.iter()
|
||||
.find(|e| matches!(e, RecoveryEvent::RecoveryAttempted { .. }))
|
||||
.expect("should have emitted RecoveryAttempted event");
|
||||
|
||||
match attempted {
|
||||
RecoveryEvent::RecoveryAttempted {
|
||||
scenario: s,
|
||||
recipe,
|
||||
result,
|
||||
} => {
|
||||
assert_eq!(*s, scenario);
|
||||
assert_eq!(recipe.scenario, scenario);
|
||||
assert!(!recipe.steps.is_empty());
|
||||
assert!(matches!(result, RecoveryResult::Recovered { .. }));
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
|
||||
// Verify the event is serializable as structured JSON
|
||||
let json = serde_json::to_string(&ctx.events()[0])
|
||||
.expect("recovery event should be serializable to JSON");
|
||||
assert!(
|
||||
json.contains("mcp_handshake_failure"),
|
||||
"serialized event should contain scenario name"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recovery_context_tracks_attempts_per_scenario() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
|
||||
// when
|
||||
assert_eq!(ctx.attempt_count(&FailureScenario::StaleBranch), 0);
|
||||
attempt_recovery(&FailureScenario::StaleBranch, &mut ctx);
|
||||
|
||||
// then
|
||||
assert_eq!(ctx.attempt_count(&FailureScenario::StaleBranch), 1);
|
||||
assert_eq!(ctx.attempt_count(&FailureScenario::PromptMisdelivery), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_recipe_has_rebase_then_clean_build() {
|
||||
// given
|
||||
let recipe = recipe_for(&FailureScenario::StaleBranch);
|
||||
|
||||
// then
|
||||
assert_eq!(recipe.steps.len(), 2);
|
||||
assert_eq!(recipe.steps[0], RecoveryStep::RebaseBranch);
|
||||
assert_eq!(recipe.steps[1], RecoveryStep::CleanBuild);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_plugin_startup_recipe_has_restart_then_handshake() {
|
||||
// given
|
||||
let recipe = recipe_for(&FailureScenario::PartialPluginStartup);
|
||||
|
||||
// then
|
||||
assert_eq!(recipe.steps.len(), 2);
|
||||
assert!(matches!(
|
||||
recipe.steps[0],
|
||||
RecoveryStep::RestartPlugin { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
recipe.steps[1],
|
||||
RecoveryStep::RetryMcpHandshake { timeout: 3000 }
|
||||
));
|
||||
assert_eq!(recipe.escalation_policy, EscalationPolicy::LogAndContinue);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failure_scenario_display_all_variants() {
|
||||
// given
|
||||
let cases = [
|
||||
(
|
||||
FailureScenario::TrustPromptUnresolved,
|
||||
"trust_prompt_unresolved",
|
||||
),
|
||||
(FailureScenario::PromptMisdelivery, "prompt_misdelivery"),
|
||||
(FailureScenario::StaleBranch, "stale_branch"),
|
||||
(
|
||||
FailureScenario::CompileRedCrossCrate,
|
||||
"compile_red_cross_crate",
|
||||
),
|
||||
(
|
||||
FailureScenario::McpHandshakeFailure,
|
||||
"mcp_handshake_failure",
|
||||
),
|
||||
(
|
||||
FailureScenario::PartialPluginStartup,
|
||||
"partial_plugin_startup",
|
||||
),
|
||||
];
|
||||
|
||||
// when / then
|
||||
for (scenario, expected) in &cases {
|
||||
assert_eq!(scenario.to_string(), *expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multi_step_success_reports_correct_steps_taken() {
|
||||
// given — StaleBranch has 2 steps, no simulated failure
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::StaleBranch;
|
||||
|
||||
// when
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
assert_eq!(result, RecoveryResult::Recovered { steps_taken: 2 });
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_handshake_recipe_uses_abort_escalation_policy() {
|
||||
// given
|
||||
let recipe = recipe_for(&FailureScenario::McpHandshakeFailure);
|
||||
|
||||
// then
|
||||
assert_eq!(recipe.escalation_policy, EscalationPolicy::Abort);
|
||||
assert_eq!(recipe.max_attempts, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_failure_kind_maps_to_failure_scenario() {
|
||||
// given / when / then — verify the bridge is correct
|
||||
assert_eq!(
|
||||
FailureScenario::from_worker_failure_kind(WorkerFailureKind::TrustGate),
|
||||
FailureScenario::TrustPromptUnresolved,
|
||||
);
|
||||
assert_eq!(
|
||||
FailureScenario::from_worker_failure_kind(WorkerFailureKind::PromptDelivery),
|
||||
FailureScenario::PromptMisdelivery,
|
||||
);
|
||||
assert_eq!(
|
||||
FailureScenario::from_worker_failure_kind(WorkerFailureKind::Protocol),
|
||||
FailureScenario::McpHandshakeFailure,
|
||||
);
|
||||
assert_eq!(
|
||||
FailureScenario::from_worker_failure_kind(WorkerFailureKind::Provider),
|
||||
FailureScenario::ProviderFailure,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_failure_recipe_uses_restart_worker_step() {
|
||||
// given
|
||||
let recipe = recipe_for(&FailureScenario::ProviderFailure);
|
||||
|
||||
// then
|
||||
assert_eq!(recipe.scenario, FailureScenario::ProviderFailure);
|
||||
assert!(recipe.steps.contains(&RecoveryStep::RestartWorker));
|
||||
assert_eq!(recipe.escalation_policy, EscalationPolicy::AlertHuman);
|
||||
assert_eq!(recipe.max_attempts, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_failure_recovery_attempt_succeeds_then_escalates() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::ProviderFailure;
|
||||
|
||||
// when — first attempt
|
||||
let first = attempt_recovery(&scenario, &mut ctx);
|
||||
assert!(matches!(first, RecoveryResult::Recovered { .. }));
|
||||
|
||||
// when — second attempt should escalate (max_attempts=1)
|
||||
let second = attempt_recovery(&scenario, &mut ctx);
|
||||
assert!(matches!(second, RecoveryResult::EscalationRequired { .. }));
|
||||
assert!(ctx
|
||||
.events()
|
||||
.iter()
|
||||
.any(|e| matches!(e, RecoveryEvent::Escalated)));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,385 @@
|
||||
use std::env;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum FilesystemIsolationMode {
|
||||
Off,
|
||||
#[default]
|
||||
WorkspaceOnly,
|
||||
AllowList,
|
||||
}
|
||||
|
||||
impl FilesystemIsolationMode {
|
||||
#[must_use]
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Off => "off",
|
||||
Self::WorkspaceOnly => "workspace-only",
|
||||
Self::AllowList => "allow-list",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
|
||||
pub struct SandboxConfig {
|
||||
pub enabled: Option<bool>,
|
||||
pub namespace_restrictions: Option<bool>,
|
||||
pub network_isolation: Option<bool>,
|
||||
pub filesystem_mode: Option<FilesystemIsolationMode>,
|
||||
pub allowed_mounts: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
|
||||
pub struct SandboxRequest {
|
||||
pub enabled: bool,
|
||||
pub namespace_restrictions: bool,
|
||||
pub network_isolation: bool,
|
||||
pub filesystem_mode: FilesystemIsolationMode,
|
||||
pub allowed_mounts: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
|
||||
pub struct ContainerEnvironment {
|
||||
pub in_container: bool,
|
||||
pub markers: Vec<String>,
|
||||
}
|
||||
|
||||
#[allow(clippy::struct_excessive_bools)]
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
|
||||
pub struct SandboxStatus {
|
||||
pub enabled: bool,
|
||||
pub requested: SandboxRequest,
|
||||
pub supported: bool,
|
||||
pub active: bool,
|
||||
pub namespace_supported: bool,
|
||||
pub namespace_active: bool,
|
||||
pub network_supported: bool,
|
||||
pub network_active: bool,
|
||||
pub filesystem_mode: FilesystemIsolationMode,
|
||||
pub filesystem_active: bool,
|
||||
pub allowed_mounts: Vec<String>,
|
||||
pub in_container: bool,
|
||||
pub container_markers: Vec<String>,
|
||||
pub fallback_reason: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SandboxDetectionInputs<'a> {
|
||||
pub env_pairs: Vec<(String, String)>,
|
||||
pub dockerenv_exists: bool,
|
||||
pub containerenv_exists: bool,
|
||||
pub proc_1_cgroup: Option<&'a str>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LinuxSandboxCommand {
|
||||
pub program: String,
|
||||
pub args: Vec<String>,
|
||||
pub env: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
impl SandboxConfig {
|
||||
#[must_use]
|
||||
pub fn resolve_request(
|
||||
&self,
|
||||
enabled_override: Option<bool>,
|
||||
namespace_override: Option<bool>,
|
||||
network_override: Option<bool>,
|
||||
filesystem_mode_override: Option<FilesystemIsolationMode>,
|
||||
allowed_mounts_override: Option<Vec<String>>,
|
||||
) -> SandboxRequest {
|
||||
SandboxRequest {
|
||||
enabled: enabled_override.unwrap_or(self.enabled.unwrap_or(true)),
|
||||
namespace_restrictions: namespace_override
|
||||
.unwrap_or(self.namespace_restrictions.unwrap_or(true)),
|
||||
network_isolation: network_override.unwrap_or(self.network_isolation.unwrap_or(false)),
|
||||
filesystem_mode: filesystem_mode_override
|
||||
.or(self.filesystem_mode)
|
||||
.unwrap_or_default(),
|
||||
allowed_mounts: allowed_mounts_override.unwrap_or_else(|| self.allowed_mounts.clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn detect_container_environment() -> ContainerEnvironment {
|
||||
let proc_1_cgroup = fs::read_to_string("/proc/1/cgroup").ok();
|
||||
detect_container_environment_from(SandboxDetectionInputs {
|
||||
env_pairs: env::vars().collect(),
|
||||
dockerenv_exists: Path::new("/.dockerenv").exists(),
|
||||
containerenv_exists: Path::new("/run/.containerenv").exists(),
|
||||
proc_1_cgroup: proc_1_cgroup.as_deref(),
|
||||
})
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn detect_container_environment_from(
|
||||
inputs: SandboxDetectionInputs<'_>,
|
||||
) -> ContainerEnvironment {
|
||||
let mut markers = Vec::new();
|
||||
if inputs.dockerenv_exists {
|
||||
markers.push("/.dockerenv".to_string());
|
||||
}
|
||||
if inputs.containerenv_exists {
|
||||
markers.push("/run/.containerenv".to_string());
|
||||
}
|
||||
for (key, value) in inputs.env_pairs {
|
||||
let normalized = key.to_ascii_lowercase();
|
||||
if matches!(
|
||||
normalized.as_str(),
|
||||
"container" | "docker" | "podman" | "kubernetes_service_host"
|
||||
) && !value.is_empty()
|
||||
{
|
||||
markers.push(format!("env:{key}={value}"));
|
||||
}
|
||||
}
|
||||
if let Some(cgroup) = inputs.proc_1_cgroup {
|
||||
for needle in ["docker", "containerd", "kubepods", "podman", "libpod"] {
|
||||
if cgroup.contains(needle) {
|
||||
markers.push(format!("/proc/1/cgroup:{needle}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
markers.sort();
|
||||
markers.dedup();
|
||||
ContainerEnvironment {
|
||||
in_container: !markers.is_empty(),
|
||||
markers,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn resolve_sandbox_status(config: &SandboxConfig, cwd: &Path) -> SandboxStatus {
|
||||
let request = config.resolve_request(None, None, None, None, None);
|
||||
resolve_sandbox_status_for_request(&request, cwd)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn resolve_sandbox_status_for_request(request: &SandboxRequest, cwd: &Path) -> SandboxStatus {
|
||||
let container = detect_container_environment();
|
||||
let namespace_supported = cfg!(target_os = "linux") && unshare_user_namespace_works();
|
||||
let network_supported = namespace_supported;
|
||||
let filesystem_active =
|
||||
request.enabled && request.filesystem_mode != FilesystemIsolationMode::Off;
|
||||
let mut fallback_reasons = Vec::new();
|
||||
|
||||
if request.enabled && request.namespace_restrictions && !namespace_supported {
|
||||
fallback_reasons
|
||||
.push("namespace isolation unavailable (requires Linux with `unshare`)".to_string());
|
||||
}
|
||||
if request.enabled && request.network_isolation && !network_supported {
|
||||
fallback_reasons
|
||||
.push("network isolation unavailable (requires Linux with `unshare`)".to_string());
|
||||
}
|
||||
if request.enabled
|
||||
&& request.filesystem_mode == FilesystemIsolationMode::AllowList
|
||||
&& request.allowed_mounts.is_empty()
|
||||
{
|
||||
fallback_reasons
|
||||
.push("filesystem allow-list requested without configured mounts".to_string());
|
||||
}
|
||||
|
||||
let active = request.enabled
|
||||
&& (!request.namespace_restrictions || namespace_supported)
|
||||
&& (!request.network_isolation || network_supported);
|
||||
|
||||
let allowed_mounts = normalize_mounts(&request.allowed_mounts, cwd);
|
||||
|
||||
SandboxStatus {
|
||||
enabled: request.enabled,
|
||||
requested: request.clone(),
|
||||
supported: namespace_supported,
|
||||
active,
|
||||
namespace_supported,
|
||||
namespace_active: request.enabled && request.namespace_restrictions && namespace_supported,
|
||||
network_supported,
|
||||
network_active: request.enabled && request.network_isolation && network_supported,
|
||||
filesystem_mode: request.filesystem_mode,
|
||||
filesystem_active,
|
||||
allowed_mounts,
|
||||
in_container: container.in_container,
|
||||
container_markers: container.markers,
|
||||
fallback_reason: (!fallback_reasons.is_empty()).then(|| fallback_reasons.join("; ")),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn build_linux_sandbox_command(
|
||||
command: &str,
|
||||
cwd: &Path,
|
||||
status: &SandboxStatus,
|
||||
) -> Option<LinuxSandboxCommand> {
|
||||
if !cfg!(target_os = "linux")
|
||||
|| !status.enabled
|
||||
|| (!status.namespace_active && !status.network_active)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut args = vec![
|
||||
"--user".to_string(),
|
||||
"--map-root-user".to_string(),
|
||||
"--mount".to_string(),
|
||||
"--ipc".to_string(),
|
||||
"--pid".to_string(),
|
||||
"--uts".to_string(),
|
||||
"--fork".to_string(),
|
||||
];
|
||||
if status.network_active {
|
||||
args.push("--net".to_string());
|
||||
}
|
||||
args.push("sh".to_string());
|
||||
args.push("-lc".to_string());
|
||||
args.push(command.to_string());
|
||||
|
||||
let sandbox_home = cwd.join(".sandbox-home");
|
||||
let sandbox_tmp = cwd.join(".sandbox-tmp");
|
||||
let mut env = vec![
|
||||
("HOME".to_string(), sandbox_home.display().to_string()),
|
||||
("TMPDIR".to_string(), sandbox_tmp.display().to_string()),
|
||||
(
|
||||
"CLAWD_SANDBOX_FILESYSTEM_MODE".to_string(),
|
||||
status.filesystem_mode.as_str().to_string(),
|
||||
),
|
||||
(
|
||||
"CLAWD_SANDBOX_ALLOWED_MOUNTS".to_string(),
|
||||
status.allowed_mounts.join(":"),
|
||||
),
|
||||
];
|
||||
if let Ok(path) = env::var("PATH") {
|
||||
env.push(("PATH".to_string(), path));
|
||||
}
|
||||
|
||||
Some(LinuxSandboxCommand {
|
||||
program: "unshare".to_string(),
|
||||
args,
|
||||
env,
|
||||
})
|
||||
}
|
||||
|
||||
fn normalize_mounts(mounts: &[String], cwd: &Path) -> Vec<String> {
|
||||
let cwd = cwd.to_path_buf();
|
||||
mounts
|
||||
.iter()
|
||||
.map(|mount| {
|
||||
let path = PathBuf::from(mount);
|
||||
if path.is_absolute() {
|
||||
path
|
||||
} else {
|
||||
cwd.join(path)
|
||||
}
|
||||
})
|
||||
.map(|path| path.display().to_string())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn command_exists(command: &str) -> bool {
|
||||
env::var_os("PATH")
|
||||
.is_some_and(|paths| env::split_paths(&paths).any(|path| path.join(command).exists()))
|
||||
}
|
||||
|
||||
/// Check whether `unshare --user` actually works on this system.
|
||||
/// On some CI environments (e.g. GitHub Actions), the binary exists but
|
||||
/// user namespaces are restricted, causing silent failures.
|
||||
fn unshare_user_namespace_works() -> bool {
|
||||
use std::sync::OnceLock;
|
||||
static RESULT: OnceLock<bool> = OnceLock::new();
|
||||
*RESULT.get_or_init(|| {
|
||||
if !command_exists("unshare") {
|
||||
return false;
|
||||
}
|
||||
std::process::Command::new("unshare")
|
||||
.args(["--user", "--map-root-user", "true"])
|
||||
.stdin(std::process::Stdio::null())
|
||||
.stdout(std::process::Stdio::null())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.status()
|
||||
.map(|s| s.success())
|
||||
.unwrap_or(false)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_linux_sandbox_command, detect_container_environment_from, FilesystemIsolationMode,
|
||||
SandboxConfig, SandboxDetectionInputs,
|
||||
};
|
||||
use std::path::Path;
|
||||
|
||||
#[test]
|
||||
fn detects_container_markers_from_multiple_sources() {
|
||||
let detected = detect_container_environment_from(SandboxDetectionInputs {
|
||||
env_pairs: vec![("container".to_string(), "docker".to_string())],
|
||||
dockerenv_exists: true,
|
||||
containerenv_exists: false,
|
||||
proc_1_cgroup: Some("12:memory:/docker/abc"),
|
||||
});
|
||||
|
||||
assert!(detected.in_container);
|
||||
assert!(detected
|
||||
.markers
|
||||
.iter()
|
||||
.any(|marker| marker == "/.dockerenv"));
|
||||
assert!(detected
|
||||
.markers
|
||||
.iter()
|
||||
.any(|marker| marker == "env:container=docker"));
|
||||
assert!(detected
|
||||
.markers
|
||||
.iter()
|
||||
.any(|marker| marker == "/proc/1/cgroup:docker"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_request_with_overrides() {
|
||||
let config = SandboxConfig {
|
||||
enabled: Some(true),
|
||||
namespace_restrictions: Some(true),
|
||||
network_isolation: Some(false),
|
||||
filesystem_mode: Some(FilesystemIsolationMode::WorkspaceOnly),
|
||||
allowed_mounts: vec!["logs".to_string()],
|
||||
};
|
||||
|
||||
let request = config.resolve_request(
|
||||
Some(true),
|
||||
Some(false),
|
||||
Some(true),
|
||||
Some(FilesystemIsolationMode::AllowList),
|
||||
Some(vec!["tmp".to_string()]),
|
||||
);
|
||||
|
||||
assert!(request.enabled);
|
||||
assert!(!request.namespace_restrictions);
|
||||
assert!(request.network_isolation);
|
||||
assert_eq!(request.filesystem_mode, FilesystemIsolationMode::AllowList);
|
||||
assert_eq!(request.allowed_mounts, vec!["tmp"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_linux_launcher_with_network_flag_when_requested() {
|
||||
let config = SandboxConfig::default();
|
||||
let status = super::resolve_sandbox_status_for_request(
|
||||
&config.resolve_request(
|
||||
Some(true),
|
||||
Some(true),
|
||||
Some(true),
|
||||
Some(FilesystemIsolationMode::WorkspaceOnly),
|
||||
None,
|
||||
),
|
||||
Path::new("/workspace"),
|
||||
);
|
||||
|
||||
if let Some(launcher) =
|
||||
build_linux_sandbox_command("printf hi", Path::new("/workspace"), &status)
|
||||
{
|
||||
assert_eq!(launcher.program, "unshare");
|
||||
assert!(launcher.args.iter().any(|arg| arg == "--mount"));
|
||||
assert!(launcher.args.iter().any(|arg| arg == "--net") == status.network_active);
|
||||
}
|
||||
}
|
||||
}
|
||||
+1095
-124
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,873 @@
|
||||
#![allow(dead_code)]
|
||||
use std::env;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use crate::session::{Session, SessionError};
|
||||
|
||||
/// Per-worktree session store that namespaces on-disk session files by
|
||||
/// workspace fingerprint so that parallel `opencode serve` instances never
|
||||
/// collide.
|
||||
///
|
||||
/// Create via [`SessionStore::from_cwd`] (derives the store path from the
|
||||
/// server's working directory) or [`SessionStore::from_data_dir`] (honours an
|
||||
/// explicit `--data-dir` flag). Both constructors produce a directory layout
|
||||
/// of `<data_dir>/sessions/<workspace_hash>/` where `<workspace_hash>` is a
|
||||
/// stable hex digest of the canonical workspace root.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SessionStore {
|
||||
/// Resolved root of the session namespace, e.g.
|
||||
/// `/home/user/project/.claw/sessions/a1b2c3d4e5f60718/`.
|
||||
sessions_root: PathBuf,
|
||||
/// The canonical workspace path that was fingerprinted.
|
||||
workspace_root: PathBuf,
|
||||
}
|
||||
|
||||
impl SessionStore {
|
||||
/// Build a store from the server's current working directory.
|
||||
///
|
||||
/// The on-disk layout becomes `<cwd>/.claw/sessions/<workspace_hash>/`.
|
||||
pub fn from_cwd(cwd: impl AsRef<Path>) -> Result<Self, SessionControlError> {
|
||||
let cwd = cwd.as_ref();
|
||||
let sessions_root = cwd
|
||||
.join(".claw")
|
||||
.join("sessions")
|
||||
.join(workspace_fingerprint(cwd));
|
||||
fs::create_dir_all(&sessions_root)?;
|
||||
Ok(Self {
|
||||
sessions_root,
|
||||
workspace_root: cwd.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a store from an explicit `--data-dir` flag.
|
||||
///
|
||||
/// The on-disk layout becomes `<data_dir>/sessions/<workspace_hash>/`
|
||||
/// where `<workspace_hash>` is derived from `workspace_root`.
|
||||
pub fn from_data_dir(
|
||||
data_dir: impl AsRef<Path>,
|
||||
workspace_root: impl AsRef<Path>,
|
||||
) -> Result<Self, SessionControlError> {
|
||||
let workspace_root = workspace_root.as_ref();
|
||||
let sessions_root = data_dir
|
||||
.as_ref()
|
||||
.join("sessions")
|
||||
.join(workspace_fingerprint(workspace_root));
|
||||
fs::create_dir_all(&sessions_root)?;
|
||||
Ok(Self {
|
||||
sessions_root,
|
||||
workspace_root: workspace_root.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
/// The fully resolved sessions directory for this namespace.
|
||||
#[must_use]
|
||||
pub fn sessions_dir(&self) -> &Path {
|
||||
&self.sessions_root
|
||||
}
|
||||
|
||||
/// The workspace root this store is bound to.
|
||||
#[must_use]
|
||||
pub fn workspace_root(&self) -> &Path {
|
||||
&self.workspace_root
|
||||
}
|
||||
|
||||
pub fn create_handle(&self, session_id: &str) -> SessionHandle {
|
||||
let id = session_id.to_string();
|
||||
let path = self
|
||||
.sessions_root
|
||||
.join(format!("{id}.{PRIMARY_SESSION_EXTENSION}"));
|
||||
SessionHandle { id, path }
|
||||
}
|
||||
|
||||
pub fn resolve_reference(&self, reference: &str) -> Result<SessionHandle, SessionControlError> {
|
||||
if is_session_reference_alias(reference) {
|
||||
let latest = self.latest_session()?;
|
||||
return Ok(SessionHandle {
|
||||
id: latest.id,
|
||||
path: latest.path,
|
||||
});
|
||||
}
|
||||
|
||||
let direct = PathBuf::from(reference);
|
||||
let candidate = if direct.is_absolute() {
|
||||
direct.clone()
|
||||
} else {
|
||||
self.workspace_root.join(&direct)
|
||||
};
|
||||
let looks_like_path = direct.extension().is_some() || direct.components().count() > 1;
|
||||
let path = if candidate.exists() {
|
||||
candidate
|
||||
} else if looks_like_path {
|
||||
return Err(SessionControlError::Format(
|
||||
format_missing_session_reference(reference),
|
||||
));
|
||||
} else {
|
||||
self.resolve_managed_path(reference)?
|
||||
};
|
||||
|
||||
Ok(SessionHandle {
|
||||
id: session_id_from_path(&path).unwrap_or_else(|| reference.to_string()),
|
||||
path,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve_managed_path(&self, session_id: &str) -> Result<PathBuf, SessionControlError> {
|
||||
for extension in [PRIMARY_SESSION_EXTENSION, LEGACY_SESSION_EXTENSION] {
|
||||
let path = self.sessions_root.join(format!("{session_id}.{extension}"));
|
||||
if path.exists() {
|
||||
return Ok(path);
|
||||
}
|
||||
}
|
||||
Err(SessionControlError::Format(
|
||||
format_missing_session_reference(session_id),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn list_sessions(&self) -> Result<Vec<ManagedSessionSummary>, SessionControlError> {
|
||||
let mut sessions = Vec::new();
|
||||
let read_result = fs::read_dir(&self.sessions_root);
|
||||
let entries = match read_result {
|
||||
Ok(entries) => entries,
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(sessions),
|
||||
Err(err) => return Err(err.into()),
|
||||
};
|
||||
for entry in entries {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if !is_managed_session_file(&path) {
|
||||
continue;
|
||||
}
|
||||
let metadata = entry.metadata()?;
|
||||
let modified_epoch_millis = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|time| time.duration_since(UNIX_EPOCH).ok())
|
||||
.map(|duration| duration.as_millis())
|
||||
.unwrap_or_default();
|
||||
let (id, message_count, parent_session_id, branch_name) =
|
||||
match Session::load_from_path(&path) {
|
||||
Ok(session) => {
|
||||
let parent_session_id = session
|
||||
.fork
|
||||
.as_ref()
|
||||
.map(|fork| fork.parent_session_id.clone());
|
||||
let branch_name = session
|
||||
.fork
|
||||
.as_ref()
|
||||
.and_then(|fork| fork.branch_name.clone());
|
||||
(
|
||||
session.session_id,
|
||||
session.messages.len(),
|
||||
parent_session_id,
|
||||
branch_name,
|
||||
)
|
||||
}
|
||||
Err(_) => (
|
||||
path.file_stem()
|
||||
.and_then(|value| value.to_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
};
|
||||
sessions.push(ManagedSessionSummary {
|
||||
id,
|
||||
path,
|
||||
modified_epoch_millis,
|
||||
message_count,
|
||||
parent_session_id,
|
||||
branch_name,
|
||||
});
|
||||
}
|
||||
sessions.sort_by(|left, right| {
|
||||
right
|
||||
.modified_epoch_millis
|
||||
.cmp(&left.modified_epoch_millis)
|
||||
.then_with(|| right.id.cmp(&left.id))
|
||||
});
|
||||
Ok(sessions)
|
||||
}
|
||||
|
||||
pub fn latest_session(&self) -> Result<ManagedSessionSummary, SessionControlError> {
|
||||
self.list_sessions()?
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| SessionControlError::Format(format_no_managed_sessions()))
|
||||
}
|
||||
|
||||
pub fn load_session(
|
||||
&self,
|
||||
reference: &str,
|
||||
) -> Result<LoadedManagedSession, SessionControlError> {
|
||||
let handle = self.resolve_reference(reference)?;
|
||||
let session = Session::load_from_path(&handle.path)?;
|
||||
Ok(LoadedManagedSession {
|
||||
handle: SessionHandle {
|
||||
id: session.session_id.clone(),
|
||||
path: handle.path,
|
||||
},
|
||||
session,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn fork_session(
|
||||
&self,
|
||||
session: &Session,
|
||||
branch_name: Option<String>,
|
||||
) -> Result<ForkedManagedSession, SessionControlError> {
|
||||
let parent_session_id = session.session_id.clone();
|
||||
let forked = session.fork(branch_name);
|
||||
let handle = self.create_handle(&forked.session_id);
|
||||
let branch_name = forked
|
||||
.fork
|
||||
.as_ref()
|
||||
.and_then(|fork| fork.branch_name.clone());
|
||||
let forked = forked.with_persistence_path(handle.path.clone());
|
||||
forked.save_to_path(&handle.path)?;
|
||||
Ok(ForkedManagedSession {
|
||||
parent_session_id,
|
||||
handle,
|
||||
session: forked,
|
||||
branch_name,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Stable hex fingerprint of a workspace path.
|
||||
///
|
||||
/// Uses FNV-1a (64-bit) to produce a 16-char hex string that partitions the
|
||||
/// on-disk session directory per workspace root.
|
||||
#[must_use]
|
||||
pub fn workspace_fingerprint(workspace_root: &Path) -> String {
|
||||
let input = workspace_root.to_string_lossy();
|
||||
let mut hash = 0xcbf2_9ce4_8422_2325_u64;
|
||||
for byte in input.as_bytes() {
|
||||
hash ^= u64::from(*byte);
|
||||
hash = hash.wrapping_mul(0x0100_0000_01b3);
|
||||
}
|
||||
format!("{hash:016x}")
|
||||
}
|
||||
|
||||
pub const PRIMARY_SESSION_EXTENSION: &str = "jsonl";
|
||||
pub const LEGACY_SESSION_EXTENSION: &str = "json";
|
||||
pub const LATEST_SESSION_REFERENCE: &str = "latest";
|
||||
|
||||
const SESSION_REFERENCE_ALIASES: &[&str] = &[LATEST_SESSION_REFERENCE, "last", "recent"];
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SessionHandle {
|
||||
pub id: String,
|
||||
pub path: PathBuf,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ManagedSessionSummary {
|
||||
pub id: String,
|
||||
pub path: PathBuf,
|
||||
pub modified_epoch_millis: u128,
|
||||
pub message_count: usize,
|
||||
pub parent_session_id: Option<String>,
|
||||
pub branch_name: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LoadedManagedSession {
|
||||
pub handle: SessionHandle,
|
||||
pub session: Session,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ForkedManagedSession {
|
||||
pub parent_session_id: String,
|
||||
pub handle: SessionHandle,
|
||||
pub session: Session,
|
||||
pub branch_name: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SessionControlError {
|
||||
Io(std::io::Error),
|
||||
Session(SessionError),
|
||||
Format(String),
|
||||
}
|
||||
|
||||
impl Display for SessionControlError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Io(error) => write!(f, "{error}"),
|
||||
Self::Session(error) => write!(f, "{error}"),
|
||||
Self::Format(error) => write!(f, "{error}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for SessionControlError {}
|
||||
|
||||
impl From<std::io::Error> for SessionControlError {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::Io(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SessionError> for SessionControlError {
|
||||
fn from(value: SessionError) -> Self {
|
||||
Self::Session(value)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sessions_dir() -> Result<PathBuf, SessionControlError> {
|
||||
managed_sessions_dir_for(env::current_dir()?)
|
||||
}
|
||||
|
||||
pub fn managed_sessions_dir_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
) -> Result<PathBuf, SessionControlError> {
|
||||
let path = base_dir.as_ref().join(".claw").join("sessions");
|
||||
fs::create_dir_all(&path)?;
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
pub fn create_managed_session_handle(
|
||||
session_id: &str,
|
||||
) -> Result<SessionHandle, SessionControlError> {
|
||||
create_managed_session_handle_for(env::current_dir()?, session_id)
|
||||
}
|
||||
|
||||
pub fn create_managed_session_handle_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
session_id: &str,
|
||||
) -> Result<SessionHandle, SessionControlError> {
|
||||
let id = session_id.to_string();
|
||||
let path =
|
||||
managed_sessions_dir_for(base_dir)?.join(format!("{id}.{PRIMARY_SESSION_EXTENSION}"));
|
||||
Ok(SessionHandle { id, path })
|
||||
}
|
||||
|
||||
pub fn resolve_session_reference(reference: &str) -> Result<SessionHandle, SessionControlError> {
|
||||
resolve_session_reference_for(env::current_dir()?, reference)
|
||||
}
|
||||
|
||||
pub fn resolve_session_reference_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
reference: &str,
|
||||
) -> Result<SessionHandle, SessionControlError> {
|
||||
let base_dir = base_dir.as_ref();
|
||||
if is_session_reference_alias(reference) {
|
||||
let latest = latest_managed_session_for(base_dir)?;
|
||||
return Ok(SessionHandle {
|
||||
id: latest.id,
|
||||
path: latest.path,
|
||||
});
|
||||
}
|
||||
|
||||
let direct = PathBuf::from(reference);
|
||||
let candidate = if direct.is_absolute() {
|
||||
direct.clone()
|
||||
} else {
|
||||
base_dir.join(&direct)
|
||||
};
|
||||
let looks_like_path = direct.extension().is_some() || direct.components().count() > 1;
|
||||
let path = if candidate.exists() {
|
||||
candidate
|
||||
} else if looks_like_path {
|
||||
return Err(SessionControlError::Format(
|
||||
format_missing_session_reference(reference),
|
||||
));
|
||||
} else {
|
||||
resolve_managed_session_path_for(base_dir, reference)?
|
||||
};
|
||||
|
||||
Ok(SessionHandle {
|
||||
id: session_id_from_path(&path).unwrap_or_else(|| reference.to_string()),
|
||||
path,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve_managed_session_path(session_id: &str) -> Result<PathBuf, SessionControlError> {
|
||||
resolve_managed_session_path_for(env::current_dir()?, session_id)
|
||||
}
|
||||
|
||||
pub fn resolve_managed_session_path_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
session_id: &str,
|
||||
) -> Result<PathBuf, SessionControlError> {
|
||||
let directory = managed_sessions_dir_for(base_dir)?;
|
||||
for extension in [PRIMARY_SESSION_EXTENSION, LEGACY_SESSION_EXTENSION] {
|
||||
let path = directory.join(format!("{session_id}.{extension}"));
|
||||
if path.exists() {
|
||||
return Ok(path);
|
||||
}
|
||||
}
|
||||
Err(SessionControlError::Format(
|
||||
format_missing_session_reference(session_id),
|
||||
))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_managed_session_file(path: &Path) -> bool {
|
||||
path.extension()
|
||||
.and_then(|ext| ext.to_str())
|
||||
.is_some_and(|extension| {
|
||||
extension == PRIMARY_SESSION_EXTENSION || extension == LEGACY_SESSION_EXTENSION
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_managed_sessions() -> Result<Vec<ManagedSessionSummary>, SessionControlError> {
|
||||
list_managed_sessions_for(env::current_dir()?)
|
||||
}
|
||||
|
||||
pub fn list_managed_sessions_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
) -> Result<Vec<ManagedSessionSummary>, SessionControlError> {
|
||||
let mut sessions = Vec::new();
|
||||
for entry in fs::read_dir(managed_sessions_dir_for(base_dir)?)? {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if !is_managed_session_file(&path) {
|
||||
continue;
|
||||
}
|
||||
let metadata = entry.metadata()?;
|
||||
let modified_epoch_millis = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|time| time.duration_since(UNIX_EPOCH).ok())
|
||||
.map(|duration| duration.as_millis())
|
||||
.unwrap_or_default();
|
||||
let (id, message_count, parent_session_id, branch_name) =
|
||||
match Session::load_from_path(&path) {
|
||||
Ok(session) => {
|
||||
let parent_session_id = session
|
||||
.fork
|
||||
.as_ref()
|
||||
.map(|fork| fork.parent_session_id.clone());
|
||||
let branch_name = session
|
||||
.fork
|
||||
.as_ref()
|
||||
.and_then(|fork| fork.branch_name.clone());
|
||||
(
|
||||
session.session_id,
|
||||
session.messages.len(),
|
||||
parent_session_id,
|
||||
branch_name,
|
||||
)
|
||||
}
|
||||
Err(_) => (
|
||||
path.file_stem()
|
||||
.and_then(|value| value.to_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
};
|
||||
sessions.push(ManagedSessionSummary {
|
||||
id,
|
||||
path,
|
||||
modified_epoch_millis,
|
||||
message_count,
|
||||
parent_session_id,
|
||||
branch_name,
|
||||
});
|
||||
}
|
||||
sessions.sort_by(|left, right| {
|
||||
right
|
||||
.modified_epoch_millis
|
||||
.cmp(&left.modified_epoch_millis)
|
||||
.then_with(|| right.id.cmp(&left.id))
|
||||
});
|
||||
Ok(sessions)
|
||||
}
|
||||
|
||||
pub fn latest_managed_session() -> Result<ManagedSessionSummary, SessionControlError> {
|
||||
latest_managed_session_for(env::current_dir()?)
|
||||
}
|
||||
|
||||
pub fn latest_managed_session_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
) -> Result<ManagedSessionSummary, SessionControlError> {
|
||||
list_managed_sessions_for(base_dir)?
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| SessionControlError::Format(format_no_managed_sessions()))
|
||||
}
|
||||
|
||||
pub fn load_managed_session(reference: &str) -> Result<LoadedManagedSession, SessionControlError> {
|
||||
load_managed_session_for(env::current_dir()?, reference)
|
||||
}
|
||||
|
||||
pub fn load_managed_session_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
reference: &str,
|
||||
) -> Result<LoadedManagedSession, SessionControlError> {
|
||||
let handle = resolve_session_reference_for(base_dir, reference)?;
|
||||
let session = Session::load_from_path(&handle.path)?;
|
||||
Ok(LoadedManagedSession {
|
||||
handle: SessionHandle {
|
||||
id: session.session_id.clone(),
|
||||
path: handle.path,
|
||||
},
|
||||
session,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn fork_managed_session(
|
||||
session: &Session,
|
||||
branch_name: Option<String>,
|
||||
) -> Result<ForkedManagedSession, SessionControlError> {
|
||||
fork_managed_session_for(env::current_dir()?, session, branch_name)
|
||||
}
|
||||
|
||||
pub fn fork_managed_session_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
session: &Session,
|
||||
branch_name: Option<String>,
|
||||
) -> Result<ForkedManagedSession, SessionControlError> {
|
||||
let parent_session_id = session.session_id.clone();
|
||||
let forked = session.fork(branch_name);
|
||||
let handle = create_managed_session_handle_for(base_dir, &forked.session_id)?;
|
||||
let branch_name = forked
|
||||
.fork
|
||||
.as_ref()
|
||||
.and_then(|fork| fork.branch_name.clone());
|
||||
let forked = forked.with_persistence_path(handle.path.clone());
|
||||
forked.save_to_path(&handle.path)?;
|
||||
Ok(ForkedManagedSession {
|
||||
parent_session_id,
|
||||
handle,
|
||||
session: forked,
|
||||
branch_name,
|
||||
})
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_session_reference_alias(reference: &str) -> bool {
|
||||
SESSION_REFERENCE_ALIASES
|
||||
.iter()
|
||||
.any(|alias| reference.eq_ignore_ascii_case(alias))
|
||||
}
|
||||
|
||||
fn session_id_from_path(path: &Path) -> Option<String> {
|
||||
path.file_name()
|
||||
.and_then(|value| value.to_str())
|
||||
.and_then(|name| {
|
||||
name.strip_suffix(&format!(".{PRIMARY_SESSION_EXTENSION}"))
|
||||
.or_else(|| name.strip_suffix(&format!(".{LEGACY_SESSION_EXTENSION}")))
|
||||
})
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn format_missing_session_reference(reference: &str) -> String {
|
||||
format!(
|
||||
"session not found: {reference}\nHint: managed sessions live in .claw/sessions/. Try `{LATEST_SESSION_REFERENCE}` for the most recent session or `/session list` in the REPL."
|
||||
)
|
||||
}
|
||||
|
||||
fn format_no_managed_sessions() -> String {
|
||||
format!(
|
||||
"no managed sessions found in .claw/sessions/\nStart `claw` to create a session, then rerun with `--resume {LATEST_SESSION_REFERENCE}`."
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
create_managed_session_handle_for, fork_managed_session_for, is_session_reference_alias,
|
||||
list_managed_sessions_for, load_managed_session_for, resolve_session_reference_for,
|
||||
workspace_fingerprint, ManagedSessionSummary, SessionStore, LATEST_SESSION_REFERENCE,
|
||||
};
|
||||
use crate::session::Session;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_dir() -> PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!("runtime-session-control-{nanos}"))
|
||||
}
|
||||
|
||||
fn persist_session(root: &Path, text: &str) -> Session {
|
||||
let mut session = Session::new();
|
||||
session
|
||||
.push_user_text(text)
|
||||
.expect("session message should save");
|
||||
let handle = create_managed_session_handle_for(root, &session.session_id)
|
||||
.expect("managed session handle should build");
|
||||
let session = session.with_persistence_path(handle.path.clone());
|
||||
session
|
||||
.save_to_path(&handle.path)
|
||||
.expect("session should persist");
|
||||
session
|
||||
}
|
||||
|
||||
fn wait_for_next_millisecond() {
|
||||
let start = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_millis();
|
||||
while SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_millis()
|
||||
<= start
|
||||
{}
|
||||
}
|
||||
|
||||
fn summary_by_id<'a>(
|
||||
summaries: &'a [ManagedSessionSummary],
|
||||
id: &str,
|
||||
) -> &'a ManagedSessionSummary {
|
||||
summaries
|
||||
.iter()
|
||||
.find(|summary| summary.id == id)
|
||||
.expect("session summary should exist")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn creates_and_lists_managed_sessions() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir should exist");
|
||||
let older = persist_session(&root, "older session");
|
||||
wait_for_next_millisecond();
|
||||
let newer = persist_session(&root, "newer session");
|
||||
|
||||
// when
|
||||
let sessions = list_managed_sessions_for(&root).expect("managed sessions should list");
|
||||
|
||||
// then
|
||||
assert_eq!(sessions.len(), 2);
|
||||
assert_eq!(sessions[0].id, newer.session_id);
|
||||
assert_eq!(summary_by_id(&sessions, &older.session_id).message_count, 1);
|
||||
assert_eq!(summary_by_id(&sessions, &newer.session_id).message_count, 1);
|
||||
fs::remove_dir_all(root).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_latest_alias_and_loads_session_from_workspace_root() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir should exist");
|
||||
let older = persist_session(&root, "older session");
|
||||
wait_for_next_millisecond();
|
||||
let newer = persist_session(&root, "newer session");
|
||||
|
||||
// when
|
||||
let handle = resolve_session_reference_for(&root, LATEST_SESSION_REFERENCE)
|
||||
.expect("latest alias should resolve");
|
||||
let loaded = load_managed_session_for(&root, "recent")
|
||||
.expect("recent alias should load the latest session");
|
||||
|
||||
// then
|
||||
assert_eq!(handle.id, newer.session_id);
|
||||
assert_eq!(loaded.handle.id, newer.session_id);
|
||||
assert_eq!(loaded.session.messages.len(), 1);
|
||||
assert_ne!(loaded.handle.id, older.session_id);
|
||||
assert!(is_session_reference_alias("last"));
|
||||
fs::remove_dir_all(root).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forks_session_into_managed_storage_with_lineage() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir should exist");
|
||||
let source = persist_session(&root, "parent session");
|
||||
|
||||
// when
|
||||
let forked = fork_managed_session_for(&root, &source, Some("incident-review".to_string()))
|
||||
.expect("session should fork");
|
||||
let sessions = list_managed_sessions_for(&root).expect("managed sessions should list");
|
||||
let summary = summary_by_id(&sessions, &forked.handle.id);
|
||||
|
||||
// then
|
||||
assert_eq!(forked.parent_session_id, source.session_id);
|
||||
assert_eq!(forked.branch_name.as_deref(), Some("incident-review"));
|
||||
assert_eq!(
|
||||
summary.parent_session_id.as_deref(),
|
||||
Some(source.session_id.as_str())
|
||||
);
|
||||
assert_eq!(summary.branch_name.as_deref(), Some("incident-review"));
|
||||
assert_eq!(
|
||||
forked.session.persistence_path(),
|
||||
Some(forked.handle.path.as_path())
|
||||
);
|
||||
fs::remove_dir_all(root).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// Per-worktree session isolation (SessionStore) tests
|
||||
// ------------------------------------------------------------------
|
||||
|
||||
fn persist_session_via_store(store: &SessionStore, text: &str) -> Session {
|
||||
let mut session = Session::new();
|
||||
session
|
||||
.push_user_text(text)
|
||||
.expect("session message should save");
|
||||
let handle = store.create_handle(&session.session_id);
|
||||
let session = session.with_persistence_path(handle.path.clone());
|
||||
session
|
||||
.save_to_path(&handle.path)
|
||||
.expect("session should persist");
|
||||
session
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_fingerprint_is_deterministic_and_differs_per_path() {
|
||||
// given
|
||||
let path_a = Path::new("/tmp/worktree-alpha");
|
||||
let path_b = Path::new("/tmp/worktree-beta");
|
||||
|
||||
// when
|
||||
let fp_a1 = workspace_fingerprint(path_a);
|
||||
let fp_a2 = workspace_fingerprint(path_a);
|
||||
let fp_b = workspace_fingerprint(path_b);
|
||||
|
||||
// then
|
||||
assert_eq!(fp_a1, fp_a2, "same path must produce the same fingerprint");
|
||||
assert_ne!(
|
||||
fp_a1, fp_b,
|
||||
"different paths must produce different fingerprints"
|
||||
);
|
||||
assert_eq!(fp_a1.len(), 16, "fingerprint must be a 16-char hex string");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_from_cwd_isolates_sessions_by_workspace() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
let workspace_a = base.join("repo-alpha");
|
||||
let workspace_b = base.join("repo-beta");
|
||||
fs::create_dir_all(&workspace_a).expect("workspace a should exist");
|
||||
fs::create_dir_all(&workspace_b).expect("workspace b should exist");
|
||||
|
||||
let store_a = SessionStore::from_cwd(&workspace_a).expect("store a should build");
|
||||
let store_b = SessionStore::from_cwd(&workspace_b).expect("store b should build");
|
||||
|
||||
// when
|
||||
let session_a = persist_session_via_store(&store_a, "alpha work");
|
||||
let _session_b = persist_session_via_store(&store_b, "beta work");
|
||||
|
||||
// then — each store only sees its own sessions
|
||||
let list_a = store_a.list_sessions().expect("list a");
|
||||
let list_b = store_b.list_sessions().expect("list b");
|
||||
assert_eq!(list_a.len(), 1, "store a should see exactly one session");
|
||||
assert_eq!(list_b.len(), 1, "store b should see exactly one session");
|
||||
assert_eq!(list_a[0].id, session_a.session_id);
|
||||
assert_ne!(
|
||||
store_a.sessions_dir(),
|
||||
store_b.sessions_dir(),
|
||||
"session directories must differ across workspaces"
|
||||
);
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_from_data_dir_namespaces_by_workspace() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
let data_dir = base.join("global-data");
|
||||
let workspace_a = PathBuf::from("/tmp/project-one");
|
||||
let workspace_b = PathBuf::from("/tmp/project-two");
|
||||
fs::create_dir_all(&data_dir).expect("data dir should exist");
|
||||
|
||||
let store_a =
|
||||
SessionStore::from_data_dir(&data_dir, &workspace_a).expect("store a should build");
|
||||
let store_b =
|
||||
SessionStore::from_data_dir(&data_dir, &workspace_b).expect("store b should build");
|
||||
|
||||
// when
|
||||
persist_session_via_store(&store_a, "work in project-one");
|
||||
persist_session_via_store(&store_b, "work in project-two");
|
||||
|
||||
// then
|
||||
assert_ne!(
|
||||
store_a.sessions_dir(),
|
||||
store_b.sessions_dir(),
|
||||
"data-dir stores must namespace by workspace"
|
||||
);
|
||||
assert_eq!(store_a.list_sessions().expect("list a").len(), 1);
|
||||
assert_eq!(store_b.list_sessions().expect("list b").len(), 1);
|
||||
assert_eq!(store_a.workspace_root(), workspace_a.as_path());
|
||||
assert_eq!(store_b.workspace_root(), workspace_b.as_path());
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_create_and_load_round_trip() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
fs::create_dir_all(&base).expect("base dir should exist");
|
||||
let store = SessionStore::from_cwd(&base).expect("store should build");
|
||||
let session = persist_session_via_store(&store, "round-trip message");
|
||||
|
||||
// when
|
||||
let loaded = store
|
||||
.load_session(&session.session_id)
|
||||
.expect("session should load via store");
|
||||
|
||||
// then
|
||||
assert_eq!(loaded.handle.id, session.session_id);
|
||||
assert_eq!(loaded.session.messages.len(), 1);
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_latest_and_resolve_reference() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
fs::create_dir_all(&base).expect("base dir should exist");
|
||||
let store = SessionStore::from_cwd(&base).expect("store should build");
|
||||
let _older = persist_session_via_store(&store, "older");
|
||||
wait_for_next_millisecond();
|
||||
let newer = persist_session_via_store(&store, "newer");
|
||||
|
||||
// when
|
||||
let latest = store.latest_session().expect("latest should resolve");
|
||||
let handle = store
|
||||
.resolve_reference("latest")
|
||||
.expect("latest alias should resolve");
|
||||
|
||||
// then
|
||||
assert_eq!(latest.id, newer.session_id);
|
||||
assert_eq!(handle.id, newer.session_id);
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_fork_stays_in_same_namespace() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
fs::create_dir_all(&base).expect("base dir should exist");
|
||||
let store = SessionStore::from_cwd(&base).expect("store should build");
|
||||
let source = persist_session_via_store(&store, "parent work");
|
||||
|
||||
// when
|
||||
let forked = store
|
||||
.fork_session(&source, Some("bugfix".to_string()))
|
||||
.expect("fork should succeed");
|
||||
let sessions = store.list_sessions().expect("list sessions");
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
sessions.len(),
|
||||
2,
|
||||
"forked session must land in the same namespace"
|
||||
);
|
||||
assert_eq!(forked.parent_session_id, source.session_id);
|
||||
assert_eq!(forked.branch_name.as_deref(), Some("bugfix"));
|
||||
assert!(
|
||||
forked.handle.path.starts_with(store.sessions_dir()),
|
||||
"forked session path must be inside the store namespace"
|
||||
);
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
}
|
||||
@@ -80,7 +80,11 @@ impl IncrementalSseParser {
|
||||
}
|
||||
|
||||
fn take_event(&mut self) -> Option<SseEvent> {
|
||||
if self.data_lines.is_empty() && self.event_name.is_none() && self.id.is_none() && self.retry.is_none() {
|
||||
if self.data_lines.is_empty()
|
||||
&& self.event_name.is_none()
|
||||
&& self.id.is_none()
|
||||
&& self.retry.is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
@@ -102,8 +106,13 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn parses_streaming_events() {
|
||||
// given
|
||||
let mut parser = IncrementalSseParser::new();
|
||||
|
||||
// when
|
||||
let first = parser.push_chunk("event: message\ndata: hel");
|
||||
|
||||
// then
|
||||
assert!(first.is_empty());
|
||||
|
||||
let second = parser.push_chunk("lo\n\nid: 1\ndata: world\n\n");
|
||||
@@ -125,4 +134,25 @@ mod tests {
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finish_flushes_a_trailing_event_without_separator() {
|
||||
// given
|
||||
let mut parser = IncrementalSseParser::new();
|
||||
parser.push_chunk("event: message\ndata: trailing");
|
||||
|
||||
// when
|
||||
let events = parser.finish();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![SseEvent {
|
||||
event: Some("message".to_string()),
|
||||
data: "trailing".to_string(),
|
||||
id: None,
|
||||
retry: None,
|
||||
}]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,429 @@
|
||||
#![allow(clippy::must_use_candidate)]
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
/// Outcome of comparing the worktree HEAD against the expected base commit.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum BaseCommitState {
|
||||
/// HEAD matches the expected base commit.
|
||||
Matches,
|
||||
/// HEAD has diverged from the expected base.
|
||||
Diverged { expected: String, actual: String },
|
||||
/// No expected base was supplied (neither flag nor file).
|
||||
NoExpectedBase,
|
||||
/// The working directory is not inside a git repository.
|
||||
NotAGitRepo,
|
||||
}
|
||||
|
||||
/// Where the expected base commit originated from.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum BaseCommitSource {
|
||||
Flag(String),
|
||||
File(String),
|
||||
}
|
||||
|
||||
/// Read the `.claw-base` file from the given directory and return the trimmed
|
||||
/// commit hash, or `None` when the file is absent or empty.
|
||||
pub fn read_claw_base_file(cwd: &Path) -> Option<String> {
|
||||
let path = cwd.join(".claw-base");
|
||||
let content = std::fs::read_to_string(path).ok()?;
|
||||
let trimmed = content.trim();
|
||||
if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(trimmed.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the expected base commit: prefer the `--base-commit` flag value,
|
||||
/// fall back to reading `.claw-base` from `cwd`.
|
||||
pub fn resolve_expected_base(flag_value: Option<&str>, cwd: &Path) -> Option<BaseCommitSource> {
|
||||
if let Some(value) = flag_value {
|
||||
let trimmed = value.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return Some(BaseCommitSource::Flag(trimmed.to_string()));
|
||||
}
|
||||
}
|
||||
read_claw_base_file(cwd).map(BaseCommitSource::File)
|
||||
}
|
||||
|
||||
/// Verify that the worktree HEAD matches `expected_base`.
|
||||
///
|
||||
/// Returns [`BaseCommitState::NoExpectedBase`] when no expected commit is
|
||||
/// provided (the check is effectively a no-op in that case).
|
||||
pub fn check_base_commit(cwd: &Path, expected_base: Option<&BaseCommitSource>) -> BaseCommitState {
|
||||
let Some(source) = expected_base else {
|
||||
return BaseCommitState::NoExpectedBase;
|
||||
};
|
||||
let expected_raw = match source {
|
||||
BaseCommitSource::Flag(value) | BaseCommitSource::File(value) => value.as_str(),
|
||||
};
|
||||
|
||||
let Some(head_sha) = resolve_head_sha(cwd) else {
|
||||
return BaseCommitState::NotAGitRepo;
|
||||
};
|
||||
|
||||
let Some(expected_sha) = resolve_rev(cwd, expected_raw) else {
|
||||
// If the expected ref cannot be resolved, compare raw strings as a
|
||||
// best-effort fallback (e.g. partial SHA provided by the caller).
|
||||
return if head_sha.starts_with(expected_raw) || expected_raw.starts_with(&head_sha) {
|
||||
BaseCommitState::Matches
|
||||
} else {
|
||||
BaseCommitState::Diverged {
|
||||
expected: expected_raw.to_string(),
|
||||
actual: head_sha,
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
if head_sha == expected_sha {
|
||||
BaseCommitState::Matches
|
||||
} else {
|
||||
BaseCommitState::Diverged {
|
||||
expected: expected_sha,
|
||||
actual: head_sha,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Format a human-readable warning when the base commit has diverged.
|
||||
///
|
||||
/// Returns `None` for non-warning states (`Matches`, `NoExpectedBase`).
|
||||
pub fn format_stale_base_warning(state: &BaseCommitState) -> Option<String> {
|
||||
match state {
|
||||
BaseCommitState::Diverged { expected, actual } => Some(format!(
|
||||
"warning: worktree HEAD ({actual}) does not match expected base commit ({expected}). \
|
||||
Session may run against a stale codebase."
|
||||
)),
|
||||
BaseCommitState::NotAGitRepo => {
|
||||
Some("warning: stale-base check skipped — not inside a git repository.".to_string())
|
||||
}
|
||||
BaseCommitState::Matches | BaseCommitState::NoExpectedBase => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_head_sha(cwd: &Path) -> Option<String> {
|
||||
resolve_rev(cwd, "HEAD")
|
||||
}
|
||||
|
||||
fn resolve_rev(cwd: &Path, rev: &str) -> Option<String> {
|
||||
let output = Command::new("git")
|
||||
.args(["rev-parse", rev])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.ok()?;
|
||||
if !output.status.success() {
|
||||
return None;
|
||||
}
|
||||
let sha = String::from_utf8(output.stdout).ok()?;
|
||||
let trimmed = sha.trim();
|
||||
if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(trimmed.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use std::process::Command;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_dir() -> 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-stale-base-{nanos}"))
|
||||
}
|
||||
|
||||
fn init_repo(path: &std::path::Path) {
|
||||
fs::create_dir_all(path).expect("create repo dir");
|
||||
run(path, &["init", "--quiet", "-b", "main"]);
|
||||
run(path, &["config", "user.email", "tests@example.com"]);
|
||||
run(path, &["config", "user.name", "Stale Base Tests"]);
|
||||
fs::write(path.join("init.txt"), "initial\n").expect("write init file");
|
||||
run(path, &["add", "."]);
|
||||
run(path, &["commit", "-m", "initial commit", "--quiet"]);
|
||||
}
|
||||
|
||||
fn run(cwd: &std::path::Path, args: &[&str]) {
|
||||
let status = Command::new("git")
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.status()
|
||||
.unwrap_or_else(|e| panic!("git {} failed to execute: {e}", args.join(" ")));
|
||||
assert!(
|
||||
status.success(),
|
||||
"git {} exited with {status}",
|
||||
args.join(" ")
|
||||
);
|
||||
}
|
||||
|
||||
fn commit_file(repo: &std::path::Path, name: &str, msg: &str) {
|
||||
fs::write(repo.join(name), format!("{msg}\n")).expect("write file");
|
||||
run(repo, &["add", name]);
|
||||
run(repo, &["commit", "-m", msg, "--quiet"]);
|
||||
}
|
||||
|
||||
fn head_sha(repo: &std::path::Path) -> String {
|
||||
let output = Command::new("git")
|
||||
.args(["rev-parse", "HEAD"])
|
||||
.current_dir(repo)
|
||||
.output()
|
||||
.expect("git rev-parse HEAD");
|
||||
String::from_utf8(output.stdout)
|
||||
.expect("valid utf8")
|
||||
.trim()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_when_head_equals_expected_base() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
let sha = head_sha(&root);
|
||||
let source = BaseCommitSource::Flag(sha);
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, Some(&source));
|
||||
|
||||
// then
|
||||
assert_eq!(state, BaseCommitState::Matches);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diverged_when_head_moved_past_expected_base() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
let old_sha = head_sha(&root);
|
||||
commit_file(&root, "extra.txt", "move head forward");
|
||||
let new_sha = head_sha(&root);
|
||||
let source = BaseCommitSource::Flag(old_sha.clone());
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, Some(&source));
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
state,
|
||||
BaseCommitState::Diverged {
|
||||
expected: old_sha,
|
||||
actual: new_sha,
|
||||
}
|
||||
);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_expected_base_when_source_is_none() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, None);
|
||||
|
||||
// then
|
||||
assert_eq!(state, BaseCommitState::NoExpectedBase);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn not_a_git_repo_when_outside_repo() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
let source = BaseCommitSource::Flag("abc1234".to_string());
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, Some(&source));
|
||||
|
||||
// then
|
||||
assert_eq!(state, BaseCommitState::NotAGitRepo);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_claw_base_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
fs::write(root.join(".claw-base"), "abc1234def5678\n").expect("write .claw-base");
|
||||
|
||||
// when
|
||||
let value = read_claw_base_file(&root);
|
||||
|
||||
// then
|
||||
assert_eq!(value, Some("abc1234def5678".to_string()));
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_for_missing_claw_base_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
|
||||
// when
|
||||
let value = read_claw_base_file(&root);
|
||||
|
||||
// then
|
||||
assert!(value.is_none());
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_for_empty_claw_base_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
fs::write(root.join(".claw-base"), " \n").expect("write empty .claw-base");
|
||||
|
||||
// when
|
||||
let value = read_claw_base_file(&root);
|
||||
|
||||
// then
|
||||
assert!(value.is_none());
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_expected_base_prefers_flag_over_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
fs::write(root.join(".claw-base"), "from_file\n").expect("write .claw-base");
|
||||
|
||||
// when
|
||||
let source = resolve_expected_base(Some("from_flag"), &root);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
source,
|
||||
Some(BaseCommitSource::Flag("from_flag".to_string()))
|
||||
);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_expected_base_falls_back_to_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
fs::write(root.join(".claw-base"), "from_file\n").expect("write .claw-base");
|
||||
|
||||
// when
|
||||
let source = resolve_expected_base(None, &root);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
source,
|
||||
Some(BaseCommitSource::File("from_file".to_string()))
|
||||
);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_expected_base_returns_none_when_nothing_available() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
|
||||
// when
|
||||
let source = resolve_expected_base(None, &root);
|
||||
|
||||
// then
|
||||
assert!(source.is_none());
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_warning_returns_message_for_diverged() {
|
||||
// given
|
||||
let state = BaseCommitState::Diverged {
|
||||
expected: "abc1234".to_string(),
|
||||
actual: "def5678".to_string(),
|
||||
};
|
||||
|
||||
// when
|
||||
let warning = format_stale_base_warning(&state);
|
||||
|
||||
// then
|
||||
let message = warning.expect("should produce warning");
|
||||
assert!(message.contains("abc1234"));
|
||||
assert!(message.contains("def5678"));
|
||||
assert!(message.contains("stale codebase"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_warning_returns_none_for_matches() {
|
||||
// given
|
||||
let state = BaseCommitState::Matches;
|
||||
|
||||
// when
|
||||
let warning = format_stale_base_warning(&state);
|
||||
|
||||
// then
|
||||
assert!(warning.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_warning_returns_none_for_no_expected_base() {
|
||||
// given
|
||||
let state = BaseCommitState::NoExpectedBase;
|
||||
|
||||
// when
|
||||
let warning = format_stale_base_warning(&state);
|
||||
|
||||
// then
|
||||
assert!(warning.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_with_claw_base_file_in_real_repo() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
let sha = head_sha(&root);
|
||||
fs::write(root.join(".claw-base"), format!("{sha}\n")).expect("write .claw-base");
|
||||
let source = resolve_expected_base(None, &root);
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, source.as_ref());
|
||||
|
||||
// then
|
||||
assert_eq!(state, BaseCommitState::Matches);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diverged_with_claw_base_file_after_new_commit() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
let old_sha = head_sha(&root);
|
||||
fs::write(root.join(".claw-base"), format!("{old_sha}\n")).expect("write .claw-base");
|
||||
commit_file(&root, "new.txt", "advance head");
|
||||
let new_sha = head_sha(&root);
|
||||
let source = resolve_expected_base(None, &root);
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, source.as_ref());
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
state,
|
||||
BaseCommitState::Diverged {
|
||||
expected: old_sha,
|
||||
actual: new_sha,
|
||||
}
|
||||
);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,417 @@
|
||||
#![allow(clippy::must_use_candidate)]
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum BranchFreshness {
|
||||
Fresh,
|
||||
Stale {
|
||||
commits_behind: usize,
|
||||
missing_fixes: Vec<String>,
|
||||
},
|
||||
Diverged {
|
||||
ahead: usize,
|
||||
behind: usize,
|
||||
missing_fixes: Vec<String>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum StaleBranchPolicy {
|
||||
AutoRebase,
|
||||
AutoMergeForward,
|
||||
WarnOnly,
|
||||
Block,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum StaleBranchEvent {
|
||||
BranchStaleAgainstMain {
|
||||
branch: String,
|
||||
commits_behind: usize,
|
||||
missing_fixes: Vec<String>,
|
||||
},
|
||||
RebaseAttempted {
|
||||
branch: String,
|
||||
result: String,
|
||||
},
|
||||
MergeForwardAttempted {
|
||||
branch: String,
|
||||
result: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum StaleBranchAction {
|
||||
Noop,
|
||||
Warn { message: String },
|
||||
Block { message: String },
|
||||
Rebase,
|
||||
MergeForward,
|
||||
}
|
||||
|
||||
pub fn check_freshness(branch: &str, main_ref: &str) -> BranchFreshness {
|
||||
check_freshness_in(branch, main_ref, Path::new("."))
|
||||
}
|
||||
|
||||
pub fn apply_policy(freshness: &BranchFreshness, policy: StaleBranchPolicy) -> StaleBranchAction {
|
||||
match freshness {
|
||||
BranchFreshness::Fresh => StaleBranchAction::Noop,
|
||||
BranchFreshness::Stale {
|
||||
commits_behind,
|
||||
missing_fixes,
|
||||
} => match policy {
|
||||
StaleBranchPolicy::WarnOnly => StaleBranchAction::Warn {
|
||||
message: format!(
|
||||
"Branch is {commits_behind} commit(s) behind main. Missing fixes: {}",
|
||||
if missing_fixes.is_empty() {
|
||||
"(none)".to_string()
|
||||
} else {
|
||||
missing_fixes.join("; ")
|
||||
}
|
||||
),
|
||||
},
|
||||
StaleBranchPolicy::Block => StaleBranchAction::Block {
|
||||
message: format!(
|
||||
"Branch is {commits_behind} commit(s) behind main and must be updated before proceeding."
|
||||
),
|
||||
},
|
||||
StaleBranchPolicy::AutoRebase => StaleBranchAction::Rebase,
|
||||
StaleBranchPolicy::AutoMergeForward => StaleBranchAction::MergeForward,
|
||||
},
|
||||
BranchFreshness::Diverged {
|
||||
ahead,
|
||||
behind,
|
||||
missing_fixes,
|
||||
} => match policy {
|
||||
StaleBranchPolicy::WarnOnly => StaleBranchAction::Warn {
|
||||
message: format!(
|
||||
"Branch has diverged: {ahead} commit(s) ahead, {behind} commit(s) behind main. Missing fixes: {}",
|
||||
format_missing_fixes(missing_fixes)
|
||||
),
|
||||
},
|
||||
StaleBranchPolicy::Block => StaleBranchAction::Block {
|
||||
message: format!(
|
||||
"Branch has diverged ({ahead} ahead, {behind} behind) and must be reconciled before proceeding. Missing fixes: {}",
|
||||
format_missing_fixes(missing_fixes)
|
||||
),
|
||||
},
|
||||
StaleBranchPolicy::AutoRebase => StaleBranchAction::Rebase,
|
||||
StaleBranchPolicy::AutoMergeForward => StaleBranchAction::MergeForward,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn check_freshness_in(
|
||||
branch: &str,
|
||||
main_ref: &str,
|
||||
repo_path: &Path,
|
||||
) -> BranchFreshness {
|
||||
let behind = rev_list_count(main_ref, branch, repo_path);
|
||||
let ahead = rev_list_count(branch, main_ref, repo_path);
|
||||
|
||||
if behind == 0 {
|
||||
return BranchFreshness::Fresh;
|
||||
}
|
||||
|
||||
if ahead > 0 {
|
||||
return BranchFreshness::Diverged {
|
||||
ahead,
|
||||
behind,
|
||||
missing_fixes: missing_fix_subjects(main_ref, branch, repo_path),
|
||||
};
|
||||
}
|
||||
|
||||
let missing_fixes = missing_fix_subjects(main_ref, branch, repo_path);
|
||||
BranchFreshness::Stale {
|
||||
commits_behind: behind,
|
||||
missing_fixes,
|
||||
}
|
||||
}
|
||||
|
||||
fn format_missing_fixes(missing_fixes: &[String]) -> String {
|
||||
if missing_fixes.is_empty() {
|
||||
"(none)".to_string()
|
||||
} else {
|
||||
missing_fixes.join("; ")
|
||||
}
|
||||
}
|
||||
|
||||
fn rev_list_count(a: &str, b: &str, repo_path: &Path) -> usize {
|
||||
let output = Command::new("git")
|
||||
.args(["rev-list", "--count", &format!("{b}..{a}")])
|
||||
.current_dir(repo_path)
|
||||
.output();
|
||||
match output {
|
||||
Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
|
||||
.trim()
|
||||
.parse::<usize>()
|
||||
.unwrap_or(0),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn missing_fix_subjects(a: &str, b: &str, repo_path: &Path) -> Vec<String> {
|
||||
let output = Command::new("git")
|
||||
.args(["log", "--format=%s", &format!("{b}..{a}")])
|
||||
.current_dir(repo_path)
|
||||
.output();
|
||||
match output {
|
||||
Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
|
||||
.lines()
|
||||
.filter(|l| !l.is_empty())
|
||||
.map(String::from)
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_dir() -> 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-stale-branch-{nanos}"))
|
||||
}
|
||||
|
||||
fn init_repo(path: &Path) {
|
||||
fs::create_dir_all(path).expect("create repo dir");
|
||||
run(path, &["init", "--quiet", "-b", "main"]);
|
||||
run(path, &["config", "user.email", "tests@example.com"]);
|
||||
run(path, &["config", "user.name", "Stale Branch Tests"]);
|
||||
fs::write(path.join("init.txt"), "initial\n").expect("write init file");
|
||||
run(path, &["add", "."]);
|
||||
run(path, &["commit", "-m", "initial commit", "--quiet"]);
|
||||
}
|
||||
|
||||
fn run(cwd: &Path, args: &[&str]) {
|
||||
let status = Command::new("git")
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.status()
|
||||
.unwrap_or_else(|e| panic!("git {} failed to execute: {e}", args.join(" ")));
|
||||
assert!(
|
||||
status.success(),
|
||||
"git {} exited with {status}",
|
||||
args.join(" ")
|
||||
);
|
||||
}
|
||||
|
||||
fn commit_file(repo: &Path, name: &str, msg: &str) {
|
||||
fs::write(repo.join(name), format!("{msg}\n")).expect("write file");
|
||||
run(repo, &["add", name]);
|
||||
run(repo, &["commit", "-m", msg, "--quiet"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fresh_branch_passes() {
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// given
|
||||
run(&root, &["checkout", "-b", "topic"]);
|
||||
|
||||
// when
|
||||
let freshness = check_freshness_in("topic", "main", &root);
|
||||
|
||||
// then
|
||||
assert_eq!(freshness, BranchFreshness::Fresh);
|
||||
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fresh_branch_ahead_of_main_still_fresh() {
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// given
|
||||
run(&root, &["checkout", "-b", "topic"]);
|
||||
commit_file(&root, "feature.txt", "add feature");
|
||||
|
||||
// when
|
||||
let freshness = check_freshness_in("topic", "main", &root);
|
||||
|
||||
// then
|
||||
assert_eq!(freshness, BranchFreshness::Fresh);
|
||||
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_detected_with_correct_behind_count_and_missing_fixes() {
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// given
|
||||
run(&root, &["checkout", "-b", "topic"]);
|
||||
run(&root, &["checkout", "main"]);
|
||||
commit_file(&root, "fix1.txt", "fix: resolve timeout");
|
||||
commit_file(&root, "fix2.txt", "fix: handle null pointer");
|
||||
|
||||
// when
|
||||
let freshness = check_freshness_in("topic", "main", &root);
|
||||
|
||||
// then
|
||||
match freshness {
|
||||
BranchFreshness::Stale {
|
||||
commits_behind,
|
||||
missing_fixes,
|
||||
} => {
|
||||
assert_eq!(commits_behind, 2);
|
||||
assert_eq!(missing_fixes.len(), 2);
|
||||
assert_eq!(missing_fixes[0], "fix: handle null pointer");
|
||||
assert_eq!(missing_fixes[1], "fix: resolve timeout");
|
||||
}
|
||||
other => panic!("expected Stale, got {other:?}"),
|
||||
}
|
||||
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diverged_branch_detection() {
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// given
|
||||
run(&root, &["checkout", "-b", "topic"]);
|
||||
commit_file(&root, "topic_work.txt", "topic work");
|
||||
run(&root, &["checkout", "main"]);
|
||||
commit_file(&root, "main_fix.txt", "main fix");
|
||||
|
||||
// when
|
||||
let freshness = check_freshness_in("topic", "main", &root);
|
||||
|
||||
// then
|
||||
match freshness {
|
||||
BranchFreshness::Diverged {
|
||||
ahead,
|
||||
behind,
|
||||
missing_fixes,
|
||||
} => {
|
||||
assert_eq!(ahead, 1);
|
||||
assert_eq!(behind, 1);
|
||||
assert_eq!(missing_fixes, vec!["main fix".to_string()]);
|
||||
}
|
||||
other => panic!("expected Diverged, got {other:?}"),
|
||||
}
|
||||
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_noop_for_fresh_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Fresh;
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::WarnOnly);
|
||||
|
||||
// then
|
||||
assert_eq!(action, StaleBranchAction::Noop);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_warn_for_stale_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Stale {
|
||||
commits_behind: 3,
|
||||
missing_fixes: vec!["fix: timeout".into(), "fix: null ptr".into()],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::WarnOnly);
|
||||
|
||||
// then
|
||||
match action {
|
||||
StaleBranchAction::Warn { message } => {
|
||||
assert!(message.contains("3 commit(s) behind"));
|
||||
assert!(message.contains("fix: timeout"));
|
||||
assert!(message.contains("fix: null ptr"));
|
||||
}
|
||||
other => panic!("expected Warn, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_block_for_stale_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Stale {
|
||||
commits_behind: 1,
|
||||
missing_fixes: vec!["hotfix".into()],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::Block);
|
||||
|
||||
// then
|
||||
match action {
|
||||
StaleBranchAction::Block { message } => {
|
||||
assert!(message.contains("1 commit(s) behind"));
|
||||
}
|
||||
other => panic!("expected Block, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_auto_rebase_for_stale_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Stale {
|
||||
commits_behind: 2,
|
||||
missing_fixes: vec![],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::AutoRebase);
|
||||
|
||||
// then
|
||||
assert_eq!(action, StaleBranchAction::Rebase);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_auto_merge_forward_for_diverged_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Diverged {
|
||||
ahead: 5,
|
||||
behind: 2,
|
||||
missing_fixes: vec!["fix: merge main".into()],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::AutoMergeForward);
|
||||
|
||||
// then
|
||||
assert_eq!(action, StaleBranchAction::MergeForward);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_warn_for_diverged_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Diverged {
|
||||
ahead: 3,
|
||||
behind: 1,
|
||||
missing_fixes: vec!["main hotfix".into()],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::WarnOnly);
|
||||
|
||||
// then
|
||||
match action {
|
||||
StaleBranchAction::Warn { message } => {
|
||||
assert!(message.contains("diverged"));
|
||||
assert!(message.contains("3 commit(s) ahead"));
|
||||
assert!(message.contains("1 commit(s) behind"));
|
||||
assert!(message.contains("main hotfix"));
|
||||
}
|
||||
other => panic!("expected Warn, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,300 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
const DEFAULT_MAX_CHARS: usize = 1_200;
|
||||
const DEFAULT_MAX_LINES: usize = 24;
|
||||
const DEFAULT_MAX_LINE_CHARS: usize = 160;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct SummaryCompressionBudget {
|
||||
pub max_chars: usize,
|
||||
pub max_lines: usize,
|
||||
pub max_line_chars: usize,
|
||||
}
|
||||
|
||||
impl Default for SummaryCompressionBudget {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_chars: DEFAULT_MAX_CHARS,
|
||||
max_lines: DEFAULT_MAX_LINES,
|
||||
max_line_chars: DEFAULT_MAX_LINE_CHARS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SummaryCompressionResult {
|
||||
pub summary: String,
|
||||
pub original_chars: usize,
|
||||
pub compressed_chars: usize,
|
||||
pub original_lines: usize,
|
||||
pub compressed_lines: usize,
|
||||
pub removed_duplicate_lines: usize,
|
||||
pub omitted_lines: usize,
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn compress_summary(
|
||||
summary: &str,
|
||||
budget: SummaryCompressionBudget,
|
||||
) -> SummaryCompressionResult {
|
||||
let original_chars = summary.chars().count();
|
||||
let original_lines = summary.lines().count();
|
||||
|
||||
let normalized = normalize_lines(summary, budget.max_line_chars);
|
||||
if normalized.lines.is_empty() || budget.max_chars == 0 || budget.max_lines == 0 {
|
||||
return SummaryCompressionResult {
|
||||
summary: String::new(),
|
||||
original_chars,
|
||||
compressed_chars: 0,
|
||||
original_lines,
|
||||
compressed_lines: 0,
|
||||
removed_duplicate_lines: normalized.removed_duplicate_lines,
|
||||
omitted_lines: normalized.lines.len(),
|
||||
truncated: original_chars > 0,
|
||||
};
|
||||
}
|
||||
|
||||
let selected = select_line_indexes(&normalized.lines, budget);
|
||||
let mut compressed_lines = selected
|
||||
.iter()
|
||||
.map(|index| normalized.lines[*index].clone())
|
||||
.collect::<Vec<_>>();
|
||||
if compressed_lines.is_empty() {
|
||||
compressed_lines.push(truncate_line(&normalized.lines[0], budget.max_chars));
|
||||
}
|
||||
let omitted_lines = normalized
|
||||
.lines
|
||||
.len()
|
||||
.saturating_sub(compressed_lines.len());
|
||||
|
||||
if omitted_lines > 0 {
|
||||
let omission_notice = omission_notice(omitted_lines);
|
||||
push_line_with_budget(&mut compressed_lines, omission_notice, budget);
|
||||
}
|
||||
|
||||
let compressed_summary = compressed_lines.join("\n");
|
||||
|
||||
SummaryCompressionResult {
|
||||
compressed_chars: compressed_summary.chars().count(),
|
||||
compressed_lines: compressed_lines.len(),
|
||||
removed_duplicate_lines: normalized.removed_duplicate_lines,
|
||||
omitted_lines,
|
||||
truncated: compressed_summary != summary.trim(),
|
||||
summary: compressed_summary,
|
||||
original_chars,
|
||||
original_lines,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn compress_summary_text(summary: &str) -> String {
|
||||
compress_summary(summary, SummaryCompressionBudget::default()).summary
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct NormalizedSummary {
|
||||
lines: Vec<String>,
|
||||
removed_duplicate_lines: usize,
|
||||
}
|
||||
|
||||
fn normalize_lines(summary: &str, max_line_chars: usize) -> NormalizedSummary {
|
||||
let mut seen = BTreeSet::new();
|
||||
let mut lines = Vec::new();
|
||||
let mut removed_duplicate_lines = 0;
|
||||
|
||||
for raw_line in summary.lines() {
|
||||
let normalized = collapse_inline_whitespace(raw_line);
|
||||
if normalized.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let truncated = truncate_line(&normalized, max_line_chars);
|
||||
let dedupe_key = dedupe_key(&truncated);
|
||||
if !seen.insert(dedupe_key) {
|
||||
removed_duplicate_lines += 1;
|
||||
continue;
|
||||
}
|
||||
|
||||
lines.push(truncated);
|
||||
}
|
||||
|
||||
NormalizedSummary {
|
||||
lines,
|
||||
removed_duplicate_lines,
|
||||
}
|
||||
}
|
||||
|
||||
fn select_line_indexes(lines: &[String], budget: SummaryCompressionBudget) -> Vec<usize> {
|
||||
let mut selected = BTreeSet::<usize>::new();
|
||||
|
||||
for priority in 0..=3 {
|
||||
for (index, line) in lines.iter().enumerate() {
|
||||
if selected.contains(&index) || line_priority(line) != priority {
|
||||
continue;
|
||||
}
|
||||
|
||||
let candidate = selected
|
||||
.iter()
|
||||
.map(|selected_index| lines[*selected_index].as_str())
|
||||
.chain(std::iter::once(line.as_str()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if candidate.len() > budget.max_lines {
|
||||
continue;
|
||||
}
|
||||
|
||||
if joined_char_count(&candidate) > budget.max_chars {
|
||||
continue;
|
||||
}
|
||||
|
||||
selected.insert(index);
|
||||
}
|
||||
}
|
||||
|
||||
selected.into_iter().collect()
|
||||
}
|
||||
|
||||
fn push_line_with_budget(lines: &mut Vec<String>, line: String, budget: SummaryCompressionBudget) {
|
||||
let candidate = lines
|
||||
.iter()
|
||||
.map(String::as_str)
|
||||
.chain(std::iter::once(line.as_str()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if candidate.len() <= budget.max_lines && joined_char_count(&candidate) <= budget.max_chars {
|
||||
lines.push(line);
|
||||
}
|
||||
}
|
||||
|
||||
fn joined_char_count(lines: &[&str]) -> usize {
|
||||
lines.iter().map(|line| line.chars().count()).sum::<usize>() + lines.len().saturating_sub(1)
|
||||
}
|
||||
|
||||
fn line_priority(line: &str) -> usize {
|
||||
if line == "Summary:" || line == "Conversation summary:" || is_core_detail(line) {
|
||||
0
|
||||
} else if is_section_header(line) {
|
||||
1
|
||||
} else if line.starts_with("- ") || line.starts_with(" - ") {
|
||||
2
|
||||
} else {
|
||||
3
|
||||
}
|
||||
}
|
||||
|
||||
fn is_core_detail(line: &str) -> bool {
|
||||
[
|
||||
"- Scope:",
|
||||
"- Current work:",
|
||||
"- Pending work:",
|
||||
"- Key files referenced:",
|
||||
"- Tools mentioned:",
|
||||
"- Recent user requests:",
|
||||
"- Previously compacted context:",
|
||||
"- Newly compacted context:",
|
||||
]
|
||||
.iter()
|
||||
.any(|prefix| line.starts_with(prefix))
|
||||
}
|
||||
|
||||
fn is_section_header(line: &str) -> bool {
|
||||
line.ends_with(':')
|
||||
}
|
||||
|
||||
fn omission_notice(omitted_lines: usize) -> String {
|
||||
format!("- … {omitted_lines} additional line(s) omitted.")
|
||||
}
|
||||
|
||||
fn collapse_inline_whitespace(line: &str) -> String {
|
||||
line.split_whitespace().collect::<Vec<_>>().join(" ")
|
||||
}
|
||||
|
||||
fn truncate_line(line: &str, max_chars: usize) -> String {
|
||||
if max_chars == 0 || line.chars().count() <= max_chars {
|
||||
return line.to_string();
|
||||
}
|
||||
|
||||
if max_chars == 1 {
|
||||
return "…".to_string();
|
||||
}
|
||||
|
||||
let mut truncated = line
|
||||
.chars()
|
||||
.take(max_chars.saturating_sub(1))
|
||||
.collect::<String>();
|
||||
truncated.push('…');
|
||||
truncated
|
||||
}
|
||||
|
||||
fn dedupe_key(line: &str) -> String {
|
||||
line.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{compress_summary, compress_summary_text, SummaryCompressionBudget};
|
||||
|
||||
#[test]
|
||||
fn collapses_whitespace_and_duplicate_lines() {
|
||||
// given
|
||||
let summary = "Conversation summary:\n\n- Scope: compact earlier messages.\n- Scope: compact earlier messages.\n- Current work: update runtime module.\n";
|
||||
|
||||
// when
|
||||
let result = compress_summary(summary, SummaryCompressionBudget::default());
|
||||
|
||||
// then
|
||||
assert_eq!(result.removed_duplicate_lines, 1);
|
||||
assert!(result
|
||||
.summary
|
||||
.contains("- Scope: compact earlier messages."));
|
||||
assert!(!result.summary.contains(" compact earlier"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_core_lines_when_budget_is_tight() {
|
||||
// given
|
||||
let summary = [
|
||||
"Conversation summary:",
|
||||
"- Scope: 18 earlier messages compacted.",
|
||||
"- Current work: finish summary compression.",
|
||||
"- Key timeline:",
|
||||
" - user: asked for a working implementation.",
|
||||
" - assistant: inspected runtime compaction flow.",
|
||||
" - tool: cargo check succeeded.",
|
||||
]
|
||||
.join("\n");
|
||||
|
||||
// when
|
||||
let result = compress_summary(
|
||||
&summary,
|
||||
SummaryCompressionBudget {
|
||||
max_chars: 120,
|
||||
max_lines: 3,
|
||||
max_line_chars: 80,
|
||||
},
|
||||
);
|
||||
|
||||
// then
|
||||
assert!(result.summary.contains("Conversation summary:"));
|
||||
assert!(result
|
||||
.summary
|
||||
.contains("- Scope: 18 earlier messages compacted."));
|
||||
assert!(result
|
||||
.summary
|
||||
.contains("- Current work: finish summary compression."));
|
||||
assert!(result.omitted_lines > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provides_a_default_text_only_helper() {
|
||||
// given
|
||||
let summary = "Summary:\n\nA short line.";
|
||||
|
||||
// when
|
||||
let compressed = compress_summary_text(summary);
|
||||
|
||||
// then
|
||||
assert_eq!(compressed, "Summary:\nA short line.");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct TaskPacket {
|
||||
pub objective: String,
|
||||
pub scope: String,
|
||||
pub repo: String,
|
||||
pub branch_policy: String,
|
||||
pub acceptance_tests: Vec<String>,
|
||||
pub commit_policy: String,
|
||||
pub reporting_contract: String,
|
||||
pub escalation_policy: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct TaskPacketValidationError {
|
||||
errors: Vec<String>,
|
||||
}
|
||||
|
||||
impl TaskPacketValidationError {
|
||||
#[must_use]
|
||||
pub fn new(errors: Vec<String>) -> Self {
|
||||
Self { errors }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn errors(&self) -> &[String] {
|
||||
&self.errors
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for TaskPacketValidationError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.errors.join("; "))
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for TaskPacketValidationError {}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ValidatedPacket(TaskPacket);
|
||||
|
||||
impl ValidatedPacket {
|
||||
#[must_use]
|
||||
pub fn packet(&self) -> &TaskPacket {
|
||||
&self.0
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn into_inner(self) -> TaskPacket {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
pub fn validate_packet(packet: TaskPacket) -> Result<ValidatedPacket, TaskPacketValidationError> {
|
||||
let mut errors = Vec::new();
|
||||
|
||||
validate_required("objective", &packet.objective, &mut errors);
|
||||
validate_required("scope", &packet.scope, &mut errors);
|
||||
validate_required("repo", &packet.repo, &mut errors);
|
||||
validate_required("branch_policy", &packet.branch_policy, &mut errors);
|
||||
validate_required("commit_policy", &packet.commit_policy, &mut errors);
|
||||
validate_required(
|
||||
"reporting_contract",
|
||||
&packet.reporting_contract,
|
||||
&mut errors,
|
||||
);
|
||||
validate_required("escalation_policy", &packet.escalation_policy, &mut errors);
|
||||
|
||||
for (index, test) in packet.acceptance_tests.iter().enumerate() {
|
||||
if test.trim().is_empty() {
|
||||
errors.push(format!(
|
||||
"acceptance_tests contains an empty value at index {index}"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if errors.is_empty() {
|
||||
Ok(ValidatedPacket(packet))
|
||||
} else {
|
||||
Err(TaskPacketValidationError::new(errors))
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_required(field: &str, value: &str, errors: &mut Vec<String>) {
|
||||
if value.trim().is_empty() {
|
||||
errors.push(format!("{field} must not be empty"));
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn sample_packet() -> TaskPacket {
|
||||
TaskPacket {
|
||||
objective: "Implement typed task packet format".to_string(),
|
||||
scope: "runtime/task system".to_string(),
|
||||
repo: "claw-code-parity".to_string(),
|
||||
branch_policy: "origin/main only".to_string(),
|
||||
acceptance_tests: vec![
|
||||
"cargo build --workspace".to_string(),
|
||||
"cargo test --workspace".to_string(),
|
||||
],
|
||||
commit_policy: "single verified commit".to_string(),
|
||||
reporting_contract: "print build result, test result, commit sha".to_string(),
|
||||
escalation_policy: "stop only on destructive ambiguity".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_packet_passes_validation() {
|
||||
let packet = sample_packet();
|
||||
let validated = validate_packet(packet.clone()).expect("packet should validate");
|
||||
assert_eq!(validated.packet(), &packet);
|
||||
assert_eq!(validated.into_inner(), packet);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_packet_accumulates_errors() {
|
||||
let packet = TaskPacket {
|
||||
objective: " ".to_string(),
|
||||
scope: String::new(),
|
||||
repo: String::new(),
|
||||
branch_policy: "\t".to_string(),
|
||||
acceptance_tests: vec!["ok".to_string(), " ".to_string()],
|
||||
commit_policy: String::new(),
|
||||
reporting_contract: String::new(),
|
||||
escalation_policy: String::new(),
|
||||
};
|
||||
|
||||
let error = validate_packet(packet).expect_err("packet should be rejected");
|
||||
|
||||
assert!(error.errors().len() >= 7);
|
||||
assert!(error
|
||||
.errors()
|
||||
.contains(&"objective must not be empty".to_string()));
|
||||
assert!(error
|
||||
.errors()
|
||||
.contains(&"scope must not be empty".to_string()));
|
||||
assert!(error
|
||||
.errors()
|
||||
.contains(&"repo must not be empty".to_string()));
|
||||
assert!(error
|
||||
.errors()
|
||||
.contains(&"acceptance_tests contains an empty value at index 1".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serialization_roundtrip_preserves_packet() {
|
||||
let packet = sample_packet();
|
||||
let serialized = serde_json::to_string(&packet).expect("packet should serialize");
|
||||
let deserialized: TaskPacket =
|
||||
serde_json::from_str(&serialized).expect("packet should deserialize");
|
||||
assert_eq!(deserialized, packet);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,503 @@
|
||||
#![allow(clippy::must_use_candidate, clippy::unnecessary_map_or)]
|
||||
//! In-memory task registry for sub-agent task lifecycle management.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{validate_packet, TaskPacket, TaskPacketValidationError};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TaskStatus {
|
||||
Created,
|
||||
Running,
|
||||
Completed,
|
||||
Failed,
|
||||
Stopped,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TaskStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Created => write!(f, "created"),
|
||||
Self::Running => write!(f, "running"),
|
||||
Self::Completed => write!(f, "completed"),
|
||||
Self::Failed => write!(f, "failed"),
|
||||
Self::Stopped => write!(f, "stopped"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Task {
|
||||
pub task_id: String,
|
||||
pub prompt: String,
|
||||
pub description: Option<String>,
|
||||
pub task_packet: Option<TaskPacket>,
|
||||
pub status: TaskStatus,
|
||||
pub created_at: u64,
|
||||
pub updated_at: u64,
|
||||
pub messages: Vec<TaskMessage>,
|
||||
pub output: String,
|
||||
pub team_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TaskMessage {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
pub timestamp: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TaskRegistry {
|
||||
inner: Arc<Mutex<RegistryInner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct RegistryInner {
|
||||
tasks: HashMap<String, Task>,
|
||||
counter: u64,
|
||||
}
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
impl TaskRegistry {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn create(&self, prompt: &str, description: Option<&str>) -> Task {
|
||||
self.create_task(prompt.to_owned(), description.map(str::to_owned), None)
|
||||
}
|
||||
|
||||
pub fn create_from_packet(
|
||||
&self,
|
||||
packet: TaskPacket,
|
||||
) -> Result<Task, TaskPacketValidationError> {
|
||||
let packet = validate_packet(packet)?.into_inner();
|
||||
Ok(self.create_task(
|
||||
packet.objective.clone(),
|
||||
Some(packet.scope.clone()),
|
||||
Some(packet),
|
||||
))
|
||||
}
|
||||
|
||||
fn create_task(
|
||||
&self,
|
||||
prompt: String,
|
||||
description: Option<String>,
|
||||
task_packet: Option<TaskPacket>,
|
||||
) -> Task {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner.counter += 1;
|
||||
let ts = now_secs();
|
||||
let task_id = format!("task_{:08x}_{}", ts, inner.counter);
|
||||
let task = Task {
|
||||
task_id: task_id.clone(),
|
||||
prompt,
|
||||
description,
|
||||
task_packet,
|
||||
status: TaskStatus::Created,
|
||||
created_at: ts,
|
||||
updated_at: ts,
|
||||
messages: Vec::new(),
|
||||
output: String::new(),
|
||||
team_id: None,
|
||||
};
|
||||
inner.tasks.insert(task_id, task.clone());
|
||||
task
|
||||
}
|
||||
|
||||
pub fn get(&self, task_id: &str) -> Option<Task> {
|
||||
let inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner.tasks.get(task_id).cloned()
|
||||
}
|
||||
|
||||
pub fn list(&self, status_filter: Option<TaskStatus>) -> Vec<Task> {
|
||||
let inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner
|
||||
.tasks
|
||||
.values()
|
||||
.filter(|t| status_filter.map_or(true, |s| t.status == s))
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn stop(&self, task_id: &str) -> Result<Task, String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
|
||||
match task.status {
|
||||
TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Stopped => {
|
||||
return Err(format!(
|
||||
"task {task_id} is already in terminal state: {}",
|
||||
task.status
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
task.status = TaskStatus::Stopped;
|
||||
task.updated_at = now_secs();
|
||||
Ok(task.clone())
|
||||
}
|
||||
|
||||
pub fn update(&self, task_id: &str, message: &str) -> Result<Task, String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
|
||||
task.messages.push(TaskMessage {
|
||||
role: String::from("user"),
|
||||
content: message.to_owned(),
|
||||
timestamp: now_secs(),
|
||||
});
|
||||
task.updated_at = now_secs();
|
||||
Ok(task.clone())
|
||||
}
|
||||
|
||||
pub fn output(&self, task_id: &str) -> Result<String, String> {
|
||||
let inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
Ok(task.output.clone())
|
||||
}
|
||||
|
||||
pub fn append_output(&self, task_id: &str, output: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
task.output.push_str(output);
|
||||
task.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn set_status(&self, task_id: &str, status: TaskStatus) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
task.status = status;
|
||||
task.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn assign_team(&self, task_id: &str, team_id: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
task.team_id = Some(team_id.to_owned());
|
||||
task.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn remove(&self, task_id: &str) -> Option<Task> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner.tasks.remove(task_id)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn len(&self) -> usize {
|
||||
let inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner.tasks.len()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn creates_and_retrieves_tasks() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Do something", Some("A test task"));
|
||||
assert_eq!(task.status, TaskStatus::Created);
|
||||
assert_eq!(task.prompt, "Do something");
|
||||
assert_eq!(task.description.as_deref(), Some("A test task"));
|
||||
assert_eq!(task.task_packet, None);
|
||||
|
||||
let fetched = registry.get(&task.task_id).expect("task should exist");
|
||||
assert_eq!(fetched.task_id, task.task_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn creates_task_from_packet() {
|
||||
let registry = TaskRegistry::new();
|
||||
let packet = TaskPacket {
|
||||
objective: "Ship task packet support".to_string(),
|
||||
scope: "runtime/task system".to_string(),
|
||||
repo: "claw-code-parity".to_string(),
|
||||
branch_policy: "origin/main only".to_string(),
|
||||
acceptance_tests: vec!["cargo test --workspace".to_string()],
|
||||
commit_policy: "single commit".to_string(),
|
||||
reporting_contract: "print commit sha".to_string(),
|
||||
escalation_policy: "manual escalation".to_string(),
|
||||
};
|
||||
|
||||
let task = registry
|
||||
.create_from_packet(packet.clone())
|
||||
.expect("packet-backed task should be created");
|
||||
|
||||
assert_eq!(task.prompt, packet.objective);
|
||||
assert_eq!(task.description.as_deref(), Some("runtime/task system"));
|
||||
assert_eq!(task.task_packet, Some(packet.clone()));
|
||||
|
||||
let fetched = registry.get(&task.task_id).expect("task should exist");
|
||||
assert_eq!(fetched.task_packet, Some(packet));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lists_tasks_with_optional_filter() {
|
||||
let registry = TaskRegistry::new();
|
||||
registry.create("Task A", None);
|
||||
let task_b = registry.create("Task B", None);
|
||||
registry
|
||||
.set_status(&task_b.task_id, TaskStatus::Running)
|
||||
.expect("set status should succeed");
|
||||
|
||||
let all = registry.list(None);
|
||||
assert_eq!(all.len(), 2);
|
||||
|
||||
let running = registry.list(Some(TaskStatus::Running));
|
||||
assert_eq!(running.len(), 1);
|
||||
assert_eq!(running[0].task_id, task_b.task_id);
|
||||
|
||||
let created = registry.list(Some(TaskStatus::Created));
|
||||
assert_eq!(created.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stops_running_task() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Stoppable", None);
|
||||
registry
|
||||
.set_status(&task.task_id, TaskStatus::Running)
|
||||
.unwrap();
|
||||
|
||||
let stopped = registry.stop(&task.task_id).expect("stop should succeed");
|
||||
assert_eq!(stopped.status, TaskStatus::Stopped);
|
||||
|
||||
// Stopping again should fail
|
||||
let result = registry.stop(&task.task_id);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn updates_task_with_messages() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Messageable", None);
|
||||
let updated = registry
|
||||
.update(&task.task_id, "Here's more context")
|
||||
.expect("update should succeed");
|
||||
assert_eq!(updated.messages.len(), 1);
|
||||
assert_eq!(updated.messages[0].content, "Here's more context");
|
||||
assert_eq!(updated.messages[0].role, "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn appends_and_retrieves_output() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Output task", None);
|
||||
registry
|
||||
.append_output(&task.task_id, "line 1\n")
|
||||
.expect("append should succeed");
|
||||
registry
|
||||
.append_output(&task.task_id, "line 2\n")
|
||||
.expect("append should succeed");
|
||||
|
||||
let output = registry.output(&task.task_id).expect("output should exist");
|
||||
assert_eq!(output, "line 1\nline 2\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assigns_team_and_removes_task() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Team task", None);
|
||||
registry
|
||||
.assign_team(&task.task_id, "team_abc")
|
||||
.expect("assign should succeed");
|
||||
|
||||
let fetched = registry.get(&task.task_id).unwrap();
|
||||
assert_eq!(fetched.team_id.as_deref(), Some("team_abc"));
|
||||
|
||||
let removed = registry.remove(&task.task_id);
|
||||
assert!(removed.is_some());
|
||||
assert!(registry.get(&task.task_id).is_none());
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_operations_on_missing_task() {
|
||||
let registry = TaskRegistry::new();
|
||||
assert!(registry.stop("nonexistent").is_err());
|
||||
assert!(registry.update("nonexistent", "msg").is_err());
|
||||
assert!(registry.output("nonexistent").is_err());
|
||||
assert!(registry.append_output("nonexistent", "data").is_err());
|
||||
assert!(registry
|
||||
.set_status("nonexistent", TaskStatus::Running)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_status_display_all_variants() {
|
||||
// given
|
||||
let cases = [
|
||||
(TaskStatus::Created, "created"),
|
||||
(TaskStatus::Running, "running"),
|
||||
(TaskStatus::Completed, "completed"),
|
||||
(TaskStatus::Failed, "failed"),
|
||||
(TaskStatus::Stopped, "stopped"),
|
||||
];
|
||||
|
||||
// when
|
||||
let rendered: Vec<_> = cases
|
||||
.into_iter()
|
||||
.map(|(status, expected)| (status.to_string(), expected))
|
||||
.collect();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
rendered,
|
||||
vec![
|
||||
("created".to_string(), "created"),
|
||||
("running".to_string(), "running"),
|
||||
("completed".to_string(), "completed"),
|
||||
("failed".to_string(), "failed"),
|
||||
("stopped".to_string(), "stopped"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_rejects_completed_task() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("done", None);
|
||||
registry
|
||||
.set_status(&task.task_id, TaskStatus::Completed)
|
||||
.expect("set status should succeed");
|
||||
|
||||
// when
|
||||
let result = registry.stop(&task.task_id);
|
||||
|
||||
// then
|
||||
let error = result.expect_err("completed task should be rejected");
|
||||
assert!(error.contains("already in terminal state"));
|
||||
assert!(error.contains("completed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_rejects_failed_task() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("failed", None);
|
||||
registry
|
||||
.set_status(&task.task_id, TaskStatus::Failed)
|
||||
.expect("set status should succeed");
|
||||
|
||||
// when
|
||||
let result = registry.stop(&task.task_id);
|
||||
|
||||
// then
|
||||
let error = result.expect_err("failed task should be rejected");
|
||||
assert!(error.contains("already in terminal state"));
|
||||
assert!(error.contains("failed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_succeeds_from_created_state() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("created task", None);
|
||||
|
||||
// when
|
||||
let stopped = registry.stop(&task.task_id).expect("stop should succeed");
|
||||
|
||||
// then
|
||||
assert_eq!(stopped.status, TaskStatus::Stopped);
|
||||
assert!(stopped.updated_at >= task.updated_at);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_registry_is_empty() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
|
||||
// when
|
||||
let all_tasks = registry.list(None);
|
||||
|
||||
// then
|
||||
assert!(registry.is_empty());
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(all_tasks.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_without_description() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
|
||||
// when
|
||||
let task = registry.create("Do the thing", None);
|
||||
|
||||
// then
|
||||
assert!(task.task_id.starts_with("task_"));
|
||||
assert_eq!(task.description, None);
|
||||
assert_eq!(task.task_packet, None);
|
||||
assert!(task.messages.is_empty());
|
||||
assert!(task.output.is_empty());
|
||||
assert_eq!(task.team_id, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_nonexistent_returns_none() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
|
||||
// when
|
||||
let removed = registry.remove("missing");
|
||||
|
||||
// then
|
||||
assert!(removed.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assign_team_rejects_missing_task() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
|
||||
// when
|
||||
let result = registry.assign_team("missing", "team_123");
|
||||
|
||||
// then
|
||||
let error = result.expect_err("missing task should be rejected");
|
||||
assert_eq!(error, "task not found: missing");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,509 @@
|
||||
#![allow(clippy::must_use_candidate)]
|
||||
//! In-memory registries for Team and Cron lifecycle management.
|
||||
//!
|
||||
//! Provides TeamCreate/Delete and CronCreate/Delete/List runtime backing
|
||||
//! to replace the stub implementations in the tools crate.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{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, Serialize, Deserialize)]
|
||||
pub struct Team {
|
||||
pub team_id: String,
|
||||
pub name: String,
|
||||
pub task_ids: Vec<String>,
|
||||
pub status: TeamStatus,
|
||||
pub created_at: u64,
|
||||
pub updated_at: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TeamStatus {
|
||||
Created,
|
||||
Running,
|
||||
Completed,
|
||||
Deleted,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TeamStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Created => write!(f, "created"),
|
||||
Self::Running => write!(f, "running"),
|
||||
Self::Completed => write!(f, "completed"),
|
||||
Self::Deleted => write!(f, "deleted"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TeamRegistry {
|
||||
inner: Arc<Mutex<TeamInner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct TeamInner {
|
||||
teams: HashMap<String, Team>,
|
||||
counter: u64,
|
||||
}
|
||||
|
||||
impl TeamRegistry {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn create(&self, name: &str, task_ids: Vec<String>) -> Team {
|
||||
let mut inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.counter += 1;
|
||||
let ts = now_secs();
|
||||
let team_id = format!("team_{:08x}_{}", ts, inner.counter);
|
||||
let team = Team {
|
||||
team_id: team_id.clone(),
|
||||
name: name.to_owned(),
|
||||
task_ids,
|
||||
status: TeamStatus::Created,
|
||||
created_at: ts,
|
||||
updated_at: ts,
|
||||
};
|
||||
inner.teams.insert(team_id, team.clone());
|
||||
team
|
||||
}
|
||||
|
||||
pub fn get(&self, team_id: &str) -> Option<Team> {
|
||||
let inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.teams.get(team_id).cloned()
|
||||
}
|
||||
|
||||
pub fn list(&self) -> Vec<Team> {
|
||||
let inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.teams.values().cloned().collect()
|
||||
}
|
||||
|
||||
pub fn delete(&self, team_id: &str) -> Result<Team, String> {
|
||||
let mut inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
let team = inner
|
||||
.teams
|
||||
.get_mut(team_id)
|
||||
.ok_or_else(|| format!("team not found: {team_id}"))?;
|
||||
team.status = TeamStatus::Deleted;
|
||||
team.updated_at = now_secs();
|
||||
Ok(team.clone())
|
||||
}
|
||||
|
||||
pub fn remove(&self, team_id: &str) -> Option<Team> {
|
||||
let mut inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.teams.remove(team_id)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn len(&self) -> usize {
|
||||
let inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.teams.len()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CronEntry {
|
||||
pub cron_id: String,
|
||||
pub schedule: String,
|
||||
pub prompt: String,
|
||||
pub description: Option<String>,
|
||||
pub enabled: bool,
|
||||
pub created_at: u64,
|
||||
pub updated_at: u64,
|
||||
pub last_run_at: Option<u64>,
|
||||
pub run_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CronRegistry {
|
||||
inner: Arc<Mutex<CronInner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct CronInner {
|
||||
entries: HashMap<String, CronEntry>,
|
||||
counter: u64,
|
||||
}
|
||||
|
||||
impl CronRegistry {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn create(&self, schedule: &str, prompt: &str, description: Option<&str>) -> CronEntry {
|
||||
let mut inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner.counter += 1;
|
||||
let ts = now_secs();
|
||||
let cron_id = format!("cron_{:08x}_{}", ts, inner.counter);
|
||||
let entry = CronEntry {
|
||||
cron_id: cron_id.clone(),
|
||||
schedule: schedule.to_owned(),
|
||||
prompt: prompt.to_owned(),
|
||||
description: description.map(str::to_owned),
|
||||
enabled: true,
|
||||
created_at: ts,
|
||||
updated_at: ts,
|
||||
last_run_at: None,
|
||||
run_count: 0,
|
||||
};
|
||||
inner.entries.insert(cron_id, entry.clone());
|
||||
entry
|
||||
}
|
||||
|
||||
pub fn get(&self, cron_id: &str) -> Option<CronEntry> {
|
||||
let inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner.entries.get(cron_id).cloned()
|
||||
}
|
||||
|
||||
pub fn list(&self, enabled_only: bool) -> Vec<CronEntry> {
|
||||
let inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner
|
||||
.entries
|
||||
.values()
|
||||
.filter(|e| !enabled_only || e.enabled)
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn delete(&self, cron_id: &str) -> Result<CronEntry, String> {
|
||||
let mut inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner
|
||||
.entries
|
||||
.remove(cron_id)
|
||||
.ok_or_else(|| format!("cron not found: {cron_id}"))
|
||||
}
|
||||
|
||||
/// Disable a cron entry without removing it.
|
||||
pub fn disable(&self, cron_id: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
let entry = inner
|
||||
.entries
|
||||
.get_mut(cron_id)
|
||||
.ok_or_else(|| format!("cron not found: {cron_id}"))?;
|
||||
entry.enabled = false;
|
||||
entry.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Record a cron run.
|
||||
pub fn record_run(&self, cron_id: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
let entry = inner
|
||||
.entries
|
||||
.get_mut(cron_id)
|
||||
.ok_or_else(|| format!("cron not found: {cron_id}"))?;
|
||||
entry.last_run_at = Some(now_secs());
|
||||
entry.run_count += 1;
|
||||
entry.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn len(&self) -> usize {
|
||||
let inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner.entries.len()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── Team tests ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn creates_and_retrieves_team() {
|
||||
let registry = TeamRegistry::new();
|
||||
let team = registry.create("Alpha Squad", vec!["task_001".into(), "task_002".into()]);
|
||||
assert_eq!(team.name, "Alpha Squad");
|
||||
assert_eq!(team.task_ids.len(), 2);
|
||||
assert_eq!(team.status, TeamStatus::Created);
|
||||
|
||||
let fetched = registry.get(&team.team_id).expect("team should exist");
|
||||
assert_eq!(fetched.team_id, team.team_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lists_and_deletes_teams() {
|
||||
let registry = TeamRegistry::new();
|
||||
let t1 = registry.create("Team A", vec![]);
|
||||
let t2 = registry.create("Team B", vec![]);
|
||||
|
||||
let all = registry.list();
|
||||
assert_eq!(all.len(), 2);
|
||||
|
||||
let deleted = registry.delete(&t1.team_id).expect("delete should succeed");
|
||||
assert_eq!(deleted.status, TeamStatus::Deleted);
|
||||
|
||||
// Team is still listable (soft delete)
|
||||
let still_there = registry.get(&t1.team_id).unwrap();
|
||||
assert_eq!(still_there.status, TeamStatus::Deleted);
|
||||
|
||||
// Hard remove
|
||||
registry.remove(&t2.team_id);
|
||||
assert_eq!(registry.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_missing_team_operations() {
|
||||
let registry = TeamRegistry::new();
|
||||
assert!(registry.delete("nonexistent").is_err());
|
||||
assert!(registry.get("nonexistent").is_none());
|
||||
}
|
||||
|
||||
// ── Cron tests ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn creates_and_retrieves_cron() {
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("0 * * * *", "Check status", Some("hourly check"));
|
||||
assert_eq!(entry.schedule, "0 * * * *");
|
||||
assert_eq!(entry.prompt, "Check status");
|
||||
assert!(entry.enabled);
|
||||
assert_eq!(entry.run_count, 0);
|
||||
assert!(entry.last_run_at.is_none());
|
||||
|
||||
let fetched = registry.get(&entry.cron_id).expect("cron should exist");
|
||||
assert_eq!(fetched.cron_id, entry.cron_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lists_with_enabled_filter() {
|
||||
let registry = CronRegistry::new();
|
||||
let c1 = registry.create("* * * * *", "Task 1", None);
|
||||
let c2 = registry.create("0 * * * *", "Task 2", None);
|
||||
registry
|
||||
.disable(&c1.cron_id)
|
||||
.expect("disable should succeed");
|
||||
|
||||
let all = registry.list(false);
|
||||
assert_eq!(all.len(), 2);
|
||||
|
||||
let enabled_only = registry.list(true);
|
||||
assert_eq!(enabled_only.len(), 1);
|
||||
assert_eq!(enabled_only[0].cron_id, c2.cron_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deletes_cron_entry() {
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("* * * * *", "To delete", None);
|
||||
let deleted = registry
|
||||
.delete(&entry.cron_id)
|
||||
.expect("delete should succeed");
|
||||
assert_eq!(deleted.cron_id, entry.cron_id);
|
||||
assert!(registry.get(&entry.cron_id).is_none());
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn records_cron_runs() {
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("*/5 * * * *", "Recurring", None);
|
||||
registry.record_run(&entry.cron_id).unwrap();
|
||||
registry.record_run(&entry.cron_id).unwrap();
|
||||
|
||||
let fetched = registry.get(&entry.cron_id).unwrap();
|
||||
assert_eq!(fetched.run_count, 2);
|
||||
assert!(fetched.last_run_at.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_missing_cron_operations() {
|
||||
let registry = CronRegistry::new();
|
||||
assert!(registry.delete("nonexistent").is_err());
|
||||
assert!(registry.disable("nonexistent").is_err());
|
||||
assert!(registry.record_run("nonexistent").is_err());
|
||||
assert!(registry.get("nonexistent").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn team_status_display_all_variants() {
|
||||
// given
|
||||
let cases = [
|
||||
(TeamStatus::Created, "created"),
|
||||
(TeamStatus::Running, "running"),
|
||||
(TeamStatus::Completed, "completed"),
|
||||
(TeamStatus::Deleted, "deleted"),
|
||||
];
|
||||
|
||||
// when
|
||||
let rendered: Vec<_> = cases
|
||||
.into_iter()
|
||||
.map(|(status, expected)| (status.to_string(), expected))
|
||||
.collect();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
rendered,
|
||||
vec![
|
||||
("created".to_string(), "created"),
|
||||
("running".to_string(), "running"),
|
||||
("completed".to_string(), "completed"),
|
||||
("deleted".to_string(), "deleted"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_team_registry_is_empty() {
|
||||
// given
|
||||
let registry = TeamRegistry::new();
|
||||
|
||||
// when
|
||||
let teams = registry.list();
|
||||
|
||||
// then
|
||||
assert!(registry.is_empty());
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(teams.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn team_remove_nonexistent_returns_none() {
|
||||
// given
|
||||
let registry = TeamRegistry::new();
|
||||
|
||||
// when
|
||||
let removed = registry.remove("missing");
|
||||
|
||||
// then
|
||||
assert!(removed.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn team_len_transitions() {
|
||||
// given
|
||||
let registry = TeamRegistry::new();
|
||||
|
||||
// when
|
||||
let alpha = registry.create("Alpha", vec![]);
|
||||
let beta = registry.create("Beta", vec![]);
|
||||
let after_create = registry.len();
|
||||
registry.remove(&alpha.team_id);
|
||||
let after_first_remove = registry.len();
|
||||
registry.remove(&beta.team_id);
|
||||
|
||||
// then
|
||||
assert_eq!(after_create, 2);
|
||||
assert_eq!(after_first_remove, 1);
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_list_all_disabled_returns_empty_for_enabled_only() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
let first = registry.create("* * * * *", "Task 1", None);
|
||||
let second = registry.create("0 * * * *", "Task 2", None);
|
||||
registry
|
||||
.disable(&first.cron_id)
|
||||
.expect("disable should succeed");
|
||||
registry
|
||||
.disable(&second.cron_id)
|
||||
.expect("disable should succeed");
|
||||
|
||||
// when
|
||||
let enabled_only = registry.list(true);
|
||||
let all_entries = registry.list(false);
|
||||
|
||||
// then
|
||||
assert!(enabled_only.is_empty());
|
||||
assert_eq!(all_entries.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_create_without_description() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
|
||||
// when
|
||||
let entry = registry.create("*/15 * * * *", "Check health", None);
|
||||
|
||||
// then
|
||||
assert!(entry.cron_id.starts_with("cron_"));
|
||||
assert_eq!(entry.description, None);
|
||||
assert!(entry.enabled);
|
||||
assert_eq!(entry.run_count, 0);
|
||||
assert_eq!(entry.last_run_at, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_cron_registry_is_empty() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
|
||||
// when
|
||||
let enabled_only = registry.list(true);
|
||||
let all_entries = registry.list(false);
|
||||
|
||||
// then
|
||||
assert!(registry.is_empty());
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(enabled_only.is_empty());
|
||||
assert!(all_entries.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_record_run_updates_timestamp_and_counter() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("*/5 * * * *", "Recurring", None);
|
||||
|
||||
// when
|
||||
registry
|
||||
.record_run(&entry.cron_id)
|
||||
.expect("first run should succeed");
|
||||
registry
|
||||
.record_run(&entry.cron_id)
|
||||
.expect("second run should succeed");
|
||||
let fetched = registry.get(&entry.cron_id).expect("entry should exist");
|
||||
|
||||
// then
|
||||
assert_eq!(fetched.run_count, 2);
|
||||
assert!(fetched.last_run_at.is_some());
|
||||
assert!(fetched.updated_at >= entry.updated_at);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_disable_updates_timestamp() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("0 0 * * *", "Nightly", None);
|
||||
|
||||
// when
|
||||
registry
|
||||
.disable(&entry.cron_id)
|
||||
.expect("disable should succeed");
|
||||
let fetched = registry.get(&entry.cron_id).expect("entry should exist");
|
||||
|
||||
// then
|
||||
assert!(!fetched.enabled);
|
||||
assert!(fetched.updated_at >= entry.updated_at);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,299 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const TRUST_PROMPT_CUES: &[&str] = &[
|
||||
"do you trust the files in this folder",
|
||||
"trust the files in this folder",
|
||||
"trust this folder",
|
||||
"allow and continue",
|
||||
"yes, proceed",
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum TrustPolicy {
|
||||
AutoTrust,
|
||||
RequireApproval,
|
||||
Deny,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum TrustEvent {
|
||||
TrustRequired { cwd: String },
|
||||
TrustResolved { cwd: String, policy: TrustPolicy },
|
||||
TrustDenied { cwd: String, reason: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TrustConfig {
|
||||
allowlisted: Vec<PathBuf>,
|
||||
denied: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
impl TrustConfig {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_allowlisted(mut self, path: impl Into<PathBuf>) -> Self {
|
||||
self.allowlisted.push(path.into());
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_denied(mut self, path: impl Into<PathBuf>) -> Self {
|
||||
self.denied.push(path.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum TrustDecision {
|
||||
NotRequired,
|
||||
Required {
|
||||
policy: TrustPolicy,
|
||||
events: Vec<TrustEvent>,
|
||||
},
|
||||
}
|
||||
|
||||
impl TrustDecision {
|
||||
#[must_use]
|
||||
pub fn policy(&self) -> Option<TrustPolicy> {
|
||||
match self {
|
||||
Self::NotRequired => None,
|
||||
Self::Required { policy, .. } => Some(*policy),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn events(&self) -> &[TrustEvent] {
|
||||
match self {
|
||||
Self::NotRequired => &[],
|
||||
Self::Required { events, .. } => events,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TrustResolver {
|
||||
config: TrustConfig,
|
||||
}
|
||||
|
||||
impl TrustResolver {
|
||||
#[must_use]
|
||||
pub fn new(config: TrustConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn resolve(&self, cwd: &str, screen_text: &str) -> TrustDecision {
|
||||
if !detect_trust_prompt(screen_text) {
|
||||
return TrustDecision::NotRequired;
|
||||
}
|
||||
|
||||
let mut events = vec![TrustEvent::TrustRequired {
|
||||
cwd: cwd.to_owned(),
|
||||
}];
|
||||
|
||||
if let Some(matched_root) = self
|
||||
.config
|
||||
.denied
|
||||
.iter()
|
||||
.find(|root| path_matches(cwd, root))
|
||||
{
|
||||
let reason = format!("cwd matches denied trust root: {}", matched_root.display());
|
||||
events.push(TrustEvent::TrustDenied {
|
||||
cwd: cwd.to_owned(),
|
||||
reason,
|
||||
});
|
||||
return TrustDecision::Required {
|
||||
policy: TrustPolicy::Deny,
|
||||
events,
|
||||
};
|
||||
}
|
||||
|
||||
if self
|
||||
.config
|
||||
.allowlisted
|
||||
.iter()
|
||||
.any(|root| path_matches(cwd, root))
|
||||
{
|
||||
events.push(TrustEvent::TrustResolved {
|
||||
cwd: cwd.to_owned(),
|
||||
policy: TrustPolicy::AutoTrust,
|
||||
});
|
||||
return TrustDecision::Required {
|
||||
policy: TrustPolicy::AutoTrust,
|
||||
events,
|
||||
};
|
||||
}
|
||||
|
||||
TrustDecision::Required {
|
||||
policy: TrustPolicy::RequireApproval,
|
||||
events,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn trusts(&self, cwd: &str) -> bool {
|
||||
!self
|
||||
.config
|
||||
.denied
|
||||
.iter()
|
||||
.any(|root| path_matches(cwd, root))
|
||||
&& self
|
||||
.config
|
||||
.allowlisted
|
||||
.iter()
|
||||
.any(|root| path_matches(cwd, root))
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn detect_trust_prompt(screen_text: &str) -> bool {
|
||||
let lowered = screen_text.to_ascii_lowercase();
|
||||
TRUST_PROMPT_CUES
|
||||
.iter()
|
||||
.any(|needle| lowered.contains(needle))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn path_matches_trusted_root(cwd: &str, trusted_root: &str) -> bool {
|
||||
path_matches(cwd, &normalize_path(Path::new(trusted_root)))
|
||||
}
|
||||
|
||||
fn path_matches(candidate: &str, root: &Path) -> bool {
|
||||
let candidate = normalize_path(Path::new(candidate));
|
||||
let root = normalize_path(root);
|
||||
candidate == root || candidate.starts_with(&root)
|
||||
}
|
||||
|
||||
fn normalize_path(path: &Path) -> PathBuf {
|
||||
std::fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
detect_trust_prompt, path_matches_trusted_root, TrustConfig, TrustDecision, TrustEvent,
|
||||
TrustPolicy, TrustResolver,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn detects_known_trust_prompt_copy() {
|
||||
// given
|
||||
let screen_text = "Do you trust the files in this folder?\n1. Yes, proceed\n2. No";
|
||||
|
||||
// when
|
||||
let detected = detect_trust_prompt(screen_text);
|
||||
|
||||
// then
|
||||
assert!(detected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_emit_events_when_prompt_is_absent() {
|
||||
// given
|
||||
let resolver = TrustResolver::new(TrustConfig::new().with_allowlisted("/tmp/worktrees"));
|
||||
|
||||
// when
|
||||
let decision = resolver.resolve("/tmp/worktrees/repo-a", "Ready for your input\n>");
|
||||
|
||||
// then
|
||||
assert_eq!(decision, TrustDecision::NotRequired);
|
||||
assert_eq!(decision.events(), &[]);
|
||||
assert_eq!(decision.policy(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_trusts_allowlisted_cwd_after_prompt_detection() {
|
||||
// given
|
||||
let resolver = TrustResolver::new(TrustConfig::new().with_allowlisted("/tmp/worktrees"));
|
||||
|
||||
// when
|
||||
let decision = resolver.resolve(
|
||||
"/tmp/worktrees/repo-a",
|
||||
"Do you trust the files in this folder?\n1. Yes, proceed\n2. No",
|
||||
);
|
||||
|
||||
// then
|
||||
assert_eq!(decision.policy(), Some(TrustPolicy::AutoTrust));
|
||||
assert_eq!(
|
||||
decision.events(),
|
||||
&[
|
||||
TrustEvent::TrustRequired {
|
||||
cwd: "/tmp/worktrees/repo-a".to_string(),
|
||||
},
|
||||
TrustEvent::TrustResolved {
|
||||
cwd: "/tmp/worktrees/repo-a".to_string(),
|
||||
policy: TrustPolicy::AutoTrust,
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requires_approval_for_unknown_cwd_after_prompt_detection() {
|
||||
// given
|
||||
let resolver = TrustResolver::new(TrustConfig::new().with_allowlisted("/tmp/worktrees"));
|
||||
|
||||
// when
|
||||
let decision = resolver.resolve(
|
||||
"/tmp/other/repo-b",
|
||||
"Do you trust the files in this folder?\n1. Yes, proceed\n2. No",
|
||||
);
|
||||
|
||||
// then
|
||||
assert_eq!(decision.policy(), Some(TrustPolicy::RequireApproval));
|
||||
assert_eq!(
|
||||
decision.events(),
|
||||
&[TrustEvent::TrustRequired {
|
||||
cwd: "/tmp/other/repo-b".to_string(),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn denied_root_takes_precedence_over_allowlist() {
|
||||
// given
|
||||
let resolver = TrustResolver::new(
|
||||
TrustConfig::new()
|
||||
.with_allowlisted("/tmp/worktrees")
|
||||
.with_denied("/tmp/worktrees/repo-c"),
|
||||
);
|
||||
|
||||
// when
|
||||
let decision = resolver.resolve(
|
||||
"/tmp/worktrees/repo-c",
|
||||
"Do you trust the files in this folder?\n1. Yes, proceed\n2. No",
|
||||
);
|
||||
|
||||
// then
|
||||
assert_eq!(decision.policy(), Some(TrustPolicy::Deny));
|
||||
assert_eq!(
|
||||
decision.events(),
|
||||
&[
|
||||
TrustEvent::TrustRequired {
|
||||
cwd: "/tmp/worktrees/repo-c".to_string(),
|
||||
},
|
||||
TrustEvent::TrustDenied {
|
||||
cwd: "/tmp/worktrees/repo-c".to_string(),
|
||||
reason: "cwd matches denied trust root: /tmp/worktrees/repo-c".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sibling_prefix_does_not_match_trusted_root() {
|
||||
// given
|
||||
let trusted_root = "/tmp/worktrees";
|
||||
let sibling_path = "/tmp/worktrees-other/repo-d";
|
||||
|
||||
// when
|
||||
let matched = path_matches_trusted_root(sibling_path, trusted_root);
|
||||
|
||||
// then
|
||||
assert!(!matched);
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user