feat: import Chinese-localized Buzz source snapshot
Docker image / Build (linux/amd64) (push) Has been cancelled
Docker image / Build (linux/arm64) (push) Has been cancelled
Docker image / Merge release multi-arch manifest (push) Has been cancelled
Docker image / Merge debug multi-arch manifest (push) Has been cancelled
Docker image / Build public push gateway (linux/amd64) (push) Has been cancelled
Docker image / Build public push gateway (linux/arm64) (push) Has been cancelled
Docker image / Publish public push gateway image (push) Has been cancelled
Sprig image / Build (linux/amd64) (push) Has been cancelled
Sprig image / Build (linux/arm64) (push) Has been cancelled
Sprig image / Merge multi-arch manifest (push) Has been cancelled
Harbor Buzz Orchestra / Python tests and lint (push) Has been cancelled
CI / Detect Changed Paths (push) Has been cancelled
CI / Rust Lint (push) Has been cancelled
CI / Unit Tests (push) Has been cancelled
CI / Desktop Core (push) Has been cancelled
CI / Desktop Smoke E2E (1) (push) Has been cancelled
CI / Desktop Smoke E2E (2) (push) Has been cancelled
CI / Desktop Smoke E2E (3) (push) Has been cancelled
CI / Desktop Smoke E2E (4) (push) Has been cancelled
CI / Desktop (push) Has been cancelled
CI / Desktop E2E Relay (push) Has been cancelled
CI / Desktop E2E Integration (1/2) (push) Has been cancelled
CI / Desktop E2E Integration (2/2) (push) Has been cancelled
CI / Desktop E2E Integration (push) Has been cancelled
CI / Backend Integration (relay e2e) (push) Has been cancelled
CI / Relay E2E (push) Has been cancelled
CI / Web (push) Has been cancelled
CI / Mobile (push) Has been cancelled
CI / Security (push) Has been cancelled
CI / Dead Token Reference Guard (push) Has been cancelled
CI / Server Cross-Compile (aarch64-unknown-linux-musl) (push) Has been cancelled
CI / Server Cross-Compile (x86_64-unknown-linux-musl) (push) Has been cancelled
CI / Windows Rust (x86_64-pc-windows-msvc) (push) Has been cancelled
CI / Desktop Build (macOS) (push) Has been cancelled
helm chart / lint + unittest + render matrix (push) Has been cancelled
helm chart / install on kind (gated) (push) Has been cancelled
helm chart / publish chart to GHCR (push) Has been cancelled
Mesh Lifecycle / Relay-Driven Mesh Lifecycle Smoke (push) Has been cancelled
Sprig / Build (aarch64-unknown-linux-musl) (push) Has been cancelled
Sprig / Build (x86_64-unknown-linux-musl) (push) Has been cancelled
Sprig / Publish rolling release (push) Has been cancelled
Sprig / Publish tagged release (push) Has been cancelled
Docker image / Build (linux/amd64) (push) Has been cancelled
Docker image / Build (linux/arm64) (push) Has been cancelled
Docker image / Merge release multi-arch manifest (push) Has been cancelled
Docker image / Merge debug multi-arch manifest (push) Has been cancelled
Docker image / Build public push gateway (linux/amd64) (push) Has been cancelled
Docker image / Build public push gateway (linux/arm64) (push) Has been cancelled
Docker image / Publish public push gateway image (push) Has been cancelled
Sprig image / Build (linux/amd64) (push) Has been cancelled
Sprig image / Build (linux/arm64) (push) Has been cancelled
Sprig image / Merge multi-arch manifest (push) Has been cancelled
Harbor Buzz Orchestra / Python tests and lint (push) Has been cancelled
CI / Detect Changed Paths (push) Has been cancelled
CI / Rust Lint (push) Has been cancelled
CI / Unit Tests (push) Has been cancelled
CI / Desktop Core (push) Has been cancelled
CI / Desktop Smoke E2E (1) (push) Has been cancelled
CI / Desktop Smoke E2E (2) (push) Has been cancelled
CI / Desktop Smoke E2E (3) (push) Has been cancelled
CI / Desktop Smoke E2E (4) (push) Has been cancelled
CI / Desktop (push) Has been cancelled
CI / Desktop E2E Relay (push) Has been cancelled
CI / Desktop E2E Integration (1/2) (push) Has been cancelled
CI / Desktop E2E Integration (2/2) (push) Has been cancelled
CI / Desktop E2E Integration (push) Has been cancelled
CI / Backend Integration (relay e2e) (push) Has been cancelled
CI / Relay E2E (push) Has been cancelled
CI / Web (push) Has been cancelled
CI / Mobile (push) Has been cancelled
CI / Security (push) Has been cancelled
CI / Dead Token Reference Guard (push) Has been cancelled
CI / Server Cross-Compile (aarch64-unknown-linux-musl) (push) Has been cancelled
CI / Server Cross-Compile (x86_64-unknown-linux-musl) (push) Has been cancelled
CI / Windows Rust (x86_64-pc-windows-msvc) (push) Has been cancelled
CI / Desktop Build (macOS) (push) Has been cancelled
helm chart / lint + unittest + render matrix (push) Has been cancelled
helm chart / install on kind (gated) (push) Has been cancelled
helm chart / publish chart to GHCR (push) Has been cancelled
Mesh Lifecycle / Relay-Driven Mesh Lifecycle Smoke (push) Has been cancelled
Sprig / Build (aarch64-unknown-linux-musl) (push) Has been cancelled
Sprig / Build (x86_64-unknown-linux-musl) (push) Has been cancelled
Sprig / Publish rolling release (push) Has been cancelled
Sprig / Publish tagged release (push) Has been cancelled
Signed-off-by: cls_宁波本机 <908705107@qq.com>
This commit is contained in:
@@ -0,0 +1,58 @@
|
||||
[package]
|
||||
name = "buzz-agent"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
rust-version.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Minimal, unbreakable ACP-compliant agent. Non-streaming. Tool-calls-as-output."
|
||||
readme = "README.md"
|
||||
keywords = ["acp", "agent", "llm", "mcp", "minimal"]
|
||||
categories = ["command-line-utilities", "web-programming"]
|
||||
|
||||
[lib]
|
||||
name = "buzz_agent"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "buzz-agent"
|
||||
path = "src/main.rs"
|
||||
|
||||
# Test-only fake MCP server. Built unconditionally because cargo can't gate
|
||||
# bins on `cfg(test)`, but it's tiny and only used by integration tests.
|
||||
[[bin]]
|
||||
name = "fake-mcp"
|
||||
path = "tests/bin/fake_mcp.rs"
|
||||
|
||||
[dependencies]
|
||||
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-std", "io-util", "sync", "process", "time", "net"] }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
serde_yaml = { workspace = true }
|
||||
reqwest = { workspace = true, features = ["json", "rustls", "form"] }
|
||||
rmcp = { version = "1", default-features = false, features = ["client", "transport-child-process"] }
|
||||
arc-swap = "1"
|
||||
getrandom = "0.4"
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
# OAuth 2.0 PKCE for Databricks (and future browser-auth providers).
|
||||
async-trait = "0.1"
|
||||
axum = { workspace = true }
|
||||
base64 = "0.22"
|
||||
hex = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
urlencoding = "2"
|
||||
webbrowser = "1"
|
||||
dirs = "6"
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
nix = { version = "0.31", default-features = false, features = ["signal", "process"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "macros", "io-std", "io-util", "sync", "process", "time", "net"] }
|
||||
nix = { version = "0.31", default-features = false, features = ["signal", "process"] }
|
||||
axum = { workspace = true }
|
||||
hex = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
tempfile = "3"
|
||||
@@ -0,0 +1,381 @@
|
||||
# buzz-agent
|
||||
|
||||
> Minimal, unbreakable ACP-compliant LLM agent. Stdio in, tool calls out. Non-streaming. No persistence. No cleverness.
|
||||
|
||||
[ACP](https://agentclientprotocol.com) is the Agent Client Protocol — JSON-RPC 2.0 over stdio between a client (Zed, JetBrains, buzz-acp, …) and an agent. [MCP](https://modelcontextprotocol.io) is how the agent talks to its tools.
|
||||
|
||||
`buzz-agent` is the agent.
|
||||
|
||||
## What It Is
|
||||
|
||||
```
|
||||
+--------+ stdio (JSON-RPC 2.0) +---------------+
|
||||
| client | <----------------------> | buzz-agent |
|
||||
+--------+ ACP frames +---------------+
|
||||
│ │
|
||||
│ │ rmcp (stdio)
|
||||
│ ▼
|
||||
│ MCP servers
|
||||
│ (your tools)
|
||||
▼
|
||||
HTTPS
|
||||
│
|
||||
▼
|
||||
Anthropic Messages API,
|
||||
OpenRouter, or any OpenAI-compat
|
||||
(vLLM, llama.cpp, Databricks,
|
||||
Block Gateway, Ollama, …)
|
||||
```
|
||||
|
||||
A client sends `session/prompt`. The agent loops: call the LLM → get tool calls → run them via MCP → feed results back → repeat. The loop terminates when the LLM stops asking for tools, the round cap is hit, or the client cancels.
|
||||
|
||||
The agent's **output is its tool calls**. Generated text is forwarded to the client as `agent_message_chunk` updates, but the real work happens in the tools. The LLM call is non-streaming — one HTTP POST, one response.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Build
|
||||
cargo build --release -p buzz-agent
|
||||
|
||||
# Run against Anthropic
|
||||
BUZZ_AGENT_PROVIDER=anthropic \
|
||||
ANTHROPIC_API_KEY=sk-ant-... \
|
||||
ANTHROPIC_MODEL=claude-sonnet-4-5 \
|
||||
./target/release/buzz-agent
|
||||
|
||||
# Or any OpenAI-compatible endpoint
|
||||
BUZZ_AGENT_PROVIDER=openai \
|
||||
OPENAI_COMPAT_API_KEY=sk-... \
|
||||
OPENAI_COMPAT_MODEL=gpt-5 \
|
||||
OPENAI_COMPAT_BASE_URL=https://api.openai.com/v1 \
|
||||
./target/release/buzz-agent
|
||||
|
||||
# Or OpenRouter
|
||||
BUZZ_AGENT_PROVIDER=openrouter \
|
||||
OPENROUTER_API_KEY=sk-or-v1-... \
|
||||
OPENROUTER_MODEL=anthropic/claude-sonnet-4.5 \
|
||||
./target/release/buzz-agent
|
||||
|
||||
# Or Databricks model serving via OAuth 2.0 PKCE
|
||||
BUZZ_AGENT_PROVIDER=databricks \
|
||||
DATABRICKS_HOST=https://dbc-...cloud.databricks.com \
|
||||
DATABRICKS_MODEL=goose-claude-4-6-sonnet \
|
||||
./target/release/buzz-agent
|
||||
```
|
||||
|
||||
That's the whole setup. The agent reads JSON-RPC frames from stdin, writes them to stdout, and logs to stderr.
|
||||
|
||||
## ACP Transcript
|
||||
|
||||
A complete round-trip. Lines starting with `→` are client→agent (stdin); `←` are agent→client (stdout). Each line is one newline-terminated JSON value. Comments are not part of the wire.
|
||||
|
||||
```jsonc
|
||||
// 1. Handshake.
|
||||
→ {"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":1,"clientCapabilities":{}}}
|
||||
← {"jsonrpc":"2.0","id":1,"result":{
|
||||
"protocolVersion":1,
|
||||
"agentCapabilities":{
|
||||
"loadSession":false,
|
||||
"promptCapabilities":{"image":false,"audio":false,"embeddedContext":false},
|
||||
"mcpCapabilities":{"http":false,"sse":false}
|
||||
},
|
||||
"agentInfo":{"name":"buzz-agent","version":"0.1.0"}
|
||||
}}
|
||||
|
||||
// 2. Open a session. The client passes the MCP servers to spawn.
|
||||
→ {"jsonrpc":"2.0","id":2,"method":"session/new","params":{
|
||||
"cwd":"/tmp",
|
||||
"mcpServers":[{"name":"echo","command":"/usr/local/bin/echo-mcp","args":[],"env":[]}]
|
||||
}}
|
||||
← {"jsonrpc":"2.0","id":2,"result":{"sessionId":"ses_a1b2c3d4e5f6a7b8"}}
|
||||
|
||||
// 3. Prompt. The agent loops until the LLM stops calling tools.
|
||||
→ {"jsonrpc":"2.0","id":3,"method":"session/prompt","params":{
|
||||
"sessionId":"ses_a1b2c3d4e5f6a7b8",
|
||||
"prompt":[{"type":"text","text":"echo hello"}]
|
||||
}}
|
||||
|
||||
// 4. Agent emits tool_call (status: pending) — visible to the UI.
|
||||
← {"jsonrpc":"2.0","method":"session/update","params":{
|
||||
"sessionId":"ses_a1b2c3d4e5f6a7b8",
|
||||
"update":{
|
||||
"sessionUpdate":"tool_call",
|
||||
"toolCallId":"toolu_01XYZ",
|
||||
"title":"echo__say",
|
||||
"kind":"other",
|
||||
"status":"pending",
|
||||
"rawInput":{"text":"hello"}
|
||||
}
|
||||
}}
|
||||
|
||||
// 5. Agent moves the call to in_progress, runs the MCP tool, then completed.
|
||||
← {"jsonrpc":"2.0","method":"session/update","params":{
|
||||
"sessionId":"ses_a1b2c3d4e5f6a7b8",
|
||||
"update":{"sessionUpdate":"tool_call_update","toolCallId":"toolu_01XYZ","status":"in_progress"}
|
||||
}}
|
||||
← {"jsonrpc":"2.0","method":"session/update","params":{
|
||||
"sessionId":"ses_a1b2c3d4e5f6a7b8",
|
||||
"update":{
|
||||
"sessionUpdate":"tool_call_update",
|
||||
"toolCallId":"toolu_01XYZ",
|
||||
"status":"completed",
|
||||
"content":[{"type":"content","content":{"type":"text","text":"hello"}}]
|
||||
}
|
||||
}}
|
||||
|
||||
// 8. The model sees the result, decides it's done, and the prompt resolves.
|
||||
← {"jsonrpc":"2.0","id":3,"result":{"stopReason":"end_turn"}}
|
||||
```
|
||||
|
||||
That's ACP. Three request methods (`initialize`, `session/new`, `session/prompt`), one inbound notification (`session/cancel`), and three outbound update variants (`agent_message_chunk`, `tool_call`, `tool_call_update`). The full server is hand-rolled in `main.rs`.
|
||||
|
||||
## Configuration
|
||||
|
||||
Everything is environment variables. No flags, no config files. (We are a subprocess; subprocess config is environment.)
|
||||
|
||||
| Variable | Default | Notes |
|
||||
|---|---|---|
|
||||
| `BUZZ_AGENT_PROVIDER` | — | Required. `anthropic`, `openai`, `openrouter`, `databricks`, or `databricks_v2`. No implicit fallback — the agent errors at startup when this is unset. |
|
||||
| `ANTHROPIC_API_KEY` | — | Required when provider=anthropic. |
|
||||
| `ANTHROPIC_MODEL` | — | Required when provider=anthropic. |
|
||||
| `ANTHROPIC_BASE_URL` | `https://api.anthropic.com` | |
|
||||
| `ANTHROPIC_API_VERSION` | `2023-06-01` | |
|
||||
| `OPENAI_COMPAT_API_KEY` | — | Required when provider=openai. |
|
||||
| `OPENAI_COMPAT_MODEL` | — | Required when provider=openai. |
|
||||
| `OPENAI_COMPAT_BASE_URL` | `https://api.openai.com/v1` | Point at vLLM, llama.cpp, Ollama, etc. |
|
||||
| `OPENAI_COMPAT_API` | `auto` | `auto` \| `chat` \| `responses`. `auto` picks Responses for `*.openai.com`, Chat Completions everywhere else. |
|
||||
| `OPENROUTER_API_KEY` | — | Required when provider=openrouter. |
|
||||
| `OPENROUTER_MODEL` | — | Required when provider=openrouter. Use OpenRouter's `vendor/model` id, e.g. `anthropic/claude-sonnet-4.5`. |
|
||||
| `OPENROUTER_BASE_URL` | `https://openrouter.ai/api/v1` | |
|
||||
| `DATABRICKS_HOST` | — | Required when provider=databricks or provider=databricks_v2. |
|
||||
| `DATABRICKS_MODEL` | — | Required when provider=databricks or provider=databricks_v2. |
|
||||
| `DATABRICKS_TOKEN` | — | Optional static bearer escape hatch. If unset, Databricks uses browser OAuth + refresh cache. |
|
||||
| `BUZZ_AGENT_SYSTEM_PROMPT` | built-in | Inline system prompt. |
|
||||
| `BUZZ_AGENT_SYSTEM_PROMPT_FILE` | — | File path. Mutually exclusive with the above. |
|
||||
| `BUZZ_AGENT_MAX_ROUNDS` | `0` | Tool-loop iteration cap. 0 = unlimited. |
|
||||
| `BUZZ_AGENT_MAX_OUTPUT_TOKENS` | `32768` | Per LLM call. Headroom for large tool-call inputs (e.g. file writes via heredoc); Sonnet 4 / Opus 4 cap at 64K. |
|
||||
| `BUZZ_AGENT_MAX_CONTEXT_TOKENS` | `200000` | Provider context window used by the handoff gate. |
|
||||
| `BUZZ_AGENT_MAX_HANDOFFS` | `10` | Max context handoffs per session before falling back to truncation. |
|
||||
| `BUZZ_AGENT_LLM_TIMEOUT_SECS` | `240` | Max seconds with no response bytes before abandoning an LLM call (per-read inactivity, not wall-clock). |
|
||||
| `BUZZ_AGENT_TOOL_TIMEOUT_SECS` | `660` | Per-tool call timeout in seconds |
|
||||
| `BUZZ_AGENT_MAX_PARALLEL_TOOLS` | `8` | Max concurrent tool calls per turn (1 = sequential) |
|
||||
| `BUZZ_AGENT_MAX_SESSIONS` | unlimited | Max concurrent ACP sessions. Sessions are cheap; default has no cap. |
|
||||
| `BUZZ_AGENT_MAX_LINE_BYTES` | `4194304` | 4 MiB. Hard cap on inbound JSON-RPC frames. |
|
||||
| `BUZZ_AGENT_MAX_HISTORY_BYTES` | `1048576` | 1 MiB. Old turns are evicted past this. |
|
||||
| `BUZZ_AGENT_MAX_TOOL_RESULT_TEXT_BYTES` | `51200` | 50 KiB. Per-result cap on tool-output text; oversize is middle-elided (head + tail kept) with an inline marker. Images are exempt. |
|
||||
| `BUZZ_AGENT_REQUIRE_REPLY` | `0` (`1` on mesh) | `1` enables the [reply guard](#reply-guard) — remind the model to publish when a turn is about to end with nothing posted to Buzz. Desktop defaults it to `1` for Buzz shared-compute agents. |
|
||||
|
||||
|
||||
## Reply Guard
|
||||
|
||||
Off by default, except on Buzz shared-compute (mesh) agents, where Buzz Desktop
|
||||
sets `BUZZ_AGENT_REQUIRE_REPLY=1` automatically. With it enabled, a turn that is
|
||||
about to end without any recognized attempt to post to Buzz gets a reminder that
|
||||
its assistant text is invisible to humans, and is rerolled.
|
||||
|
||||
This exists because a Buzz agent's reasoning and tool output are not shown to
|
||||
anyone. A turn that does real work and never posts is a silent failure — the
|
||||
requester waits on a result that was produced and thrown away.
|
||||
|
||||
Mesh agents get it by default because they run on small local models, which are
|
||||
the ones most likely to do the work and then end the turn without publishing it.
|
||||
Setting `BUZZ_AGENT_REQUIRE_REPLY=0` on the agent, persona, or global env opts a
|
||||
mesh agent back out; the default never overrides an explicit value.
|
||||
|
||||
**Advisory, never a trap.** At most two reminders, then the turn ends whether or
|
||||
not anything was published. The guard catches accidental omission; it does not
|
||||
compel speech. The reminder text explicitly licenses silence, because the
|
||||
built-in system prompt says publishing is optional and silence is often the
|
||||
correct outcome.
|
||||
|
||||
**Recognition contract.** A turn counts as having replied when it issues a call
|
||||
that:
|
||||
|
||||
- resolves to a registered, non-hook tool (a hallucinated tool name is rejected
|
||||
at preflight and never runs, so it must not disarm the guard),
|
||||
- whose qualified name ends in `__shell` — i.e. the bare tool name is exactly
|
||||
`shell`, which is `buzz-dev-mcp`'s shell tool and any other server's, and
|
||||
- whose `command` argument contains `messages send` or `reactions add`.
|
||||
|
||||
`messages send` also covers `messages send-diff`. Reactions count because the
|
||||
built-in prompt directs agents to react rather than post a bare
|
||||
acknowledgement, so nagging an agent that reacted would punish documented
|
||||
behavior.
|
||||
|
||||
Detection is checked **after** the per-turn tool-call cap
|
||||
(`MAX_TOOL_CALLS_PER_TURN`) is applied: a publish-shaped call that was discarded
|
||||
never ran.
|
||||
|
||||
**It recognizes an attempt, not a successful publish.** Only the command text is
|
||||
inspected, never the exit status. A send that fails still satisfies the guard —
|
||||
which is fine, since a failed send already returns a non-zero exit and error
|
||||
JSON to the model, louder feedback than a reminder.
|
||||
|
||||
**Known limits**, both deliberate. A command assembled at runtime (`$CMD`) or
|
||||
buried in a wrapper script is missed, so that turn is reminded despite having
|
||||
posted. Text that merely quotes a send (`echo "buzz messages send"`) matches, so
|
||||
that turn is not reminded. Missing a real post is the expensive direction, and
|
||||
substring matching is the forgiving one there. Neither edge is pinned by a test;
|
||||
the matcher is free to improve.
|
||||
|
||||
**Budget.** Reminders ride the existing `_Stop` gate and share
|
||||
`BUZZ_AGENT_STOP_MAX_REJECTIONS` — the outer cap on every end-turn objection.
|
||||
At the default 3 both reminders fit; at 1 only one does; at 0 the guard is off
|
||||
along with the hooks. A round carrying both a `_Stop` hook objection and a
|
||||
reminder costs one rejection and delivers both texts. This is not a new
|
||||
lifecycle hook — see [MCP_DRIVEN_HOOKS.md](../../docs/MCP_DRIVEN_HOOKS.md).
|
||||
|
||||
|
||||
## Providers
|
||||
|
||||
`buzz-agent` speaks a few HTTP dialects. Pick with `BUZZ_AGENT_PROVIDER`.
|
||||
|
||||
| Provider | `BUZZ_AGENT_PROVIDER` | Endpoint (auto) | Tested with |
|
||||
|---|---|---|---|
|
||||
| Anthropic | `anthropic` | `POST {base}/v1/messages` | claude-sonnet-4-5, claude-opus-4 |
|
||||
| OpenAI | `openai` | `POST {base}/responses` | gpt-5, gpt-5-mini, o4-mini, gpt-4o |
|
||||
| vLLM | `openai` | `POST {base}/chat/completions` | any tool-calling model |
|
||||
| llama.cpp | `openai` | `POST {base}/chat/completions` | any tool-calling GGUF |
|
||||
| Ollama | `openai` | `POST {base}/chat/completions` | llama3.1, qwen2.5-coder |
|
||||
| Block Gateway | `openai` | `POST {base}/chat/completions` | gpt-5, claude |
|
||||
| OpenRouter | `openrouter` | `POST {base}/chat/completions` | anything they route (extended-thinking replay, provider-agnostic tool calling) |
|
||||
| Databricks | `databricks` | `POST {host}/serving-endpoints/{model}/invocations` | goose-claude-4-6-sonnet |
|
||||
| Databricks AI Gateway v2 | `databricks_v2` | `POST {host}/ai-gateway/{provider}/v1/...` | databricks-gpt-5-5, databricks-claude-opus-4-7 |
|
||||
|
||||
If `BUZZ_AGENT_PROVIDER=anthropic` is selected without `ANTHROPIC_API_KEY`, `BUZZ_AGENT_PROVIDER=openai` is selected without `OPENAI_COMPAT_API_KEY`, or `BUZZ_AGENT_PROVIDER=openrouter` is selected without `OPENROUTER_API_KEY`, the agent returns an error — there is no implicit fallback to another provider.
|
||||
|
||||
`provider=openai` speaks two HTTP dialects: the [Responses API](https://platform.openai.com/docs/api-reference/responses) (`/v1/responses`, required for GPT-5 / o-series tool-calling on OpenAI's own service) and the [Chat Completions API](https://platform.openai.com/docs/api-reference/chat) (`/chat/completions`, the broadly-supported OpenAI-compatible wire format).
|
||||
|
||||
By default (`OPENAI_COMPAT_API=auto`) the agent picks **Responses** when `OPENAI_COMPAT_BASE_URL` points at an `*.openai.com` host and **Chat Completions** everywhere else. Pin the choice explicitly with `OPENAI_COMPAT_API=chat` or `OPENAI_COMPAT_API=responses` for providers that diverge from the default (e.g. a Responses-compatible self-hosted gateway).
|
||||
|
||||
`provider=openrouter` is first-class, not routed through `provider=openai`: it speaks OpenAI's Chat Completions wire format but with OpenRouter-specific extensions layered on top —
|
||||
|
||||
- `reasoning.effort` is set on the request when reasoning effort is configured. The request deliberately carries no `provider.require_parameters` filter: that filter routes only to endpoints advertising every parameter in the body, and 83 of 274 tools-capable OpenRouter models do not advertise `reasoning`, so it turns an effort setting into a hard 404 on a valid model id. A model that cannot reason answers without reasoning instead.
|
||||
- The response's `reasoning_details` array (opaque extended-thinking payload) is captured and replayed byte-for-byte on the next turn's assistant message, so multi-turn tool use keeps the model's chain-of-thought.
|
||||
- `anthropic/*` models get Anthropic-style `cache_control` breakpoints injected on the system message and the last two user messages.
|
||||
- Retryable statuses (429 and typed `provider_overloaded` 503) honor the documented `Retry-After` header (clamped to a small ceiling — see `RETRY_AFTER_CAP_SECS` in `llm.rs` — since the sleep happens outside `BUZZ_AGENT_LLM_TIMEOUT_SECS`); 502 and untyped 503 retry with jittered backoff instead. `401` is treated as an expired/invalid key and refreshed once, while `402` (no credits) and `403` (guardrail/moderation/permission) fail immediately without retry.
|
||||
|
||||
`Provider` is a Rust `enum` with one `match` in `Llm::complete`. There is no trait, no `Box<dyn>`, no async-trait. Adding a provider is a `match` arm and one `body`/`parse` pair in `llm.rs`.
|
||||
|
||||
## MCP Servers
|
||||
|
||||
The client passes MCP server specs in `session/new`. The agent spawns each one as a stdio subprocess, calls `tools/list`, and merges everything into a single tool catalog the LLM sees. Tool names are namespaced as `server__tool` (double underscore separator). Bare tool names containing `__` are rejected at registration.
|
||||
|
||||
Example: a single echo MCP server.
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 2,
|
||||
"method": "session/new",
|
||||
"params": {
|
||||
"cwd": "/work",
|
||||
"mcpServers": [
|
||||
{
|
||||
"name": "echo",
|
||||
"command": "/usr/local/bin/echo-mcp",
|
||||
"args": ["--mode", "stdio"],
|
||||
"env": [
|
||||
{ "name": "ECHO_VERBOSE", "value": "1" }
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Multiple servers: just add more entries. Tool calls fan out to the right server by namespace prefix.
|
||||
|
||||
**Transport: stdio only.** No HTTP, no SSE. We advertise this in `agentCapabilities` (`mcpCapabilities.http: false`, `mcpCapabilities.sse: false`); spec-compliant clients won't ask for what we don't have.
|
||||
|
||||
## Security Model
|
||||
|
||||
The trust boundary is **the operator who launched the agent**. The harness, MCP server binaries, and API keys are all trusted. Untrusted input — model output, tool results, prompts — is bounded.
|
||||
|
||||
| Boundary | Mechanism |
|
||||
|---|---|
|
||||
| Stdout discipline | Single-consumer `mpsc` channel feeding stdout. No two tasks can interleave bytes. All logs go to stderr. |
|
||||
| MCP child env | Whitelist (`PATH`, `HOME`, `TERM`, `LANG`, `LC_ALL`, `TMPDIR`) plus what the client explicitly passes. Your `ANTHROPIC_API_KEY` does not leak into MCP children. |
|
||||
| MCP child lifetime | Process group via `setpgid(0,0)` in `pre_exec`. On transport break or shutdown: `killpg(SIGKILL)`. Grandchildren die too. |
|
||||
| Server poisoning | After a timeout or transport break, the offending server is marked dead. Future calls trigger a lazy restart with exponential backoff. Other servers keep working. |
|
||||
| Frame size | `BUZZ_AGENT_MAX_LINE_BYTES` (default 4 MiB). Oversize → connection killed. |
|
||||
| LLM response size | 16 MiB hard cap. Both `Content-Length` precheck and streaming-buffer cap. |
|
||||
| Cancellation | `tokio::select! { biased; _ = cancel.changed() => ... }` at every loop boundary. Cancel always wins the race. |
|
||||
| Session isolation | Unlimited concurrent sessions by default (configurable via `BUZZ_AGENT_MAX_SESSIONS`). One prompt per session at a time. Each session gets its own MCP servers. |
|
||||
| `tool_use ↔ tool_result` pairing | Encoded in the type system. Every `ToolCall` and `ToolResult` carries a `provider_id: String` (not `Option`). |
|
||||
|
||||
### Bounded Everything
|
||||
|
||||
| Limit | Default | Where |
|
||||
|---|---|---|
|
||||
| Inbound JSON-RPC frame | 4 MiB | `BUZZ_AGENT_MAX_LINE_BYTES` |
|
||||
| Single prompt | 1 MiB | `MAX_PROMPT_BYTES` |
|
||||
| History window | 1 MiB | `BUZZ_AGENT_MAX_HISTORY_BYTES` |
|
||||
| LLM response body | 16 MiB | `MAX_LLM_RESPONSE_BYTES` |
|
||||
| LLM error body | 4 KiB | `MAX_LLM_ERROR_BODY_BYTES` |
|
||||
| Tool result body (total, incl. images) | 8 MiB | `MAX_TOOL_RESULT_BYTES` |
|
||||
| Tool result text | 50 KiB | `BUZZ_AGENT_MAX_TOOL_RESULT_TEXT_BYTES` |
|
||||
| MCP servers / session | 16 | `MAX_MCP_SERVERS` |
|
||||
| Tools / session | 128 | `MAX_TOOLS_PER_SESSION` |
|
||||
| Tool description bytes | 1 KiB | `MAX_DESCRIPTION_BYTES` |
|
||||
| Tool schema bytes | 4 KiB | `MAX_SCHEMA_BYTES` (oversize → replaced with `{}`) |
|
||||
| Tool calls per turn | 64 | `MAX_TOOL_CALLS_PER_TURN` |
|
||||
| Loop rounds | 0 (unlimited) | `BUZZ_AGENT_MAX_ROUNDS` |
|
||||
| LLM read inactivity timeout | 240 s | `BUZZ_AGENT_LLM_TIMEOUT_SECS` |
|
||||
| Tool call timeout | 660 s | `BUZZ_AGENT_TOOL_TIMEOUT_SECS` |
|
||||
|
||||
## What This Is NOT
|
||||
|
||||
A short list, because the answer is mostly "no":
|
||||
|
||||
- **Not a framework.** No plugins, no recipes, no slash commands, no modes. MCP servers can participate in agent lifecycle via [hook tools](../../docs/MCP_DRIVEN_HOOKS.md) (`_Stop`, `_PostCompact`), but these are advisory, fail-open, and budget-bounded — not a plugin system.
|
||||
- **Not streaming.** One non-streaming HTTP POST per round. The LLM's generated text is forwarded to the client as `agent_message_chunk`, but there is no token-level streaming.
|
||||
- **Not persistent.** Everything is in-memory, per-process. No SQLite. When context fills, the agent summarizes its own history and continues (context handoff). No external persistence.
|
||||
- **Not an SDK.** This is a binary. The protocol seam is stdin/stdout. Use it from any language.
|
||||
- **Not a UI.** No TUI, no web, no notifications. The client renders.
|
||||
- **Not authenticated.** API keys come from env. Use systemd, Docker secrets, or a wrapper.
|
||||
- **Not networked MCP.** Stdio transport only. No HTTP/SSE MCP transport.
|
||||
- **Not load-able.** No `session/load`. We advertise `loadSession: false`.
|
||||
- **Not a router.** No agent-to-agent, no fan-out, no orchestration. One model. One loop.
|
||||
|
||||
**Concurrency model:**
|
||||
|
||||
```
|
||||
┌──── reader task ──────────┐
|
||||
│ (stdin → JSON-RPC → ...) │
|
||||
│ │
|
||||
stdin ─────────┤ dispatch │
|
||||
│ │ │
|
||||
│ ├── initialize │ (sync reply)
|
||||
│ ├── session/new │ (sync reply)
|
||||
│ ├── session/prompt ───┼─── spawn ──> prompt task
|
||||
│ │ │ │
|
||||
│ ├── session/cancel ───┼─> watch::send│ (biased select wins)
|
||||
│ │ │ │
|
||||
└───────────────────────────┘ │
|
||||
│
|
||||
┌── writer task ────────────────┐ │
|
||||
stdout ────────┤ mpsc<WireMsg> consumer │<─────────┘
|
||||
│ (the only stdout writer) │
|
||||
└───────────────────────────────┘
|
||||
```
|
||||
|
||||
One reader, one writer, up to 8 concurrent prompt tasks (one per session).
|
||||
|
||||
## Building
|
||||
|
||||
```bash
|
||||
cargo build --release -p buzz-agent
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
cargo test -p buzz-agent
|
||||
```
|
||||
|
||||
Test strategy is **real subprocess, no mocks**:
|
||||
|
||||
- **Fake LLM** — `tests/fake_llm.rs` and the helpers in `tests/regressions.rs` spin up a real `tokio::net::TcpListener` on port 0, parse `Content-Length`, and return scripted JSON. No HTTP mocking library.
|
||||
- **Fake MCP server** — `tests/bin/fake_mcp.rs` is a separate binary controlled by env vars: `FAKE_MCP_HANG_INIT`, `FAKE_MCP_TOOL_DELAY`, `FAKE_MCP_SPAWN_GRANDCHILD`, etc. Each fault path is a real process being abused.
|
||||
- **Regression tests are the changelog.** Each `#[test]` in `regressions.rs` is named for the bug it locks down: `assistant_text_preserved_across_prompts`, `cancel_leaves_history_valid_for_next_prompt`, `mcp_init_timeout_kills_child`, `oversize_line_kills_connection`. Read them in order to learn the protocol's failure modes.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.6 MiB |
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,847 @@
|
||||
//! Token sources for the LLM transport layer.
|
||||
//!
|
||||
//! [`TokenSource`] decouples request auth from `Config::api_key`: providers
|
||||
//! can supply a static string ([`StaticTokenSource`]) or a refreshable OAuth
|
||||
//! 2.0 PKCE engine ([`PkceOAuthTokenSource`]). Engines own their own cache
|
||||
//! and refresh logic; the [`Llm`] just asks for a bearer per request.
|
||||
//!
|
||||
//! The PKCE engine implements RFC 6749 + RFC 7636 with on-disk token
|
||||
//! caching keyed by `sha256(discovery_url|client_id|scopes)`. It's the
|
||||
//! same shape goose uses for Databricks, but we own the wire format and
|
||||
//! cache directory so the two are independently upgradable.
|
||||
//!
|
||||
//! First-use (cache empty) requires a browser: the engine opens
|
||||
//! `authorization_endpoint` in `webbrowser`, listens on `127.0.0.1:0`,
|
||||
//! captures the redirect, and exchanges the code for a token. Subsequent
|
||||
//! calls hit the cache and silently refresh when expired.
|
||||
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::Engine;
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sha2::Digest;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::types::AgentError;
|
||||
|
||||
/// Buffer before `expires_at` to consider a cached token "still good".
|
||||
/// Keeps us off the cliff if the clock or the server's clock drifts.
|
||||
const TOKEN_REFRESH_LEEWAY: Duration = Duration::from_secs(60);
|
||||
|
||||
/// Wall-clock budget for the interactive browser dance. Goose uses 60s.
|
||||
/// We match: any longer and the user has gone to lunch.
|
||||
const BROWSER_AUTH_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
|
||||
/// Asynchronous source of a bearer token. The [`Llm`] calls this per
|
||||
/// request, so impls are expected to be cheap on the cache-hit path.
|
||||
#[async_trait]
|
||||
pub trait TokenSource: Send + Sync {
|
||||
async fn bearer(&self) -> Result<String, AgentError>;
|
||||
|
||||
/// Return a bearer token from cache or refresh, **never** opening a browser.
|
||||
///
|
||||
/// The default delegates to [`bearer`](Self::bearer) — correct for token
|
||||
/// sources (e.g. static API keys) that can never trigger a browser flow.
|
||||
/// [`PkceOAuthTokenSource`] overrides this to stop before the browser step.
|
||||
async fn bearer_no_browser(&self) -> Result<String, AgentError> {
|
||||
self.bearer().await
|
||||
}
|
||||
|
||||
/// Force a fresh bearer after the server rejected the current one (401).
|
||||
///
|
||||
/// `rejected` is the exact access token that just got the 401. Unlike
|
||||
/// [`bearer`](Self::bearer), which trusts the local expiry clock, this is
|
||||
/// driven by the server's verdict: the cached token looked valid to us
|
||||
/// (well within its local expiry) but the provider rejected it — clock
|
||||
/// skew, server-side revocation, or a node that never saw it. The clock
|
||||
/// therefore can't decide whether to refresh; the caller passes the
|
||||
/// rejected token so the impl can refresh unless a concurrent caller has
|
||||
/// *already* replaced it. Implementations must obtain a new token without
|
||||
/// any interactive step, so a headless harness never hangs. The default
|
||||
/// returns the existing bearer — correct for sources that can't refresh
|
||||
/// (a static key); the caller's retry then fails terminally rather than
|
||||
/// looping.
|
||||
async fn refresh_now(&self, _rejected: &str) -> Result<String, AgentError> {
|
||||
self.bearer().await
|
||||
}
|
||||
}
|
||||
|
||||
/// A token that never changes for the life of the process.
|
||||
pub struct StaticTokenSource(String);
|
||||
|
||||
impl StaticTokenSource {
|
||||
pub fn new(token: impl Into<String>) -> Self {
|
||||
Self(token.into())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenSource for StaticTokenSource {
|
||||
async fn bearer(&self) -> Result<String, AgentError> {
|
||||
Ok(self.0.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Static config for an OAuth 2.0 Authorization Code + PKCE provider.
|
||||
///
|
||||
/// The `discovery_url` must return a JSON document with at least
|
||||
/// `authorization_endpoint` and `token_endpoint` (RFC 8414). The
|
||||
/// `cache_namespace` is the directory under `~/.config/buzz-agent/oauth/`
|
||||
/// the token JSON lives in — separates providers' caches cleanly.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PkceOAuthConfig {
|
||||
pub discovery_url: String,
|
||||
pub client_id: String,
|
||||
pub scopes: Vec<String>,
|
||||
pub cache_namespace: String,
|
||||
/// When `Some`, the engine writes tokens here instead of
|
||||
/// `~/.config/buzz-agent/oauth/<cache_namespace>/`. Production code
|
||||
/// leaves this `None`. Integration tests use it to avoid stomping on
|
||||
/// a shared `$HOME` when running in parallel.
|
||||
pub cache_dir_override: Option<PathBuf>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
struct CachedToken {
|
||||
access_token: String,
|
||||
refresh_token: Option<String>,
|
||||
/// Unix seconds. `None` means the server didn't advertise an expiry;
|
||||
/// we use it without checking and rely on refresh on 401.
|
||||
expires_at: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct OidcEndpoints {
|
||||
authorization_endpoint: String,
|
||||
token_endpoint: String,
|
||||
}
|
||||
|
||||
/// PKCE OAuth token source with on-disk refresh cache.
|
||||
///
|
||||
/// First call:
|
||||
/// 1. Loads from cache if present and unexpired.
|
||||
/// 2. Otherwise tries `refresh_token` if cached.
|
||||
/// 3. Otherwise runs the full browser flow.
|
||||
///
|
||||
/// Subsequent calls hit an in-memory copy of the cached token and only
|
||||
/// touch disk/network if the access token is past `expires_at`.
|
||||
pub struct PkceOAuthTokenSource {
|
||||
cfg: PkceOAuthConfig,
|
||||
http: Client,
|
||||
cache_path: PathBuf,
|
||||
/// Single-flight guard: only one refresh/browser flow at a time, even
|
||||
/// if many tool calls land concurrently.
|
||||
state: Mutex<Option<CachedToken>>,
|
||||
}
|
||||
|
||||
impl PkceOAuthTokenSource {
|
||||
pub fn new(cfg: PkceOAuthConfig) -> Result<Arc<Self>, AgentError> {
|
||||
let cache_path = cache_path_for(&cfg)?;
|
||||
if let Some(parent) = cache_path.parent() {
|
||||
fs::create_dir_all(parent)
|
||||
.map_err(|e| AgentError::Llm(format!("oauth cache dir {parent:?}: {e}")))?;
|
||||
}
|
||||
let initial = read_cache(&cache_path);
|
||||
Ok(Arc::new(Self {
|
||||
cfg,
|
||||
http: Client::new(),
|
||||
cache_path,
|
||||
state: Mutex::new(initial),
|
||||
}))
|
||||
}
|
||||
|
||||
/// Discover authorization + token endpoints from the well-known URL.
|
||||
async fn endpoints(&self) -> Result<OidcEndpoints, AgentError> {
|
||||
let v: Value = self
|
||||
.http
|
||||
.get(&self.cfg.discovery_url)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth discovery: {e}")))?
|
||||
.error_for_status()
|
||||
.map_err(|e| AgentError::Llm(format!("oauth discovery status: {e}")))?
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth discovery json: {e}")))?;
|
||||
let auth = v
|
||||
.get("authorization_endpoint")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| {
|
||||
AgentError::Llm("oauth discovery: authorization_endpoint missing".into())
|
||||
})?
|
||||
.to_string();
|
||||
let token = v
|
||||
.get("token_endpoint")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| AgentError::Llm("oauth discovery: token_endpoint missing".into()))?
|
||||
.to_string();
|
||||
Ok(OidcEndpoints {
|
||||
authorization_endpoint: auth,
|
||||
token_endpoint: token,
|
||||
})
|
||||
}
|
||||
|
||||
/// Persist a token to disk and the in-memory cell.
|
||||
fn save(&self, state: &mut Option<CachedToken>, token: CachedToken) -> Result<(), AgentError> {
|
||||
let body = serde_json::to_vec_pretty(&token)
|
||||
.map_err(|e| AgentError::Llm(format!("oauth cache serialize: {e}")))?;
|
||||
// Atomic rename so a concurrent reader never sees a partial write.
|
||||
let tmp = self.cache_path.with_extension("json.tmp");
|
||||
fs::write(&tmp, &body)
|
||||
.map_err(|e| AgentError::Llm(format!("oauth cache write {tmp:?}: {e}")))?;
|
||||
fs::rename(&tmp, &self.cache_path)
|
||||
.map_err(|e| AgentError::Llm(format!("oauth cache rename: {e}")))?;
|
||||
*state = Some(token);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Exchange a refresh token for a fresh access token.
|
||||
async fn refresh(
|
||||
&self,
|
||||
endpoints: &OidcEndpoints,
|
||||
refresh_token: &str,
|
||||
) -> Result<CachedToken, AgentError> {
|
||||
let params = [
|
||||
("grant_type", "refresh_token"),
|
||||
("refresh_token", refresh_token),
|
||||
("client_id", &self.cfg.client_id),
|
||||
];
|
||||
let resp = self
|
||||
.http
|
||||
.post(&endpoints.token_endpoint)
|
||||
.form(¶ms)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth refresh: {e}")))?;
|
||||
if !resp.status().is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(AgentError::Llm(format!("oauth refresh failed: {body}")));
|
||||
}
|
||||
let v: Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth refresh json: {e}")))?;
|
||||
token_from_response(&v, Some(refresh_token))
|
||||
}
|
||||
|
||||
/// Run the full browser-mediated Authorization Code + PKCE flow.
|
||||
/// Caller must hold a TTY/browser: this opens a window and blocks.
|
||||
pub async fn interactive_login(&self) -> Result<(), AgentError> {
|
||||
let endpoints = self.endpoints().await?;
|
||||
let token = browser_pkce_flow(&self.http, &self.cfg, &endpoints).await?;
|
||||
let mut state = self.state.lock().await;
|
||||
self.save(&mut state, token)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenSource for PkceOAuthTokenSource {
|
||||
async fn bearer(&self) -> Result<String, AgentError> {
|
||||
let mut state = self.state.lock().await;
|
||||
|
||||
// 1. In-memory cache hit, still fresh.
|
||||
if let Some(tok) = state.as_ref() {
|
||||
if !is_expired(tok) {
|
||||
return Ok(tok.access_token.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Re-read disk — another process may have refreshed already.
|
||||
if let Some(disk_tok) = read_cache(&self.cache_path) {
|
||||
if !is_expired(&disk_tok) {
|
||||
let bearer = disk_tok.access_token.clone();
|
||||
*state = Some(disk_tok);
|
||||
return Ok(bearer);
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Try refresh if we have a refresh token. Discover endpoints once
|
||||
// here — deliberately hoisted above the refresh-token check so the
|
||||
// browser flow at step 5 (which also needs them) reuses this call.
|
||||
let endpoints = self.endpoints().await?;
|
||||
let refresh = state.as_ref().and_then(|t| t.refresh_token.clone());
|
||||
if let Some(rt) = refresh {
|
||||
match self.refresh(&endpoints, &rt).await {
|
||||
Ok(fresh) => {
|
||||
let bearer = fresh.access_token.clone();
|
||||
self.save(&mut state, fresh)?;
|
||||
return Ok(bearer);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "oauth refresh failed; falling back to browser flow");
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Re-read disk after refresh failure — another process may have won the race.
|
||||
if let Some(disk_tok) = read_cache(&self.cache_path) {
|
||||
if !is_expired(&disk_tok) {
|
||||
let bearer = disk_tok.access_token.clone();
|
||||
*state = Some(disk_tok);
|
||||
return Ok(bearer);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 5. No usable cache: full browser dance.
|
||||
let fresh = browser_pkce_flow(&self.http, &self.cfg, &endpoints).await?;
|
||||
let bearer = fresh.access_token.clone();
|
||||
self.save(&mut state, fresh)?;
|
||||
Ok(bearer)
|
||||
}
|
||||
|
||||
async fn bearer_no_browser(&self) -> Result<String, AgentError> {
|
||||
self.try_bearer_no_browser().await
|
||||
}
|
||||
|
||||
/// Force-refresh after a 401, never touching the browser flow.
|
||||
///
|
||||
/// `rejected` is the access token the server just 401'd. Coalescing keys
|
||||
/// off token *identity*, not the expiry clock: a 401 means the token was
|
||||
/// rejected while it still looked locally fresh, so `is_expired()` would
|
||||
/// say "keep it" and no grant would ever run. Instead, under the lock we
|
||||
/// compare the current cached token to `rejected` — if they differ, a
|
||||
/// concurrent caller (this process or a sibling) already refreshed, so we
|
||||
/// return the new token without burning a second grant. If they still
|
||||
/// match, this is the rejected token and we run the refresh-token grant
|
||||
/// unconditionally. The whole check→refresh→save runs under one lock hold
|
||||
/// so concurrent callers serialize. On any failure the refresh token is
|
||||
/// preserved (never nulled) and the error is terminal `LlmAuth` — no
|
||||
/// browser, no hang.
|
||||
async fn refresh_now(&self, rejected: &str) -> Result<String, AgentError> {
|
||||
let mut state = self.state.lock().await;
|
||||
|
||||
// 1. Coalesce by identity: if the cached token (in-memory, then disk)
|
||||
// is no longer the one the server rejected, someone already
|
||||
// refreshed it. Return that instead of grabbing another grant.
|
||||
if let Some(tok) = state.as_ref() {
|
||||
if tok.access_token != rejected {
|
||||
return Ok(tok.access_token.clone());
|
||||
}
|
||||
}
|
||||
if let Some(disk_tok) = read_cache(&self.cache_path) {
|
||||
if disk_tok.access_token != rejected {
|
||||
let bearer = disk_tok.access_token.clone();
|
||||
*state = Some(disk_tok);
|
||||
return Ok(bearer);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. The cached token is still the rejected one. Run the refresh-token
|
||||
// grant unconditionally — the expiry clock can't be trusted here, a
|
||||
// locally-fresh token is exactly what got 401'd.
|
||||
let refresh = state.as_ref().and_then(|t| t.refresh_token.clone());
|
||||
let Some(rt) = refresh else {
|
||||
return Err(AgentError::LlmAuth(
|
||||
"token rejected and no refresh token available".into(),
|
||||
));
|
||||
};
|
||||
let endpoints = self.endpoints().await?;
|
||||
match self.refresh(&endpoints, &rt).await {
|
||||
Ok(fresh) => {
|
||||
let bearer = fresh.access_token.clone();
|
||||
self.save(&mut state, fresh)?;
|
||||
Ok(bearer)
|
||||
}
|
||||
// 3. Refresh token is itself dead. Terminal — surfacing LlmAuth
|
||||
// stops the retry loop instead of falling to the browser flow,
|
||||
// which would hang a headless harness.
|
||||
Err(e) => Err(AgentError::LlmAuth(format!("token refresh failed: {e}"))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PkceOAuthTokenSource {
|
||||
/// Return a bearer token from cache or refresh, **never** opening a browser.
|
||||
///
|
||||
/// Follows the same steps as [`bearer`](TokenSource::bearer) but stops at
|
||||
/// step 4 — if no usable token is available after cache + refresh attempts,
|
||||
/// returns `Err(LlmAuth(...))` instead of launching the browser PKCE flow.
|
||||
/// Used by model-discovery paths that must not block on user interaction.
|
||||
pub(crate) async fn try_bearer_no_browser(&self) -> Result<String, AgentError> {
|
||||
let mut state = self.state.lock().await;
|
||||
|
||||
// 1. In-memory cache hit, still fresh.
|
||||
if let Some(tok) = state.as_ref() {
|
||||
if !is_expired(tok) {
|
||||
return Ok(tok.access_token.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Re-read disk — another process may have refreshed already.
|
||||
if let Some(disk_tok) = read_cache(&self.cache_path) {
|
||||
if !is_expired(&disk_tok) {
|
||||
let bearer = disk_tok.access_token.clone();
|
||||
*state = Some(disk_tok);
|
||||
return Ok(bearer);
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Try refresh if we have a refresh token. Endpoints are discovered
|
||||
// lazily here — only when a refresh token is actually present — so
|
||||
// that an unreachable OIDC discovery URL cannot prevent the
|
||||
// no-token/no-cache path from returning `LlmAuth` (graceful
|
||||
// fallback) instead of `Llm` (hard error).
|
||||
let refresh = state.as_ref().and_then(|t| t.refresh_token.clone());
|
||||
if let Some(rt) = refresh {
|
||||
let endpoints = self.endpoints().await?;
|
||||
match self.refresh(&endpoints, &rt).await {
|
||||
Ok(fresh) => {
|
||||
let bearer = fresh.access_token.clone();
|
||||
self.save(&mut state, fresh)?;
|
||||
return Ok(bearer);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "oauth refresh failed during model discovery");
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Re-read disk after refresh failure.
|
||||
if let Some(disk_tok) = read_cache(&self.cache_path) {
|
||||
if !is_expired(&disk_tok) {
|
||||
let bearer = disk_tok.access_token.clone();
|
||||
*state = Some(disk_tok);
|
||||
return Ok(bearer);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// No usable token — return error instead of opening a browser.
|
||||
Err(AgentError::LlmAuth(
|
||||
"no cached Databricks token; run `buzz-agent auth databricks` first".into(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
// ---- helpers -------------------------------------------------------------
|
||||
|
||||
/// Aborts a spawned task when dropped. Used to guarantee the localhost
|
||||
/// callback server doesn't outlive a failed/abandoned PKCE attempt.
|
||||
struct AbortOnDrop(tokio::task::JoinHandle<()>);
|
||||
|
||||
impl Drop for AbortOnDrop {
|
||||
fn drop(&mut self) {
|
||||
self.0.abort();
|
||||
}
|
||||
}
|
||||
|
||||
fn is_expired(t: &CachedToken) -> bool {
|
||||
let Some(exp) = t.expires_at else {
|
||||
return false;
|
||||
};
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
now + TOKEN_REFRESH_LEEWAY.as_secs() >= exp
|
||||
}
|
||||
|
||||
fn cache_path_for(cfg: &PkceOAuthConfig) -> Result<PathBuf, AgentError> {
|
||||
let mut h = sha2::Sha256::new();
|
||||
h.update(cfg.discovery_url.as_bytes());
|
||||
h.update(b"|");
|
||||
h.update(cfg.client_id.as_bytes());
|
||||
h.update(b"|");
|
||||
h.update(cfg.scopes.join(",").as_bytes());
|
||||
let hash = hex::encode(h.finalize());
|
||||
|
||||
let dir = match &cfg.cache_dir_override {
|
||||
Some(p) => p.join(&cfg.cache_namespace),
|
||||
None => dirs::home_dir()
|
||||
.ok_or_else(|| AgentError::Llm("oauth cache: home directory not found".into()))?
|
||||
.join(".config")
|
||||
.join("buzz-agent")
|
||||
.join("oauth")
|
||||
.join(&cfg.cache_namespace),
|
||||
};
|
||||
Ok(dir.join(format!("{hash}.json")))
|
||||
}
|
||||
|
||||
fn read_cache(path: &PathBuf) -> Option<CachedToken> {
|
||||
let body = fs::read(path).ok()?;
|
||||
serde_json::from_slice(&body).ok()
|
||||
}
|
||||
|
||||
/// Parse a token-endpoint JSON response. Fails loudly when `access_token`
|
||||
/// is missing or empty — without this, a malformed server response would
|
||||
/// be cached and `bearer()` would silently return `""` until the entry
|
||||
/// expires or is deleted by hand.
|
||||
fn token_from_response(
|
||||
v: &Value,
|
||||
fallback_refresh: Option<&str>,
|
||||
) -> Result<CachedToken, AgentError> {
|
||||
let access_token = v
|
||||
.get("access_token")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|s| !s.is_empty())
|
||||
.ok_or_else(|| AgentError::Llm("oauth: token response missing/empty access_token".into()))?
|
||||
.to_string();
|
||||
let refresh_token = v
|
||||
.get("refresh_token")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
.or_else(|| fallback_refresh.map(str::to_string));
|
||||
let expires_at = v.get("expires_in").and_then(Value::as_u64).map(|secs| {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
+ secs
|
||||
});
|
||||
Ok(CachedToken {
|
||||
access_token,
|
||||
refresh_token,
|
||||
expires_at,
|
||||
})
|
||||
}
|
||||
|
||||
/// PKCE pieces: URL-safe random verifier (~64 chars) and its SHA-256
|
||||
/// challenge (RFC 7636 §4.2).
|
||||
fn pkce_pair() -> Result<(String, String), AgentError> {
|
||||
let mut bytes = [0u8; 48];
|
||||
getrandom::fill(&mut bytes).map_err(|e| AgentError::Llm(format!("pkce rng: {e}")))?;
|
||||
let verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
|
||||
let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.encode(sha2::Sha256::digest(verifier.as_bytes()));
|
||||
Ok((verifier, challenge))
|
||||
}
|
||||
|
||||
fn random_state() -> Result<String, AgentError> {
|
||||
let mut bytes = [0u8; 16];
|
||||
getrandom::fill(&mut bytes).map_err(|e| AgentError::Llm(format!("state rng: {e}")))?;
|
||||
Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes))
|
||||
}
|
||||
|
||||
/// Spin up a localhost callback server, open the authorize URL in a
|
||||
/// browser, wait up to [`BROWSER_AUTH_TIMEOUT`] for the redirect, then
|
||||
/// exchange the code for a token.
|
||||
async fn browser_pkce_flow(
|
||||
http: &Client,
|
||||
cfg: &PkceOAuthConfig,
|
||||
endpoints: &OidcEndpoints,
|
||||
) -> Result<CachedToken, AgentError> {
|
||||
use axum::{extract::Query, response::Html, routing::get, Router};
|
||||
use std::collections::HashMap;
|
||||
use std::net::SocketAddr;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
let (verifier, challenge) = pkce_pair()?;
|
||||
let state = random_state()?;
|
||||
|
||||
let (tx, rx) = oneshot::channel::<Result<String, String>>();
|
||||
let tx = Arc::new(Mutex::new(Some(tx)));
|
||||
|
||||
let expected_state = state.clone();
|
||||
let app = Router::new().route(
|
||||
"/",
|
||||
get(move |Query(params): Query<HashMap<String, String>>| {
|
||||
let tx = Arc::clone(&tx);
|
||||
let expected = expected_state.clone();
|
||||
async move {
|
||||
let result = match (params.get("code"), params.get("state")) {
|
||||
(Some(code), Some(st)) if st == &expected => Ok(code.clone()),
|
||||
(Some(_), Some(_)) => Err("state mismatch".to_string()),
|
||||
_ => Err(params
|
||||
.get("error")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "missing code".into())),
|
||||
};
|
||||
if let Some(sender) = tx.lock().await.take() {
|
||||
let _ = sender.send(result.clone());
|
||||
}
|
||||
match result {
|
||||
Ok(_) => Html(
|
||||
"<h2>Buzz: signed in</h2><p>You can close this window.</p>".to_string(),
|
||||
),
|
||||
Err(e) => Html(format!("<h2>Buzz auth failed</h2><pre>{e}</pre>")),
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth callback bind: {e}")))?;
|
||||
let port = listener
|
||||
.local_addr()
|
||||
.map_err(|e| AgentError::Llm(format!("oauth callback addr: {e}")))?
|
||||
.port();
|
||||
let redirect_uri = format!("http://localhost:{port}");
|
||||
|
||||
// `_server` is held until this function returns; the drop guard aborts
|
||||
// the axum task on every exit path (timeout, callback error, token
|
||||
// exchange failure, or success), so we never leak a listener bound to
|
||||
// 127.0.0.1 past the auth attempt.
|
||||
let _server = AbortOnDrop(tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app).await;
|
||||
}));
|
||||
|
||||
let auth_url = format!(
|
||||
"{}?response_type=code&client_id={}&redirect_uri={}&scope={}&state={}&code_challenge={}&code_challenge_method=S256",
|
||||
endpoints.authorization_endpoint,
|
||||
urlencoding::encode(&cfg.client_id),
|
||||
urlencoding::encode(&redirect_uri),
|
||||
urlencoding::encode(&cfg.scopes.join(" ")),
|
||||
urlencoding::encode(&state),
|
||||
urlencoding::encode(&challenge),
|
||||
);
|
||||
|
||||
eprintln!("Opening browser for authentication. If it doesn't open, visit:\n {auth_url}");
|
||||
let _ = webbrowser::open(&auth_url);
|
||||
|
||||
let code = tokio::time::timeout(BROWSER_AUTH_TIMEOUT, rx)
|
||||
.await
|
||||
.map_err(|_| AgentError::Llm("oauth: browser auth timed out".into()))?
|
||||
.map_err(|_| AgentError::Llm("oauth: callback sender dropped".into()))?
|
||||
.map_err(|e| AgentError::Llm(format!("oauth callback: {e}")))?;
|
||||
|
||||
// Exchange code for token.
|
||||
let params = [
|
||||
("grant_type", "authorization_code"),
|
||||
("code", &code),
|
||||
("redirect_uri", &redirect_uri),
|
||||
("code_verifier", &verifier),
|
||||
("client_id", &cfg.client_id),
|
||||
];
|
||||
let resp = http
|
||||
.post(&endpoints.token_endpoint)
|
||||
.form(¶ms)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth exchange: {e}")))?;
|
||||
if !resp.status().is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(AgentError::Llm(format!("oauth exchange failed: {body}")));
|
||||
}
|
||||
let v: Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth exchange json: {e}")))?;
|
||||
token_from_response(&v, None)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn pkce_pair_produces_valid_challenge() {
|
||||
let (verifier, challenge) = pkce_pair().unwrap();
|
||||
assert!(verifier.len() >= 43);
|
||||
let expected = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.encode(sha2::Sha256::digest(verifier.as_bytes()));
|
||||
assert_eq!(expected, challenge);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_token_no_expiry_is_not_expired() {
|
||||
let t = CachedToken {
|
||||
access_token: "x".into(),
|
||||
refresh_token: None,
|
||||
expires_at: None,
|
||||
};
|
||||
assert!(!is_expired(&t));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_token_far_future_is_not_expired() {
|
||||
let future = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
+ 3600;
|
||||
let t = CachedToken {
|
||||
access_token: "x".into(),
|
||||
refresh_token: None,
|
||||
expires_at: Some(future),
|
||||
};
|
||||
assert!(!is_expired(&t));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_token_within_leeway_is_expired() {
|
||||
let near = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
+ 10; // 10s away, leeway is 60s → counts as expired
|
||||
let t = CachedToken {
|
||||
access_token: "x".into(),
|
||||
refresh_token: None,
|
||||
expires_at: Some(near),
|
||||
};
|
||||
assert!(is_expired(&t));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_path_uses_platform_home_directory() {
|
||||
let cfg = PkceOAuthConfig {
|
||||
discovery_url: "https://example.com/.well-known".into(),
|
||||
client_id: "abc".into(),
|
||||
scopes: vec!["a".into(), "b".into()],
|
||||
cache_namespace: "demo".into(),
|
||||
cache_dir_override: None,
|
||||
};
|
||||
let p = cache_path_for(&cfg).unwrap();
|
||||
let expected_dir = dirs::home_dir()
|
||||
.unwrap()
|
||||
.join(".config")
|
||||
.join("buzz-agent")
|
||||
.join("oauth")
|
||||
.join("demo");
|
||||
assert_eq!(p.parent(), Some(expected_dir.as_path()));
|
||||
assert_eq!(p.extension().and_then(|s| s.to_str()), Some("json"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_from_response_uses_fallback_refresh() {
|
||||
let v: Value = serde_json::from_str(r#"{"access_token":"abc","expires_in":3600}"#).unwrap();
|
||||
let t = token_from_response(&v, Some("old-refresh")).unwrap();
|
||||
assert_eq!(t.access_token, "abc");
|
||||
assert_eq!(t.refresh_token.as_deref(), Some("old-refresh"));
|
||||
assert!(t.expires_at.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_from_response_rejects_missing_access_token() {
|
||||
let v: Value = serde_json::from_str(r#"{"expires_in":3600}"#).unwrap();
|
||||
assert!(token_from_response(&v, None).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_from_response_rejects_empty_access_token() {
|
||||
let v: Value = serde_json::from_str(r#"{"access_token":""}"#).unwrap();
|
||||
assert!(token_from_response(&v, None).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bearer_reuses_disk_token_after_expiry() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let cfg = PkceOAuthConfig {
|
||||
discovery_url: "https://example.com/.well-known".into(),
|
||||
client_id: "test-client".into(),
|
||||
scopes: vec!["offline_access".into()],
|
||||
cache_namespace: "test".into(),
|
||||
cache_dir_override: Some(dir.path().to_path_buf()),
|
||||
};
|
||||
let source = PkceOAuthTokenSource::new(cfg).unwrap();
|
||||
|
||||
// Expire the in-memory state.
|
||||
{
|
||||
let mut state = source.state.lock().await;
|
||||
*state = Some(CachedToken {
|
||||
access_token: "stale".into(),
|
||||
refresh_token: None,
|
||||
expires_at: Some(0), // long expired
|
||||
});
|
||||
}
|
||||
|
||||
// Write a valid token to disk (simulating another process refreshing).
|
||||
let future_exp = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
+ 7200;
|
||||
let fresh_token = CachedToken {
|
||||
access_token: "fresh-from-disk".into(),
|
||||
refresh_token: Some("rt".into()),
|
||||
expires_at: Some(future_exp),
|
||||
};
|
||||
let body = serde_json::to_vec_pretty(&fresh_token).unwrap();
|
||||
fs::write(&source.cache_path, &body).unwrap();
|
||||
|
||||
// bearer() should pick up the disk token without any network call.
|
||||
let result = source.bearer().await.unwrap();
|
||||
assert_eq!(result, "fresh-from-disk");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bearer_falls_through_to_browser_when_disk_also_expired() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let cfg = PkceOAuthConfig {
|
||||
discovery_url: "https://example.com/.well-known".into(),
|
||||
client_id: "test-client".into(),
|
||||
scopes: vec!["offline_access".into()],
|
||||
cache_namespace: "test".into(),
|
||||
cache_dir_override: Some(dir.path().to_path_buf()),
|
||||
};
|
||||
let source = PkceOAuthTokenSource::new(cfg).unwrap();
|
||||
|
||||
// Expire the in-memory state.
|
||||
{
|
||||
let mut state = source.state.lock().await;
|
||||
*state = Some(CachedToken {
|
||||
access_token: "stale".into(),
|
||||
refresh_token: None,
|
||||
expires_at: Some(0),
|
||||
});
|
||||
}
|
||||
|
||||
// Write an expired token to disk too.
|
||||
let expired_token = CachedToken {
|
||||
access_token: "also-stale".into(),
|
||||
refresh_token: None,
|
||||
expires_at: Some(0),
|
||||
};
|
||||
let body = serde_json::to_vec_pretty(&expired_token).unwrap();
|
||||
fs::write(&source.cache_path, &body).unwrap();
|
||||
|
||||
// bearer() should fall through past the disk check.
|
||||
// It will fail at the endpoints() discovery call since there's no server,
|
||||
// which proves it didn't short-circuit on the expired disk token.
|
||||
let result = source.bearer().await;
|
||||
assert!(result.is_err());
|
||||
let err_msg = format!("{}", result.unwrap_err());
|
||||
assert!(
|
||||
err_msg.contains("oauth discovery"),
|
||||
"expected discovery error, got: {err_msg}"
|
||||
);
|
||||
}
|
||||
|
||||
/// `try_bearer_no_browser` with an empty cache and no refresh token must
|
||||
/// return `LlmAuth` immediately — it must NOT attempt OIDC discovery even
|
||||
/// when the `discovery_url` is unreachable/invalid. This guards the
|
||||
/// regression where `endpoints()` was called unconditionally before the
|
||||
/// refresh-token check, causing an `Llm` error (hard failure) instead of
|
||||
/// the intended graceful `LlmAuth` fallback.
|
||||
#[tokio::test]
|
||||
async fn test_try_bearer_no_browser_empty_cache_no_refresh_returns_llm_auth_without_discovery()
|
||||
{
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
// Intentionally invalid/unreachable discovery URL — if endpoints() is
|
||||
// called, the test will get an `Llm` error and the assertion below fails.
|
||||
let cfg = PkceOAuthConfig {
|
||||
discovery_url: "https://invalid.example.test/.well-known/oauth-authorization-server"
|
||||
.into(),
|
||||
client_id: "test-client".into(),
|
||||
scopes: vec!["offline_access".into()],
|
||||
cache_namespace: "test".into(),
|
||||
cache_dir_override: Some(dir.path().to_path_buf()),
|
||||
};
|
||||
let source = PkceOAuthTokenSource::new(cfg).unwrap();
|
||||
|
||||
// Empty in-memory state (no token, no refresh token).
|
||||
{
|
||||
let mut state = source.state.lock().await;
|
||||
*state = None;
|
||||
}
|
||||
|
||||
// No disk cache file either — dir is empty.
|
||||
|
||||
let result = source.try_bearer_no_browser().await;
|
||||
assert!(result.is_err(), "expected Err, got Ok");
|
||||
match result.unwrap_err() {
|
||||
AgentError::LlmAuth(_) => {} // correct: graceful fallback
|
||||
other => panic!(
|
||||
"expected LlmAuth (no discovery attempted), got: {other:?}\n\
|
||||
This means endpoints() was called before the refresh-token check."
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,575 @@
|
||||
//! Built-in tools that run in-process, bypassing MCP.
|
||||
//!
|
||||
//! Currently: `load_skill` — reads a skill's full SKILL.md body from disk
|
||||
//! and returns it so the agent can load skill content on demand rather than
|
||||
//! having every skill inlined into the system prompt at session start.
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::hints::{strip_frontmatter, SkillEntry, MAX_SKILL_BODY_BYTES};
|
||||
use crate::mcp::truncate_at_boundary;
|
||||
use crate::types::{ToolDef, ToolResult, ToolResultContent};
|
||||
|
||||
pub const LOAD_SKILL_TOOL: &str = "load_skill";
|
||||
|
||||
/// Return the `ToolDef` for `load_skill` to include in the LLM tool list.
|
||||
pub fn load_skill_def() -> ToolDef {
|
||||
ToolDef {
|
||||
name: LOAD_SKILL_TOOL.to_owned(),
|
||||
description: "Load the full content of a skill by name. \
|
||||
Call this before using a skill — the system prompt lists skill names \
|
||||
and descriptions only; the full instructions are loaded on demand. \
|
||||
To load a supporting file within a skill, use the form \
|
||||
\"skill-name/relative/path\" (e.g. \"my-skill/references/foo.md\")."
|
||||
.to_owned(),
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {
|
||||
"type": "string",
|
||||
"description": "The skill name as listed in the Available Skills section, \
|
||||
or \"skill-name/relative/path\" to load a supporting file."
|
||||
}
|
||||
},
|
||||
"required": ["name"]
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a `load_skill` call. Returns a `ToolResult` on success or a
|
||||
/// user-visible error result if the skill is not found or cannot be read.
|
||||
pub async fn call_load_skill(arguments: &Value, skills: &[SkillEntry]) -> ToolResult {
|
||||
let name = match arguments.get("name").and_then(Value::as_str) {
|
||||
Some(n) => n,
|
||||
None => {
|
||||
return error_result("load_skill: missing required argument \"name\"");
|
||||
}
|
||||
};
|
||||
|
||||
// Two forms:
|
||||
// "skill-name" → load SKILL.md body + ## Supporting Files section
|
||||
// "skill-name/rel/path" → load a specific supporting file
|
||||
if let Some((skill_name, rel_path)) = name.split_once('/') {
|
||||
return load_supporting_file(skill_name, rel_path, skills).await;
|
||||
}
|
||||
|
||||
// Plain skill-name form: load SKILL.md body.
|
||||
let entry = match skills.iter().find(|s| s.name == name) {
|
||||
Some(e) => e,
|
||||
None => {
|
||||
let available: Vec<&str> = skills.iter().map(|s| s.name.as_str()).collect();
|
||||
return error_result(&format!(
|
||||
"load_skill: skill {name:?} not found. Available: {available:?}"
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Read the file off the async executor to avoid blocking a Tokio worker.
|
||||
let skill_path = entry.path.clone();
|
||||
let raw = match tokio::task::spawn_blocking(move || std::fs::read_to_string(&skill_path))
|
||||
.await
|
||||
.unwrap_or_else(|e| Err(std::io::Error::other(e)))
|
||||
{
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
return error_result(&format!("load_skill: could not read {:?}: {e}", entry.path));
|
||||
}
|
||||
};
|
||||
|
||||
// Strip the YAML frontmatter — the agent already knows name/description
|
||||
// from the system prompt; return only the body.
|
||||
let body = strip_frontmatter(&raw);
|
||||
|
||||
let mut output = body.to_owned();
|
||||
|
||||
// Append ## Supporting Files section if this skill has any.
|
||||
if !entry.supporting_files.is_empty() {
|
||||
let skill_dir = entry.path.parent().unwrap_or(&entry.path);
|
||||
output.push_str("\n\n## Supporting Files\n\n");
|
||||
for file in &entry.supporting_files {
|
||||
if let Ok(rel) = file.strip_prefix(skill_dir) {
|
||||
let rel_str = rel.to_string_lossy().replace('\\', "/");
|
||||
output.push_str(&format!(
|
||||
"- {} (load_skill(name: \"{}/{}\"))\n",
|
||||
rel_str, entry.name, rel_str
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply the size cap to the full output (body + Supporting Files section)
|
||||
// so the total tool result stays within MAX_SKILL_BODY_BYTES.
|
||||
let output = if output.len() > MAX_SKILL_BODY_BYTES {
|
||||
truncate_at_boundary(&output, MAX_SKILL_BODY_BYTES).to_owned()
|
||||
} else {
|
||||
output
|
||||
};
|
||||
|
||||
ToolResult {
|
||||
provider_id: String::new(),
|
||||
content: vec![ToolResultContent::Text(output)],
|
||||
is_error: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Load a supporting file identified by `skill_name/rel_path`.
|
||||
/// Matches against the pre-enumerated `supporting_files` list and applies a
|
||||
/// canonicalize-based traversal guard before reading.
|
||||
async fn load_supporting_file(
|
||||
skill_name: &str,
|
||||
rel_path: &str,
|
||||
skills: &[SkillEntry],
|
||||
) -> ToolResult {
|
||||
let rel_path = rel_path.replace('\\', "/");
|
||||
|
||||
let entry = match skills.iter().find(|s| s.name == skill_name) {
|
||||
Some(e) => e,
|
||||
None => {
|
||||
let available: Vec<&str> = skills.iter().map(|s| s.name.as_str()).collect();
|
||||
return error_result(&format!(
|
||||
"load_skill: skill {skill_name:?} not found. Available: {available:?}"
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let skill_dir = match entry.path.parent() {
|
||||
Some(d) => d,
|
||||
None => {
|
||||
return error_result(&format!(
|
||||
"load_skill: could not determine skill directory for {skill_name:?}"
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Match rel_path against the pre-enumerated supporting_files list.
|
||||
let matched = entry.supporting_files.iter().find(|f| {
|
||||
f.strip_prefix(skill_dir)
|
||||
.map(|r| r.to_string_lossy().replace('\\', "/") == rel_path)
|
||||
.unwrap_or(false)
|
||||
});
|
||||
|
||||
let file_path = match matched {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
let available: Vec<String> = entry
|
||||
.supporting_files
|
||||
.iter()
|
||||
.filter_map(|f| {
|
||||
f.strip_prefix(skill_dir)
|
||||
.ok()
|
||||
.map(|r| r.to_string_lossy().replace('\\', "/"))
|
||||
})
|
||||
.collect();
|
||||
if available.is_empty() {
|
||||
return error_result(&format!(
|
||||
"load_skill: skill {skill_name:?} has no supporting files."
|
||||
));
|
||||
}
|
||||
return error_result(&format!(
|
||||
"load_skill: file {rel_path:?} not found in skill {skill_name:?}. \
|
||||
Available: {available:?}"
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Traversal guard: canonicalize both paths and verify the file stays inside
|
||||
// the skill directory. Fail hard if the skill directory itself can't be
|
||||
// canonicalized — a degraded guard is worse than no guard.
|
||||
let canonical_skill_dir = match skill_dir.canonicalize() {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
return error_result(&format!(
|
||||
"load_skill: could not canonicalize skill directory for {skill_name:?}: {e}"
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Clone the path so we can move it into spawn_blocking.
|
||||
let file_path = file_path.clone();
|
||||
let skill_name = skill_name.to_owned();
|
||||
let rel_path_owned = rel_path.clone();
|
||||
|
||||
match tokio::task::spawn_blocking(move || file_path.canonicalize().map(|c| (c, file_path)))
|
||||
.await
|
||||
.unwrap_or_else(|e| Err(std::io::Error::other(e)))
|
||||
{
|
||||
Ok((canonical_file, resolved_path)) if canonical_file.starts_with(&canonical_skill_dir) => {
|
||||
match tokio::task::spawn_blocking(move || std::fs::read_to_string(&resolved_path))
|
||||
.await
|
||||
.unwrap_or_else(|e| Err(std::io::Error::other(e)))
|
||||
{
|
||||
Ok(content) => {
|
||||
let output = format!(
|
||||
"# Loaded: {}/{}\n\n{}\n\n---\nFile loaded into context.",
|
||||
skill_name, rel_path_owned, content
|
||||
);
|
||||
let output = if output.len() > MAX_SKILL_BODY_BYTES {
|
||||
truncate_at_boundary(&output, MAX_SKILL_BODY_BYTES).to_owned()
|
||||
} else {
|
||||
output
|
||||
};
|
||||
ToolResult {
|
||||
provider_id: String::new(),
|
||||
content: vec![ToolResultContent::Text(output)],
|
||||
is_error: false,
|
||||
}
|
||||
}
|
||||
Err(e) => error_result(&format!(
|
||||
"load_skill: could not read {skill_name:?}/{rel_path_owned}: {e}"
|
||||
)),
|
||||
}
|
||||
}
|
||||
Ok(_) => error_result(&format!(
|
||||
"load_skill: refusing to load {skill_name:?}/{rel_path_owned}: \
|
||||
resolves outside the skill directory"
|
||||
)),
|
||||
Err(e) => error_result(&format!(
|
||||
"load_skill: could not resolve {skill_name:?}/{rel_path_owned}: {e}"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn error_result(msg: &str) -> ToolResult {
|
||||
ToolResult {
|
||||
provider_id: String::new(),
|
||||
content: vec![ToolResultContent::Text(msg.to_owned())],
|
||||
is_error: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::path::PathBuf;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn text_content(result: &ToolResult) -> String {
|
||||
match &result.content[0] {
|
||||
ToolResultContent::Text(t) => t.clone(),
|
||||
ToolResultContent::Image { .. } => panic!("unexpected Image content in test"),
|
||||
}
|
||||
}
|
||||
|
||||
fn make_skill(name: &str, description: &str, path: PathBuf) -> SkillEntry {
|
||||
SkillEntry {
|
||||
name: name.to_owned(),
|
||||
description: description.to_owned(),
|
||||
path,
|
||||
supporting_files: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn make_skill_with_files(
|
||||
name: &str,
|
||||
description: &str,
|
||||
path: PathBuf,
|
||||
supporting_files: Vec<PathBuf>,
|
||||
) -> SkillEntry {
|
||||
SkillEntry {
|
||||
name: name.to_owned(),
|
||||
description: description.to_owned(),
|
||||
path,
|
||||
supporting_files,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_missing_name_arg() {
|
||||
let result = call_load_skill(&serde_json::json!({}), &[]).await;
|
||||
assert!(result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(text.contains("missing required argument"), "got: {text}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_skill_not_found() {
|
||||
let result = call_load_skill(&serde_json::json!({"name": "no-such"}), &[]).await;
|
||||
assert!(result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(text.contains("not found"), "got: {text}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_returns_body_strips_frontmatter() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_md = tmp.path().join("SKILL.md");
|
||||
std::fs::write(
|
||||
&skill_md,
|
||||
"---\nname: test\ndescription: A test\n---\nSkill body here.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let skills = vec![make_skill("test", "A test", skill_md)];
|
||||
let result = call_load_skill(&serde_json::json!({"name": "test"}), &skills).await;
|
||||
assert!(!result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(text.contains("Skill body here."), "got: {text}");
|
||||
assert!(
|
||||
!text.contains("---"),
|
||||
"frontmatter should be stripped: {text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_appends_supporting_files_section() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path();
|
||||
let skill_md = skill_dir.join("SKILL.md");
|
||||
std::fs::write(
|
||||
&skill_md,
|
||||
"---\nname: my-skill\ndescription: desc\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let refs_dir = skill_dir.join("references");
|
||||
std::fs::create_dir_all(&refs_dir).unwrap();
|
||||
let ref_file = refs_dir.join("foo.md");
|
||||
std::fs::write(&ref_file, "Reference content.").unwrap();
|
||||
|
||||
let skills = vec![make_skill_with_files(
|
||||
"my-skill",
|
||||
"desc",
|
||||
skill_md,
|
||||
vec![ref_file],
|
||||
)];
|
||||
let result = call_load_skill(&serde_json::json!({"name": "my-skill"}), &skills).await;
|
||||
assert!(!result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(text.contains("Body."), "body missing: {text}");
|
||||
assert!(
|
||||
text.contains("## Supporting Files"),
|
||||
"missing Supporting Files section: {text}"
|
||||
);
|
||||
assert!(
|
||||
text.contains("references/foo.md"),
|
||||
"missing file listing: {text}"
|
||||
);
|
||||
assert!(
|
||||
text.contains("load_skill(name: \"my-skill/references/foo.md\")"),
|
||||
"missing load_skill hint: {text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_no_supporting_files_section_when_empty() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_md = tmp.path().join("SKILL.md");
|
||||
std::fs::write(
|
||||
&skill_md,
|
||||
"---\nname: bare\ndescription: desc\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let skills = vec![make_skill("bare", "desc", skill_md)];
|
||||
let result = call_load_skill(&serde_json::json!({"name": "bare"}), &skills).await;
|
||||
assert!(!result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(
|
||||
!text.contains("## Supporting Files"),
|
||||
"should not have Supporting Files section when none: {text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_supporting_file_returns_content() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path();
|
||||
let skill_md = skill_dir.join("SKILL.md");
|
||||
std::fs::write(
|
||||
&skill_md,
|
||||
"---\nname: my-skill\ndescription: desc\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let refs_dir = skill_dir.join("references");
|
||||
std::fs::create_dir_all(&refs_dir).unwrap();
|
||||
let ref_file = refs_dir.join("foo.md");
|
||||
std::fs::write(&ref_file, "Reference content here.").unwrap();
|
||||
|
||||
let skills = vec![make_skill_with_files(
|
||||
"my-skill",
|
||||
"desc",
|
||||
skill_md,
|
||||
vec![ref_file],
|
||||
)];
|
||||
let result = call_load_skill(
|
||||
&serde_json::json!({"name": "my-skill/references/foo.md"}),
|
||||
&skills,
|
||||
)
|
||||
.await;
|
||||
assert!(!result.is_error, "expected success, got error");
|
||||
let text = text_content(&result);
|
||||
assert!(
|
||||
text.contains("Reference content here."),
|
||||
"file content missing: {text}"
|
||||
);
|
||||
assert!(
|
||||
text.contains("# Loaded: my-skill/references/foo.md"),
|
||||
"missing header: {text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_supporting_file_not_found_lists_available() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path();
|
||||
let skill_md = skill_dir.join("SKILL.md");
|
||||
std::fs::write(
|
||||
&skill_md,
|
||||
"---\nname: my-skill\ndescription: desc\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let refs_dir = skill_dir.join("references");
|
||||
std::fs::create_dir_all(&refs_dir).unwrap();
|
||||
let ref_file = refs_dir.join("foo.md");
|
||||
std::fs::write(&ref_file, "content").unwrap();
|
||||
|
||||
let skills = vec![make_skill_with_files(
|
||||
"my-skill",
|
||||
"desc",
|
||||
skill_md,
|
||||
vec![ref_file],
|
||||
)];
|
||||
let result = call_load_skill(
|
||||
&serde_json::json!({"name": "my-skill/references/missing.md"}),
|
||||
&skills,
|
||||
)
|
||||
.await;
|
||||
assert!(result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(text.contains("not found"), "got: {text}");
|
||||
assert!(
|
||||
text.contains("references/foo.md"),
|
||||
"should list available: {text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_no_supporting_files_error_message() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_md = tmp.path().join("SKILL.md");
|
||||
std::fs::write(
|
||||
&skill_md,
|
||||
"---\nname: bare\ndescription: desc\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let skills = vec![make_skill("bare", "desc", skill_md)];
|
||||
let result =
|
||||
call_load_skill(&serde_json::json!({"name": "bare/anything.md"}), &skills).await;
|
||||
assert!(result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(text.contains("no supporting files"), "got: {text}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_traversal_guard_rejects_escape() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path().join("my-skill");
|
||||
std::fs::create_dir_all(&skill_dir).unwrap();
|
||||
let skill_md = skill_dir.join("SKILL.md");
|
||||
std::fs::write(
|
||||
&skill_md,
|
||||
"---\nname: my-skill\ndescription: desc\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Create a file outside the skill dir that we'll try to reference.
|
||||
let outside_file = tmp.path().join("secret.txt");
|
||||
std::fs::write(&outside_file, "secret content").unwrap();
|
||||
|
||||
// Manually construct a SkillEntry with a supporting_files entry that
|
||||
// points outside the skill dir — simulating a crafted/malicious entry.
|
||||
// The traversal guard should catch this.
|
||||
let skills = vec![make_skill_with_files(
|
||||
"my-skill",
|
||||
"desc",
|
||||
skill_md.clone(),
|
||||
vec![outside_file.clone()],
|
||||
)];
|
||||
|
||||
// The slash form splits "my-skill/../secret.txt" into skill_name="my-skill"
|
||||
// and rel_path="../secret.txt". strip_prefix(skill_dir) on outside_file
|
||||
// fails, so it won't match any supporting_files entry — the pre-enumeration
|
||||
// guard rejects it before the canonicalize guard even fires.
|
||||
let result = call_load_skill(
|
||||
&serde_json::json!({"name": "my-skill/../secret.txt"}),
|
||||
&skills,
|
||||
)
|
||||
.await;
|
||||
assert!(result.is_error, "traversal attempt should be rejected");
|
||||
let text = text_content(&result);
|
||||
assert!(
|
||||
!text.contains("secret content"),
|
||||
"secret content must not be returned: {text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_truncates_large_body() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path();
|
||||
let skill_md = skill_dir.join("SKILL.md");
|
||||
// Build a body that exceeds MAX_SKILL_BODY_BYTES (32 KiB).
|
||||
let large_body = "x".repeat(40 * 1024);
|
||||
std::fs::write(
|
||||
&skill_md,
|
||||
format!("---\nname: big\ndescription: desc\n---\n{large_body}\n"),
|
||||
)
|
||||
.unwrap();
|
||||
// Add a supporting file so the Supporting Files section is also appended
|
||||
// before the cap is applied.
|
||||
let refs_dir = skill_dir.join("references");
|
||||
std::fs::create_dir_all(&refs_dir).unwrap();
|
||||
let ref_file = refs_dir.join("extra.md");
|
||||
std::fs::write(&ref_file, "extra content").unwrap();
|
||||
|
||||
let skills = vec![make_skill_with_files(
|
||||
"big",
|
||||
"desc",
|
||||
skill_md,
|
||||
vec![ref_file],
|
||||
)];
|
||||
let result = call_load_skill(&serde_json::json!({"name": "big"}), &skills).await;
|
||||
assert!(!result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(
|
||||
text.len() <= MAX_SKILL_BODY_BYTES,
|
||||
"output length {} exceeds MAX_SKILL_BODY_BYTES {}",
|
||||
text.len(),
|
||||
MAX_SKILL_BODY_BYTES
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_load_skill_truncates_large_supporting_file() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path();
|
||||
let skill_md = skill_dir.join("SKILL.md");
|
||||
std::fs::write(&skill_md, "---\nname: big\ndescription: desc\n---\nBody.\n").unwrap();
|
||||
|
||||
let refs_dir = skill_dir.join("references");
|
||||
std::fs::create_dir_all(&refs_dir).unwrap();
|
||||
let ref_file = refs_dir.join("huge.md");
|
||||
std::fs::write(&ref_file, "x".repeat(MAX_SKILL_BODY_BYTES * 2)).unwrap();
|
||||
|
||||
let skills = vec![make_skill_with_files(
|
||||
"big",
|
||||
"desc",
|
||||
skill_md,
|
||||
vec![ref_file],
|
||||
)];
|
||||
let result = call_load_skill(
|
||||
&serde_json::json!({"name": "big/references/huge.md"}),
|
||||
&skills,
|
||||
)
|
||||
.await;
|
||||
assert!(!result.is_error);
|
||||
let text = text_content(&result);
|
||||
assert!(
|
||||
text.len() <= MAX_SKILL_BODY_BYTES,
|
||||
"output length {} exceeds MAX_SKILL_BODY_BYTES {}",
|
||||
text.len(),
|
||||
MAX_SKILL_BODY_BYTES
|
||||
);
|
||||
assert!(
|
||||
text.starts_with("# Loaded: big/references/huge.md"),
|
||||
"missing supporting-file header: {text}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,668 @@
|
||||
//! Databricks model catalog discovery.
|
||||
//!
|
||||
//! Exposes [`discover_databricks_models`] — an async helper that lists
|
||||
//! available models for the `databricks` and `databricks_v2` providers
|
||||
//! without triggering a browser OAuth flow. Auth is acquired in-process via
|
||||
//! [`build_token_source`](crate::llm::build_token_source):
|
||||
//!
|
||||
//! - Static bearer (`DATABRICKS_TOKEN`): returned immediately.
|
||||
//! - PKCE cache hit: returned from disk without a network round-trip.
|
||||
//! - PKCE cache empty / no token: returns `Err(AgentError::LlmAuth)`.
|
||||
//!
|
||||
//! This helper never opens a browser. Callers choose whether to reject, degrade,
|
||||
//! or start a separate interactive authentication flow.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use reqwest::Client;
|
||||
|
||||
use crate::{
|
||||
auth::TokenSource,
|
||||
config::{Config, Provider},
|
||||
llm::build_token_source,
|
||||
types::AgentError,
|
||||
};
|
||||
|
||||
/// A discovered model entry: `id` is the picker value, `name` is the display
|
||||
/// label (same as `id` for Databricks — the API has no separate display name).
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ModelEntry {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
/// Known Databricks AI Gateway v2 models — used only when an authenticated
|
||||
/// `api/ai-gateway/v2/endpoints` call succeeds with an empty list.
|
||||
/// Mirrors goose's `DATABRICKS_V2_KNOWN_MODELS`.
|
||||
pub const DATABRICKS_V2_KNOWN_MODELS: &[&str] =
|
||||
&["databricks-gpt-5-5", "databricks-claude-opus-4-7"];
|
||||
|
||||
const AUTHENTICATED_EMPTY_CATALOG_SUFFIX: &str = " (default catalog)";
|
||||
|
||||
fn authenticated_empty_v2_catalog() -> Vec<ModelEntry> {
|
||||
DATABRICKS_V2_KNOWN_MODELS
|
||||
.iter()
|
||||
.map(|id| ModelEntry {
|
||||
id: id.to_string(),
|
||||
name: format!("{id}{AUTHENTICATED_EMPTY_CATALOG_SUFFIX}"),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Heuristic: `true` when a v2 AI Gateway endpoint name looks like it serves
|
||||
/// chat/completions traffic.
|
||||
///
|
||||
/// The v1 `serving-endpoints` payload carries `task`, so [`parse_v1_endpoints`]
|
||||
/// can filter on it directly. The v2 `ai-gateway/v2/endpoints` payload carries
|
||||
/// no task or readiness field at all, so the only signal available here is the
|
||||
/// endpoint name. Embedding endpoints are the one family that reliably cannot
|
||||
/// serve a chat request — they reject it with
|
||||
/// `API type 'mlflow/v1/chat/completions' is not supported by '<name>'` — so
|
||||
/// they are dropped rather than offered as selectable models.
|
||||
///
|
||||
/// Deliberately narrow: image-capable endpoints (e.g.
|
||||
/// `databricks-gemini-3-pro-image`) do answer chat requests, so they stay. Any
|
||||
/// name this heuristic does not recognise is kept — preferring to include over
|
||||
/// silently dropping, matching [`parse_v1_endpoints`].
|
||||
pub(crate) fn is_chat_capable_endpoint(name: &str) -> bool {
|
||||
let lower = name.to_ascii_lowercase();
|
||||
if lower.contains("embedding") {
|
||||
return false;
|
||||
}
|
||||
// Segment match so `bge`/`gte` cannot fire on a substring of a longer word.
|
||||
!lower
|
||||
.split('-')
|
||||
.any(|segment| matches!(segment, "bge" | "gte"))
|
||||
}
|
||||
|
||||
/// Discover available models for a Databricks provider.
|
||||
///
|
||||
/// Returns a non-empty `Vec<ModelEntry>` on success. Returns
|
||||
/// `Err(AgentError::LlmAuth)` when no token is available (no static token,
|
||||
/// no PKCE cache). The helper itself never starts interactive authentication.
|
||||
///
|
||||
/// # Panics
|
||||
/// Never panics.
|
||||
pub async fn discover_databricks_models(cfg: &Config) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
discover_databricks_models_with_token_source(cfg, build_token_source(cfg)?).await
|
||||
}
|
||||
|
||||
async fn discover_databricks_models_with_token_source(
|
||||
cfg: &Config,
|
||||
token_source: Arc<dyn TokenSource>,
|
||||
) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
let mut bearer = token_source.bearer_no_browser().await?;
|
||||
let http = Client::new();
|
||||
let host = cfg.base_url.trim_end_matches('/');
|
||||
let mut refreshed = false;
|
||||
|
||||
loop {
|
||||
let result = match cfg.provider {
|
||||
Provider::Databricks => fetch_v1_models(&http, host, &bearer).await,
|
||||
Provider::DatabricksV2 => fetch_v2_models(&http, host, &bearer).await,
|
||||
_ => {
|
||||
return Err(AgentError::InvalidParams(
|
||||
"discover_databricks_models called for non-Databricks provider".into(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
match result {
|
||||
Err(AgentError::LlmAuth(_)) if !refreshed => {
|
||||
refreshed = true;
|
||||
let fresh = token_source.refresh_now(&bearer).await?;
|
||||
if fresh == bearer {
|
||||
return Err(AgentError::LlmAuth(
|
||||
"Databricks rejected the configured credential".into(),
|
||||
));
|
||||
}
|
||||
bearer = fresh;
|
||||
}
|
||||
result => return result,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// v1 — api/2.0/serving-endpoints
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn fetch_v1_models(
|
||||
http: &Client,
|
||||
host: &str,
|
||||
bearer: &str,
|
||||
) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
let url = format!("{host}/api/2.0/serving-endpoints");
|
||||
let response = http
|
||||
.get(&url)
|
||||
.bearer_auth(bearer)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("Databricks model discovery request failed: {e}")))?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
if status.as_u16() == 401 {
|
||||
return Err(AgentError::LlmAuth(format!(
|
||||
"Databricks model discovery HTTP {status}"
|
||||
)));
|
||||
}
|
||||
return Err(AgentError::Llm(format!(
|
||||
"Databricks model discovery HTTP {status}: {body}"
|
||||
)));
|
||||
}
|
||||
|
||||
let json: serde_json::Value = response.json().await.map_err(|e| {
|
||||
AgentError::Llm(format!(
|
||||
"Databricks model discovery response parse failed: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
parse_v1_endpoints(&json)
|
||||
}
|
||||
|
||||
/// Parse a `GET api/2.0/serving-endpoints` response.
|
||||
///
|
||||
/// Filters to endpoints that are READY and serve an LLM chat/completions task.
|
||||
/// When `state.ready` or `task` is absent the endpoint is included — prefer
|
||||
/// including over silently dropping, per spec.
|
||||
pub(crate) fn parse_v1_endpoints(json: &serde_json::Value) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
let endpoints = json
|
||||
.get("endpoints")
|
||||
.and_then(|v| v.as_array())
|
||||
.ok_or_else(|| {
|
||||
AgentError::Llm(
|
||||
"Databricks model discovery: unexpected response (missing 'endpoints' array)"
|
||||
.into(),
|
||||
)
|
||||
})?;
|
||||
|
||||
let models = endpoints
|
||||
.iter()
|
||||
.filter_map(|endpoint| {
|
||||
let name = endpoint.get("name")?.as_str()?.to_string();
|
||||
|
||||
// Require READY state when present; include when absent.
|
||||
let state_ready = endpoint
|
||||
.get("state")
|
||||
.and_then(|s| s.get("ready"))
|
||||
.and_then(|r| r.as_str())
|
||||
.map(|r| r == "READY")
|
||||
.unwrap_or(true);
|
||||
if !state_ready {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Require LLM chat or completions task when present.
|
||||
let task_ok = endpoint
|
||||
.get("task")
|
||||
.and_then(|t| t.as_str())
|
||||
.map(|t| t == "llm/v1/chat" || t == "llm/v1/completions")
|
||||
.unwrap_or(true);
|
||||
if !task_ok {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(ModelEntry {
|
||||
id: name.clone(),
|
||||
name,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(models)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// v2 — api/ai-gateway/v2/endpoints (paginated)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Percent-encode a string for use as a URL query parameter value.
|
||||
/// Only encodes characters that are not unreserved (RFC 3986).
|
||||
fn percent_encode(s: &str) -> String {
|
||||
s.bytes()
|
||||
.flat_map(|b| match b {
|
||||
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
|
||||
vec![b as char]
|
||||
}
|
||||
_ => format!("%{b:02X}").chars().collect(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn fetch_v2_models(
|
||||
http: &Client,
|
||||
host: &str,
|
||||
bearer: &str,
|
||||
) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
let mut all_endpoints: Vec<V2Endpoint> = Vec::new();
|
||||
let mut page_token: Option<String> = None;
|
||||
let base_url = format!("{host}/api/ai-gateway/v2/endpoints");
|
||||
|
||||
// Cap at 20 pages (2 000 endpoints) to bound execution time.
|
||||
for _ in 0..20 {
|
||||
// Build URL with query params manually — avoids requiring the `query`
|
||||
// reqwest feature in buzz-agent's Cargo.toml.
|
||||
let url = match &page_token {
|
||||
Some(tok) => format!(
|
||||
"{base_url}?page_size=100&page_token={}",
|
||||
percent_encode(tok)
|
||||
),
|
||||
None => format!("{base_url}?page_size=100"),
|
||||
};
|
||||
let response = http
|
||||
.get(&url)
|
||||
.bearer_auth(bearer)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
AgentError::Llm(format!("Databricks v2 model discovery request failed: {e}"))
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
if status.as_u16() == 401 {
|
||||
return Err(AgentError::LlmAuth(format!(
|
||||
"Databricks v2 model discovery HTTP {status}"
|
||||
)));
|
||||
}
|
||||
return Err(AgentError::Llm(format!(
|
||||
"Databricks v2 model discovery HTTP {status}: {body}"
|
||||
)));
|
||||
}
|
||||
|
||||
let json: serde_json::Value = response.json().await.map_err(|e| {
|
||||
AgentError::Llm(format!(
|
||||
"Databricks v2 model discovery response parse failed: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let (page_endpoints, next) = parse_v2_endpoints_page(&json)?;
|
||||
all_endpoints.extend(page_endpoints);
|
||||
|
||||
match next {
|
||||
Some(tok) if Some(&tok) != page_token.as_ref() => page_token = Some(tok),
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to known-model list if the API returned nothing.
|
||||
if all_endpoints.is_empty() {
|
||||
return Ok(authenticated_empty_v2_catalog());
|
||||
}
|
||||
|
||||
sort_v2_endpoints_newest_first(&mut all_endpoints);
|
||||
|
||||
Ok(all_endpoints
|
||||
.into_iter()
|
||||
.map(|endpoint| endpoint.entry)
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// A v2 gateway endpoint plus the key discovery orders the catalog by.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct V2Endpoint {
|
||||
pub(crate) entry: ModelEntry,
|
||||
/// `created_timestamp` as epoch milliseconds. `None` when the field is
|
||||
/// absent or unparseable — those sort last rather than jumping the queue.
|
||||
pub(crate) created_ms: Option<i64>,
|
||||
}
|
||||
|
||||
/// Read `created_timestamp` from one endpoint object.
|
||||
///
|
||||
/// The gateway sends epoch milliseconds as a JSON *string*
|
||||
/// (`"created_timestamp": "1699610000000"`); accept a bare number too, so a
|
||||
/// wire-shape change doesn't silently drop every endpoint to the bottom.
|
||||
fn endpoint_created_ms(endpoint: &serde_json::Value) -> Option<i64> {
|
||||
let value = endpoint.get("created_timestamp")?;
|
||||
value
|
||||
.as_i64()
|
||||
.or_else(|| value.as_str()?.trim().parse::<i64>().ok())
|
||||
}
|
||||
|
||||
/// Order the catalog newest-first, breaking ties by name.
|
||||
///
|
||||
/// The gateway returns endpoints in two phases — Databricks-managed first, then
|
||||
/// workspace-created — each alphabetical by name, which buries a brand-new
|
||||
/// frontier model deep in the list. Newest-first puts the models people are
|
||||
/// reaching for at the top of the picker.
|
||||
///
|
||||
/// Endpoints with no usable timestamp sort last, and the name tiebreak keeps the
|
||||
/// result stable: several managed endpoints share one placeholder timestamp, so
|
||||
/// without it their relative order would be arbitrary.
|
||||
pub(crate) fn sort_v2_endpoints_newest_first(endpoints: &mut [V2Endpoint]) {
|
||||
endpoints.sort_by(|a, b| {
|
||||
// `None` < `Some(_)`, so reversing puts timestamped endpoints first.
|
||||
b.created_ms
|
||||
.cmp(&a.created_ms)
|
||||
.then_with(|| a.entry.name.cmp(&b.entry.name))
|
||||
});
|
||||
}
|
||||
|
||||
/// Parse one page of a `GET api/ai-gateway/v2/endpoints` response.
|
||||
///
|
||||
/// Returns `(endpoints, next_page_token)`. An empty or absent `next_page_token`
|
||||
/// signals the last page. Endpoints that cannot serve chat traffic are dropped
|
||||
/// (see [`is_chat_capable_endpoint`]) so the model picker only offers models the
|
||||
/// agent can actually run. Page order is preserved here; the caller sorts once
|
||||
/// every page is in (see [`sort_v2_endpoints_newest_first`]).
|
||||
pub(crate) fn parse_v2_endpoints_page(
|
||||
json: &serde_json::Value,
|
||||
) -> Result<(Vec<V2Endpoint>, Option<String>), AgentError> {
|
||||
let endpoints = json
|
||||
.get("endpoints")
|
||||
.and_then(|v| v.as_array())
|
||||
.ok_or_else(|| {
|
||||
AgentError::Llm(
|
||||
"Databricks v2 model discovery: unexpected response (missing 'endpoints' array)"
|
||||
.into(),
|
||||
)
|
||||
})?;
|
||||
|
||||
let models = endpoints
|
||||
.iter()
|
||||
.filter_map(|endpoint| {
|
||||
let name = endpoint.get("name")?.as_str()?.to_string();
|
||||
if !is_chat_capable_endpoint(&name) {
|
||||
return None;
|
||||
}
|
||||
Some(V2Endpoint {
|
||||
entry: ModelEntry {
|
||||
id: name.clone(),
|
||||
name,
|
||||
},
|
||||
created_ms: endpoint_created_ms(endpoint),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let next_page_token = json
|
||||
.get("next_page_token")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|token| !token.is_empty())
|
||||
.map(str::to_string);
|
||||
|
||||
Ok((models, next_page_token))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
struct RefreshingTestTokenSource {
|
||||
refreshes: AtomicUsize,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenSource for RefreshingTestTokenSource {
|
||||
async fn bearer(&self) -> Result<String, AgentError> {
|
||||
Ok("rejected".into())
|
||||
}
|
||||
|
||||
async fn refresh_now(&self, rejected: &str) -> Result<String, AgentError> {
|
||||
assert_eq!(rejected, "rejected");
|
||||
self.refreshes.fetch_add(1, Ordering::SeqCst);
|
||||
Ok("fresh".into())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discovery_refreshes_rejected_bearer_once_then_retries_successfully() {
|
||||
use axum::{
|
||||
extract::Query,
|
||||
http::{HeaderMap, StatusCode},
|
||||
routing::get,
|
||||
Json, Router,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let host = format!("http://{}", listener.local_addr().unwrap());
|
||||
let requests = Arc::new(AtomicUsize::new(0));
|
||||
let requests_for_route = requests.clone();
|
||||
let app = Router::new().route(
|
||||
"/api/ai-gateway/v2/endpoints",
|
||||
get(
|
||||
move |headers: HeaderMap, Query(_query): Query<HashMap<String, String>>| {
|
||||
let requests = requests_for_route.clone();
|
||||
async move {
|
||||
requests.fetch_add(1, Ordering::SeqCst);
|
||||
match headers
|
||||
.get("authorization")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
{
|
||||
Some("Bearer fresh") => Ok(Json(serde_json::json!({
|
||||
"endpoints": [{"name": "discovered-model"}],
|
||||
"next_page_token": null,
|
||||
}))),
|
||||
_ => Err((StatusCode::UNAUTHORIZED, "rejected")),
|
||||
}
|
||||
}
|
||||
},
|
||||
),
|
||||
);
|
||||
tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app).await;
|
||||
});
|
||||
|
||||
let source = Arc::new(RefreshingTestTokenSource {
|
||||
refreshes: AtomicUsize::new(0),
|
||||
});
|
||||
let cfg = Config::for_discovery(Provider::DatabricksV2, String::new(), host);
|
||||
let models = discover_databricks_models_with_token_source(&cfg, source.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(models[0].id, "discovered-model");
|
||||
assert_eq!(source.refreshes.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(requests.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v1_parse_filters_ready_chat_endpoints() {
|
||||
let json = serde_json::json!({
|
||||
"endpoints": [
|
||||
// included: READY + llm/v1/chat
|
||||
{"name": "my-llm", "state": {"ready": "READY"}, "task": "llm/v1/chat"},
|
||||
// included: READY + llm/v1/completions
|
||||
{"name": "my-completions", "state": {"ready": "READY"}, "task": "llm/v1/completions"},
|
||||
// excluded: NOT_READY
|
||||
{"name": "dead-endpoint", "state": {"ready": "NOT_READY"}, "task": "llm/v1/chat"},
|
||||
// excluded: wrong task
|
||||
{"name": "embedding-ep", "state": {"ready": "READY"}, "task": "llm/v1/embedding"},
|
||||
// included: no state field → include by default
|
||||
{"name": "no-state", "task": "llm/v1/chat"},
|
||||
// included: no task field → include by default
|
||||
{"name": "no-task", "state": {"ready": "READY"}},
|
||||
]
|
||||
});
|
||||
|
||||
let models = parse_v1_endpoints(&json).unwrap();
|
||||
let ids: Vec<&str> = models.iter().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(ids, vec!["my-llm", "my-completions", "no-state", "no-task"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v1_parse_errors_on_missing_endpoints_array() {
|
||||
let json = serde_json::json!({"data": []});
|
||||
let err = parse_v1_endpoints(&json).unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("missing 'endpoints' array"),
|
||||
"got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v1_parse_empty_endpoints_returns_empty_vec() {
|
||||
let json = serde_json::json!({"endpoints": []});
|
||||
let models = parse_v1_endpoints(&json).unwrap();
|
||||
assert!(models.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v2_parse_extracts_names_and_page_token() {
|
||||
let json = serde_json::json!({
|
||||
"endpoints": [
|
||||
{"name": "databricks-claude-opus-4-7"},
|
||||
{"name": "databricks-gpt-5-5"},
|
||||
{"name": "custom-model"}
|
||||
],
|
||||
"next_page_token": "tok123"
|
||||
});
|
||||
|
||||
let (models, next) = parse_v2_endpoints_page(&json).unwrap();
|
||||
let ids: Vec<&str> = models.iter().map(|m| m.entry.id.as_str()).collect();
|
||||
assert_eq!(
|
||||
ids,
|
||||
vec![
|
||||
"databricks-claude-opus-4-7",
|
||||
"databricks-gpt-5-5",
|
||||
"custom-model"
|
||||
]
|
||||
);
|
||||
assert_eq!(next.as_deref(), Some("tok123"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v2_parse_empty_token_signals_last_page() {
|
||||
let json = serde_json::json!({
|
||||
"endpoints": [{"name": "only-model"}],
|
||||
"next_page_token": ""
|
||||
});
|
||||
|
||||
let (models, next) = parse_v2_endpoints_page(&json).unwrap();
|
||||
assert_eq!(models.len(), 1);
|
||||
assert!(
|
||||
next.is_none(),
|
||||
"empty token should be treated as no more pages"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v2_parse_absent_token_signals_last_page() {
|
||||
let json = serde_json::json!({"endpoints": [{"name": "only-model"}]});
|
||||
let (_, next) = parse_v2_endpoints_page(&json).unwrap();
|
||||
assert!(next.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v2_parse_errors_on_missing_endpoints_array() {
|
||||
let json = serde_json::json!({"data": []});
|
||||
let err = parse_v2_endpoints_page(&json).unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("missing 'endpoints' array"),
|
||||
"got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v2_parse_drops_embedding_endpoints() {
|
||||
// The v2 payload carries no `task`, so embedding endpoints are only
|
||||
// recognisable by name. They reject chat requests, so offering them in
|
||||
// the picker can only produce a 400 at send time.
|
||||
let json = serde_json::json!({
|
||||
"endpoints": [
|
||||
{"name": "databricks-bge-large-en"},
|
||||
{"name": "databricks-gte-large-en"},
|
||||
{"name": "databricks-qwen3-embedding-0-6b"},
|
||||
{"name": "databricks-claude-opus-5"},
|
||||
{"name": "databricks-gemini-3-pro-image"},
|
||||
]
|
||||
});
|
||||
|
||||
let (models, _) = parse_v2_endpoints_page(&json).unwrap();
|
||||
let ids: Vec<&str> = models.iter().map(|m| m.entry.id.as_str()).collect();
|
||||
// Image endpoints DO answer chat requests, so they are retained.
|
||||
assert_eq!(
|
||||
ids,
|
||||
vec!["databricks-claude-opus-5", "databricks-gemini-3-pro-image"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v2_parse_reads_created_timestamp_in_either_wire_shape() {
|
||||
// The gateway sends epoch ms as a string; a bare number must work too.
|
||||
let json = serde_json::json!({
|
||||
"endpoints": [
|
||||
{"name": "string-ts", "created_timestamp": "1784932442251"},
|
||||
{"name": "number-ts", "created_timestamp": 1784932442251i64},
|
||||
{"name": "junk-ts", "created_timestamp": "not-a-number"},
|
||||
{"name": "no-ts"},
|
||||
]
|
||||
});
|
||||
|
||||
let (models, _) = parse_v2_endpoints_page(&json).unwrap();
|
||||
let stamps: Vec<Option<i64>> = models.iter().map(|m| m.created_ms).collect();
|
||||
assert_eq!(
|
||||
stamps,
|
||||
vec![Some(1784932442251), Some(1784932442251), None, None,]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn v2_endpoints_sort_newest_first_then_by_name() {
|
||||
// Mirrors the real catalog: the gateway pages Databricks-managed
|
||||
// endpoints first, then workspace-created ones, each alphabetical — so
|
||||
// the newest model is buried mid-list until this sort runs.
|
||||
let json = serde_json::json!({
|
||||
"endpoints": [
|
||||
{"name": "databricks-claude-opus-5", "created_timestamp": "1784851200000"},
|
||||
{"name": "databricks-gpt-5-6-sol", "created_timestamp": "1784073600000"},
|
||||
{"name": "databricks-gpt-5-6-luna", "created_timestamp": "1784073600000"},
|
||||
{"name": "databricks-llama-4-maverick", "created_timestamp": "1699610000000"},
|
||||
{"name": "goose-claude-opus-5", "created_timestamp": "1784932442251"},
|
||||
{"name": "endpoint-without-timestamp"},
|
||||
]
|
||||
});
|
||||
|
||||
let (mut models, _) = parse_v2_endpoints_page(&json).unwrap();
|
||||
sort_v2_endpoints_newest_first(&mut models);
|
||||
|
||||
let ids: Vec<&str> = models.iter().map(|m| m.entry.id.as_str()).collect();
|
||||
assert_eq!(
|
||||
ids,
|
||||
vec![
|
||||
// Newest first, across both pagination phases.
|
||||
"goose-claude-opus-5",
|
||||
"databricks-claude-opus-5",
|
||||
// Same timestamp — the name tiebreak keeps this deterministic.
|
||||
"databricks-gpt-5-6-luna",
|
||||
"databricks-gpt-5-6-sol",
|
||||
"databricks-llama-4-maverick",
|
||||
// No usable timestamp sorts last, never first.
|
||||
"endpoint-without-timestamp",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn authenticated_empty_v2_catalog_marks_fallback_provenance() {
|
||||
let models = authenticated_empty_v2_catalog();
|
||||
let ids: Vec<&str> = models.iter().map(|model| model.id.as_str()).collect();
|
||||
|
||||
assert_eq!(ids, DATABRICKS_V2_KNOWN_MODELS);
|
||||
assert!(models.iter().all(|model| {
|
||||
model.name == format!("{}{AUTHENTICATED_EMPTY_CATALOG_SUFFIX}", model.id)
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_chat_capable_endpoint_keeps_unrecognised_names() {
|
||||
// Prefer including over silently dropping — an unknown family is kept.
|
||||
assert!(is_chat_capable_endpoint("databricks-glm-5-2"));
|
||||
assert!(is_chat_capable_endpoint("some-teams-custom-endpoint"));
|
||||
// `bge`/`gte` match as whole segments only, never as substrings.
|
||||
assert!(is_chat_capable_endpoint("databricks-budget-gtex-model"));
|
||||
assert!(!is_chat_capable_endpoint("databricks-bge-large-en"));
|
||||
assert!(!is_chat_capable_endpoint("databricks-gte-large-en"));
|
||||
assert!(!is_chat_capable_endpoint("databricks-qwen3-embedding-0-6b"));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,650 @@
|
||||
use crate::agent::RunCtx;
|
||||
use crate::config::{
|
||||
HANDOFF_MAX_OUTPUT_TOKENS, HANDOFF_MAX_TOOL_NAMES, HANDOFF_MIN_PROMPT_BUDGET_BYTES,
|
||||
HANDOFF_ORIGINAL_TASK_MAX_BYTES, MAX_CONTEXT_RECOVERIES_PER_RUN,
|
||||
};
|
||||
use crate::llm::summary_completion_cap;
|
||||
use crate::types::HistoryItem;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
struct HandoffTokenCounts {
|
||||
before: u64,
|
||||
after: u64,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for HandoffTokenCounts {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{} -> {} tokens", self.before, self.after)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) enum HandoffOutcome {
|
||||
Performed,
|
||||
Skipped,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
/// Result of the reactive context-recovery ladder.
|
||||
pub(crate) enum ContextRecovery {
|
||||
/// History was reset; the caller should retry the request.
|
||||
Recovered,
|
||||
/// Cancelled mid-recovery.
|
||||
Cancelled,
|
||||
/// No rescue remains — the caller must surface the provider error. Either
|
||||
/// the per-`run()` budget is spent or the prompt budget fell below the
|
||||
/// floor where a summary can still be useful.
|
||||
Exhausted,
|
||||
}
|
||||
|
||||
/// System prompt for the handoff summarizer. `LazyLock` + `format!` so the
|
||||
/// token figure is derived from [`HANDOFF_MAX_OUTPUT_TOKENS`] instead of a
|
||||
/// duplicated literal, and "visible plain-text summary" makes explicit that
|
||||
/// the limit is on summary text, not on any hidden reasoning the model does
|
||||
/// first (which is budgeted separately on the wire — see
|
||||
/// `openrouter_summary_body`).
|
||||
static HANDOFF_SYSTEM_PROMPT: std::sync::LazyLock<String> = std::sync::LazyLock::new(|| {
|
||||
format!(
|
||||
"You are generating a context handoff summary for the next turn of an autonomous agent. \
|
||||
Be concise but thorough. Cover: what the original task was, what you accomplished, key \
|
||||
decisions made, what remains, and one concrete next step. Output plain text only — no \
|
||||
tool calls, no JSON. Keep the visible plain-text summary under \
|
||||
{HANDOFF_MAX_OUTPUT_TOKENS} tokens."
|
||||
)
|
||||
});
|
||||
|
||||
impl RunCtx<'_> {
|
||||
pub(crate) async fn maybe_handoff(&mut self, handoff_attempts: &mut usize) -> HandoffOutcome {
|
||||
if !self.should_handoff() {
|
||||
return HandoffOutcome::Skipped;
|
||||
}
|
||||
if *handoff_attempts >= self.cfg.max_handoffs {
|
||||
let projected = self.projected_handoff_input_tokens();
|
||||
let threshold =
|
||||
token_threshold(self.cfg.max_context_tokens, self.cfg.max_output_tokens);
|
||||
tracing::warn!(
|
||||
session_id = self.session_id,
|
||||
reason = "preflight",
|
||||
handoff_attempts = *handoff_attempts,
|
||||
max_handoffs = self.cfg.max_handoffs,
|
||||
projected_tokens = projected,
|
||||
threshold_tokens = threshold,
|
||||
"handoff cap reached; using truncation",
|
||||
);
|
||||
return HandoffOutcome::Skipped;
|
||||
}
|
||||
// Consume one attempt slot before calling handoff(). This ensures
|
||||
// that empty-summary, summarize-error, and cancellation outcomes all
|
||||
// burn budget — not just successful compactions — so the cap cannot
|
||||
// be bypassed by a flaky summarizer.
|
||||
*handoff_attempts += 1;
|
||||
self.handoff(None).await
|
||||
}
|
||||
|
||||
/// Handoff forced by a provider context-window rejection, bypassing both
|
||||
/// gates in [`Self::maybe_handoff`].
|
||||
///
|
||||
/// The gates exist to *predict* overflow; a 400 naming a context-length
|
||||
/// overflow is overflow already observed, so neither prediction applies.
|
||||
/// `should_handoff()` reads a token count frozen at the last SUCCESSFUL
|
||||
/// request (a failed request reports no usage), so it is under threshold by
|
||||
/// construction — that frozen reading is the permanent stick. And
|
||||
/// `max_handoffs` is a cost cap whose only alternative here is a request
|
||||
/// that cannot succeed.
|
||||
///
|
||||
/// `history_budget_bytes` is explicit rather than derived from
|
||||
/// `cfg.max_context_tokens`: that window is the quantity the provider just
|
||||
/// contradicted, so the recovery ladder must not be computed from it.
|
||||
pub(crate) async fn forced_handoff(&mut self, history_budget_bytes: usize) -> HandoffOutcome {
|
||||
tracing::warn!(
|
||||
"provider reported context overflow; forcing handoff (history budget {history_budget_bytes} bytes)"
|
||||
);
|
||||
self.handoff(Some(history_budget_bytes)).await
|
||||
}
|
||||
|
||||
/// The reactive context-recovery ladder, run after the provider rejected a
|
||||
/// request with a context-window 400.
|
||||
///
|
||||
/// `attempts` is the caller's per-`run()` recovery counter, advanced here as
|
||||
/// rungs are consumed. The caller owns it so the budget spans every
|
||||
/// context-400 in the turn, not just the rungs of one ladder.
|
||||
///
|
||||
/// The shrink schedule is anchored on the history that was just *observed*
|
||||
/// to be too large, halving from there — not on `cfg.max_context_tokens`,
|
||||
/// which the provider just contradicted and which may be overstated by an
|
||||
/// unknown factor. Halving needs no calibration: by the third rung it is at
|
||||
/// 1/8 of the rejected size.
|
||||
///
|
||||
/// Loops rather than returning after one rung because the summarize call
|
||||
/// travels the same provider path and can be rejected for the same reason.
|
||||
/// Treating that as unrecoverable would reproduce the very stick this fixes:
|
||||
/// the next rung halves the summarizer's own prompt, which is the only way
|
||||
/// out.
|
||||
///
|
||||
/// Gives up when the next budget would fall below
|
||||
/// [`HANDOFF_MIN_PROMPT_BUDGET_BYTES`]. That can happen on the FIRST rung
|
||||
/// when history is already small — correct, not premature: if a few KiB of
|
||||
/// history still overflows the window, the overflow is dominated by what a
|
||||
/// handoff cannot shrink (system prompt, tool schemas, the live user
|
||||
/// prompt), so further halving would only issue smaller doomed requests in
|
||||
/// place of a clear error.
|
||||
pub(crate) async fn recover_from_context_overflow(
|
||||
&mut self,
|
||||
attempts: &mut u32,
|
||||
) -> ContextRecovery {
|
||||
let rejected_bytes: usize = self
|
||||
.history
|
||||
.iter()
|
||||
.map(HistoryItem::context_pressure_bytes)
|
||||
.sum();
|
||||
loop {
|
||||
if *attempts >= MAX_CONTEXT_RECOVERIES_PER_RUN {
|
||||
tracing::error!(
|
||||
"context recovery budget spent ({MAX_CONTEXT_RECOVERIES_PER_RUN} attempts this turn); surfacing provider error"
|
||||
);
|
||||
return ContextRecovery::Exhausted;
|
||||
}
|
||||
// Shift by `attempts + 1`: the first rung already halves, since
|
||||
// rebuilding the rejected size would just fail again.
|
||||
let shift = (*attempts + 1).min(usize::BITS - 1);
|
||||
let budget = rejected_bytes >> shift;
|
||||
*attempts += 1;
|
||||
if budget < HANDOFF_MIN_PROMPT_BUDGET_BYTES {
|
||||
tracing::error!(
|
||||
"context recovery would shrink the handoff prompt to {budget} bytes, below \
|
||||
the {HANDOFF_MIN_PROMPT_BUDGET_BYTES}-byte floor (history {rejected_bytes} \
|
||||
bytes); surfacing provider error"
|
||||
);
|
||||
return ContextRecovery::Exhausted;
|
||||
}
|
||||
match self.forced_handoff(budget).await {
|
||||
HandoffOutcome::Performed => return ContextRecovery::Recovered,
|
||||
HandoffOutcome::Cancelled => return ContextRecovery::Cancelled,
|
||||
// Summarizer errored or returned nothing — possibly because its
|
||||
// own prompt overflowed. Truncation is not a usable fallback
|
||||
// (it sizes against the request-body budget, not context
|
||||
// pressure), so take the next rung with a smaller prompt.
|
||||
HandoffOutcome::Skipped => {
|
||||
tracing::warn!(
|
||||
"forced handoff at {budget} bytes did not run; shrinking further"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The handoff mechanism itself: summarize, reset, re-seat the live prompt.
|
||||
/// Holds no gate — callers decide whether a handoff is warranted.
|
||||
async fn handoff(&mut self, history_budget_bytes: Option<usize>) -> HandoffOutcome {
|
||||
let prompt = self.build_handoff_prompt(history_budget_bytes);
|
||||
let tokens_before = self.projected_handoff_input_tokens();
|
||||
let summary = tokio::select! {
|
||||
biased;
|
||||
_ = self.cancel.changed() => return HandoffOutcome::Cancelled,
|
||||
r = self.llm.summarize(
|
||||
self.cfg,
|
||||
&HANDOFF_SYSTEM_PROMPT,
|
||||
&prompt,
|
||||
HANDOFF_MAX_OUTPUT_TOKENS,
|
||||
self.effective_model,
|
||||
) => match r {
|
||||
Ok(s) if !s.trim().is_empty() => s,
|
||||
Ok(_) => {
|
||||
tracing::warn!("handoff returned empty summary; truncating");
|
||||
return HandoffOutcome::Skipped;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("handoff failed: {e}; truncating");
|
||||
return HandoffOutcome::Skipped;
|
||||
}
|
||||
},
|
||||
};
|
||||
let current_prompt = self.history.iter().rev().find_map(|item| match item {
|
||||
HistoryItem::User(s) => Some(s.clone()),
|
||||
_ => None,
|
||||
});
|
||||
let prior = self.history.len();
|
||||
// Reset history first; the _PostCompact hook is meant to inject
|
||||
// state into the FRESH context, not the old one we're discarding.
|
||||
self.history.clear();
|
||||
let post_compact = self
|
||||
.mcp
|
||||
.call_hooks(
|
||||
"_PostCompact",
|
||||
&serde_json::json!({}),
|
||||
self.cfg.hook_timeout,
|
||||
&self.cfg.hook_servers,
|
||||
)
|
||||
.await;
|
||||
// Handoff summary and hook output are injected as a synthetic user
|
||||
// message in one block. This keeps `_PostCompact` untrusted while also
|
||||
// avoiding orphan tool-result messages in the fresh context: OpenAI
|
||||
// Chat/Responses require tool outputs to follow an assistant tool call,
|
||||
// but handoff reset intentionally discards the old assistant turn.
|
||||
let mut handoff_text = format!("[Context Handoff]\n{summary}");
|
||||
if !post_compact.is_empty() {
|
||||
handoff_text.push_str("\n\n[Post-compact hook output — untrusted]\n");
|
||||
handoff_text.push_str(&hook_outputs_text(&post_compact));
|
||||
}
|
||||
self.history.push(HistoryItem::User(handoff_text));
|
||||
if let Some(prompt) = current_prompt {
|
||||
self.history.push(HistoryItem::User(prompt));
|
||||
}
|
||||
*self.handoff_count += 1;
|
||||
let token_counts = HandoffTokenCounts {
|
||||
before: tokens_before,
|
||||
after: estimate_history_tokens(self.history),
|
||||
};
|
||||
tracing::info!(
|
||||
"handoff #{} (history {prior} -> {} items; {token_counts})",
|
||||
*self.handoff_count,
|
||||
self.history.len()
|
||||
);
|
||||
HandoffOutcome::Performed
|
||||
}
|
||||
|
||||
fn should_handoff(&self) -> bool {
|
||||
match *self.last_request_input_tokens {
|
||||
Some(_) => {
|
||||
self.projected_handoff_input_tokens()
|
||||
>= token_threshold(self.cfg.max_context_tokens, self.cfg.max_output_tokens)
|
||||
}
|
||||
None => {
|
||||
let bytes: usize = self
|
||||
.history
|
||||
.iter()
|
||||
.map(HistoryItem::context_pressure_bytes)
|
||||
.sum();
|
||||
bytes
|
||||
> byte_fallback_threshold(
|
||||
self.cfg.max_context_tokens,
|
||||
self.cfg.max_output_tokens,
|
||||
self.cfg.max_history_bytes,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn projected_handoff_input_tokens(&self) -> u64 {
|
||||
let current_tokens = estimate_history_tokens(self.history);
|
||||
match *self.last_request_input_tokens {
|
||||
// Token-first: the provider told us exactly how many input tokens
|
||||
// the PREVIOUS request used. But history has grown since that
|
||||
// measurement — new assistant text, tool results, and the next
|
||||
// user prompt are appended before the next `complete()`. The exact
|
||||
// count alone would miss "previous request was under threshold, but
|
||||
// newly appended content pushes the next one over" (the stale-usage
|
||||
// cousin of the original stale-bytes bug). So we add a conservative
|
||||
// token estimate of the bytes added since the measurement.
|
||||
Some(measured_tokens) => {
|
||||
let measured_bytes = self.last_request_history_bytes.unwrap_or(0);
|
||||
let current_bytes: usize = self
|
||||
.history
|
||||
.iter()
|
||||
.map(HistoryItem::context_pressure_bytes)
|
||||
.sum();
|
||||
let grown = current_bytes.saturating_sub(measured_bytes);
|
||||
measured_tokens.saturating_add(estimate_tokens_from_bytes(grown))
|
||||
}
|
||||
// No usage yet (first request, or just after a handoff reset).
|
||||
// Fall back to the byte heuristic, capped conservatively so a
|
||||
// single pre-usage request can't blow the window. We map the token
|
||||
// threshold to bytes using a deliberately LOW bytes/token ratio:
|
||||
// a low ratio implies more tokens per byte, so the byte cap is
|
||||
// small and the handoff fires early rather than late. Never raise
|
||||
// the cap above the configured byte budget.
|
||||
//
|
||||
// Caveat: this can't shrink a single oversized current prompt,
|
||||
// since a handoff re-adds the current prompt verbatim — that is a
|
||||
// prompt-cap concern (MAX_PROMPT_BYTES), not this gate.
|
||||
None => current_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the summarizer prompt. `history_budget_bytes` overrides the
|
||||
/// budget normally derived from `cfg.max_context_tokens`; `None` keeps the
|
||||
/// derived value, which is what the proactive path uses.
|
||||
fn build_handoff_prompt(&self, history_budget_bytes: Option<usize>) -> String {
|
||||
let mut head = String::new();
|
||||
head.push_str(&format!(
|
||||
"[Internal handoff #{} — context reset]\n\n",
|
||||
*self.handoff_count + 1
|
||||
));
|
||||
head.push_str("# Original Task\n");
|
||||
let task = self.original_task.as_deref().unwrap_or("(unknown)");
|
||||
head.push_str(&clamp_bytes(task, HANDOFF_ORIGINAL_TASK_MAX_BYTES));
|
||||
head.push_str("\n\n# Available Tools\n");
|
||||
let all_tools = self.mcp.tools();
|
||||
let total = all_tools.len();
|
||||
if total == 0 {
|
||||
head.push_str("(none)\n");
|
||||
} else {
|
||||
let shown = total.min(HANDOFF_MAX_TOOL_NAMES);
|
||||
let names: Vec<&str> = all_tools[..shown].iter().map(|t| t.name.as_str()).collect();
|
||||
head.push_str(&names.join(", "));
|
||||
if shown < total {
|
||||
head.push_str(&format!(", … (+{} more)", total - shown));
|
||||
}
|
||||
head.push('\n');
|
||||
}
|
||||
let tail = "\n# Instructions\n\
|
||||
Produce a context handoff summary covering: (1) original task, \
|
||||
(2) what was accomplished, (3) key decisions, (4) what remains, \
|
||||
(5) one concrete next step. Be concise but thorough. Plain text.\n";
|
||||
let history_header = "\n# Session History (oldest first)\n";
|
||||
let fixed_bytes = head.len() + history_header.len() + tail.len();
|
||||
// An explicit budget is the allowance for the whole prompt, so subtract
|
||||
// the fixed frame from it exactly as the derived path does — otherwise
|
||||
// a caller's ceiling would be silently exceeded by the frame. When the
|
||||
// frame alone is larger than the budget, history drops to zero and the
|
||||
// frame is what remains: it is already independently clamped
|
||||
// (`HANDOFF_ORIGINAL_TASK_MAX_BYTES`, `HANDOFF_MAX_TOOL_NAMES`) and is
|
||||
// not reducible from here.
|
||||
let prompt_budget = match history_budget_bytes {
|
||||
Some(explicit) => explicit.saturating_sub(fixed_bytes),
|
||||
None => handoff_prompt_budget_bytes(
|
||||
self.cfg.max_context_tokens,
|
||||
summary_completion_cap(self.cfg.provider, HANDOFF_MAX_OUTPUT_TOKENS),
|
||||
fixed_bytes,
|
||||
),
|
||||
};
|
||||
|
||||
let mut snippets: Vec<String> = Vec::new();
|
||||
let mut snippets_bytes = 0usize;
|
||||
let mut dropped = 0usize;
|
||||
for item in self.history.iter().rev() {
|
||||
let mut snippet = String::new();
|
||||
push_history_snippet(&mut snippet, item);
|
||||
let snippet_bytes = snippet.len();
|
||||
if snippets_bytes.saturating_add(snippet_bytes) > prompt_budget {
|
||||
if snippets.is_empty() {
|
||||
snippets.push(clamp_bytes(&snippet, prompt_budget));
|
||||
snippets_bytes = prompt_budget;
|
||||
}
|
||||
dropped += 1;
|
||||
continue;
|
||||
}
|
||||
snippets_bytes += snippet_bytes;
|
||||
snippets.push(snippet);
|
||||
}
|
||||
snippets.reverse();
|
||||
if dropped > 0 {
|
||||
tracing::info!(
|
||||
"handoff prompt budget, dropped {dropped} oldest snippets; kept {} bytes",
|
||||
snippets_bytes
|
||||
);
|
||||
}
|
||||
|
||||
let mut out = String::with_capacity(
|
||||
head.len()
|
||||
+ history_header.len()
|
||||
+ tail.len()
|
||||
+ snippets_bytes
|
||||
+ if dropped > 0 { 32 } else { 0 },
|
||||
);
|
||||
out.push_str(&head);
|
||||
out.push_str(history_header);
|
||||
if dropped > 0 {
|
||||
out.push_str(&format!("(… {dropped} older items omitted)\n"));
|
||||
}
|
||||
for s in &snippets {
|
||||
out.push_str(s);
|
||||
}
|
||||
out.push_str(tail);
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
fn hook_outputs_text(outputs: &[(String, String)]) -> String {
|
||||
outputs
|
||||
.iter()
|
||||
.map(|(name, text)| format!("[{name}]\n{text}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n")
|
||||
}
|
||||
|
||||
fn push_history_snippet(out: &mut String, item: &HistoryItem) {
|
||||
match item {
|
||||
HistoryItem::User(s) => {
|
||||
out.push_str("[user] ");
|
||||
out.push_str(s);
|
||||
out.push('\n');
|
||||
}
|
||||
HistoryItem::Assistant {
|
||||
text,
|
||||
tool_calls,
|
||||
reasoning_details: _,
|
||||
} => {
|
||||
out.push_str("[assistant] ");
|
||||
if !text.is_empty() {
|
||||
out.push_str(text);
|
||||
}
|
||||
for c in tool_calls {
|
||||
out.push_str(&format!(" tool:{}", c.name));
|
||||
}
|
||||
out.push('\n');
|
||||
}
|
||||
HistoryItem::ToolResult(r) => {
|
||||
out.push_str(if r.is_error { "[tool_err] " } else { "[tool] " });
|
||||
out.push_str(&r.text());
|
||||
out.push('\n');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Byte budget for session-history text inside the handoff prompt. The
|
||||
/// summarizer uses the same provider/model config as normal completion, so
|
||||
/// derive the input budget from the model context window instead of applying a
|
||||
/// separate fixed prompt cap. We keep the same 1 byte/token upper-bound
|
||||
/// estimate used by the handoff gate, which is conservative: it may drop old
|
||||
/// history early for unusually large sessions, but it should not build a prompt
|
||||
/// that exceeds the configured context window.
|
||||
fn handoff_prompt_budget_bytes(
|
||||
max_context_tokens: u64,
|
||||
max_output_tokens: u32,
|
||||
fixed_prompt_bytes: usize,
|
||||
) -> usize {
|
||||
max_context_tokens
|
||||
.saturating_sub(u64::from(max_output_tokens))
|
||||
.saturating_mul(CONSERVATIVE_BYTES_PER_TOKEN)
|
||||
.saturating_sub(u64::try_from(fixed_prompt_bytes).unwrap_or(u64::MAX))
|
||||
.try_into()
|
||||
.unwrap_or(usize::MAX)
|
||||
}
|
||||
|
||||
pub(crate) fn clamp_bytes(s: &str, max_bytes: usize) -> String {
|
||||
if s.len() <= max_bytes {
|
||||
return s.to_owned();
|
||||
}
|
||||
if max_bytes < 4 {
|
||||
let mut cut = max_bytes.min(s.len());
|
||||
while cut > 0 && !s.is_char_boundary(cut) {
|
||||
cut -= 1;
|
||||
}
|
||||
return s[..cut].to_owned();
|
||||
}
|
||||
let target = max_bytes - "…".len();
|
||||
let mut cut = target;
|
||||
while cut > 0 && !s.is_char_boundary(cut) {
|
||||
cut -= 1;
|
||||
}
|
||||
format!("{}…", &s[..cut])
|
||||
}
|
||||
|
||||
/// Conservative bytes-per-token ratio used when estimating tokens from raw
|
||||
/// history bytes. We use 1: a token is always at least one byte, so treating
|
||||
/// every byte as a whole token is an unconditional UPPER bound on the true
|
||||
/// token count — it can never undercount, regardless of content density (even
|
||||
/// the densest real content sits at ~1.4 bytes/token). That over-estimate is
|
||||
/// exactly what a fail-early preflight gate wants: it hands off sooner rather
|
||||
/// than risk the next request exceeding the window.
|
||||
const CONSERVATIVE_BYTES_PER_TOKEN: u64 = 1;
|
||||
|
||||
fn estimate_history_tokens(history: &[HistoryItem]) -> u64 {
|
||||
estimate_tokens_from_bytes(
|
||||
history
|
||||
.iter()
|
||||
.map(HistoryItem::context_pressure_bytes)
|
||||
.sum(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Estimate tokens from a byte count at the conservative ratio (rounding up,
|
||||
/// so a partial token still counts). At a 1:1 ratio this is just the byte
|
||||
/// count — a guaranteed upper bound on tokens.
|
||||
fn estimate_tokens_from_bytes(bytes: usize) -> u64 {
|
||||
(bytes as u64).div_ceil(CONSERVATIVE_BYTES_PER_TOKEN)
|
||||
}
|
||||
|
||||
/// Input-token count at which to hand off. Caps at the configured fraction of
|
||||
/// the window and also leaves room for `max_output_tokens`, so input + output
|
||||
/// can't together exceed the window. Free function so the policy math is unit
|
||||
/// testable without constructing a [`RunCtx`].
|
||||
fn token_threshold(max_context_tokens: u64, max_output_tokens: u32) -> u64 {
|
||||
// Integer math: handoff threshold is 90%, i.e. window * 9 / 10.
|
||||
let fractional = max_context_tokens / 10 * 9;
|
||||
let output_reserved = max_context_tokens.saturating_sub(u64::from(max_output_tokens));
|
||||
fractional.min(output_reserved)
|
||||
}
|
||||
|
||||
/// Conservative byte cap used only before any usage is known. Maps the token
|
||||
/// threshold to bytes at the conservative bytes/token ratio (so the cap is
|
||||
/// small and the handoff fires early), clamped to the configured byte budget
|
||||
/// so it can only ever be more conservative than the old byte-only behavior.
|
||||
fn byte_fallback_threshold(
|
||||
max_context_tokens: u64,
|
||||
max_output_tokens: u32,
|
||||
max_history_bytes: usize,
|
||||
) -> usize {
|
||||
let derived = token_threshold(max_context_tokens, max_output_tokens)
|
||||
.saturating_mul(CONSERVATIVE_BYTES_PER_TOKEN);
|
||||
let byte_cap = max_history_bytes / 10 * 9;
|
||||
usize::try_from(derived).unwrap_or(usize::MAX).min(byte_cap)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
byte_fallback_threshold, estimate_tokens_from_bytes, handoff_prompt_budget_bytes,
|
||||
summary_completion_cap, token_threshold, HANDOFF_SYSTEM_PROMPT,
|
||||
};
|
||||
use crate::config::{Provider, HANDOFF_MAX_OUTPUT_TOKENS};
|
||||
|
||||
#[test]
|
||||
fn handoff_prompt_budget_reserves_summary_output_and_fixed_prompt() {
|
||||
assert_eq!(handoff_prompt_budget_bytes(25_000, 8_192, 1_000), 15_808);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handoff_prompt_budget_saturates_when_fixed_prompt_exceeds_window() {
|
||||
assert_eq!(handoff_prompt_budget_bytes(1_000, 2_000, 10_000), 0);
|
||||
}
|
||||
|
||||
/// OpenRouter's summary request grants reasoning an equal budget on top of
|
||||
/// the visible-text budget, so its completion cap is 2× the handoff text
|
||||
/// budget; the input budget must reserve that doubled cap. At the
|
||||
/// 1-byte/token upper bound, prompt bytes bound prompt tokens, so the join
|
||||
/// to pin is: (budget + fixed prompt) + actual completion cap ≤ window.
|
||||
/// Reserving only `HANDOFF_MAX_OUTPUT_TOKENS` would break this by exactly
|
||||
/// one extra reasoning budget at the maximum constructed prompt.
|
||||
#[test]
|
||||
fn openrouter_prompt_budget_reserves_doubled_completion_cap() {
|
||||
let cap = summary_completion_cap(Provider::OpenRouter, HANDOFF_MAX_OUTPUT_TOKENS);
|
||||
assert_eq!(
|
||||
cap,
|
||||
2 * HANDOFF_MAX_OUTPUT_TOKENS,
|
||||
"OpenRouter doubles: text + reasoning"
|
||||
);
|
||||
let window = 200_000u64;
|
||||
let fixed = 1_000usize;
|
||||
let budget = handoff_prompt_budget_bytes(window, cap, fixed);
|
||||
assert_eq!(budget, 182_616); // 200_000 - 16_384 - 1_000
|
||||
let max_prompt_tokens = estimate_tokens_from_bytes(budget + fixed);
|
||||
assert!(
|
||||
max_prompt_tokens + u64::from(cap) <= window,
|
||||
"input + completion allowance must fit the configured window"
|
||||
);
|
||||
// The old single reservation violates the same join — the regression
|
||||
// this guards against.
|
||||
let stale_budget = handoff_prompt_budget_bytes(window, HANDOFF_MAX_OUTPUT_TOKENS, fixed);
|
||||
assert!(
|
||||
estimate_tokens_from_bytes(stale_budget + fixed) + u64::from(cap) > window,
|
||||
"reserving only the text budget must be observable as an overflow here"
|
||||
);
|
||||
}
|
||||
|
||||
/// Anthropic/OpenAI/Databricks summary bodies request exactly the caller's
|
||||
/// budget, so their input reservation is unchanged.
|
||||
#[test]
|
||||
fn non_openrouter_completion_cap_is_the_callers_budget() {
|
||||
for provider in [
|
||||
Provider::Anthropic,
|
||||
Provider::OpenAi,
|
||||
Provider::Databricks,
|
||||
Provider::DatabricksV2,
|
||||
] {
|
||||
assert_eq!(
|
||||
summary_completion_cap(provider, HANDOFF_MAX_OUTPUT_TOKENS),
|
||||
HANDOFF_MAX_OUTPUT_TOKENS
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// The prompt's token figure is derived from `HANDOFF_MAX_OUTPUT_TOKENS`
|
||||
/// and names the *visible plain-text summary* as its target, so hidden
|
||||
/// reasoning (budgeted separately on the wire) is not the referent.
|
||||
#[test]
|
||||
fn handoff_system_prompt_derives_limit_and_targets_visible_text() {
|
||||
let expected = format!(
|
||||
"Keep the visible plain-text summary under {HANDOFF_MAX_OUTPUT_TOKENS} tokens."
|
||||
);
|
||||
assert!(
|
||||
HANDOFF_SYSTEM_PROMPT.contains(&expected),
|
||||
"prompt must derive its token figure from HANDOFF_MAX_OUTPUT_TOKENS: {}",
|
||||
*HANDOFF_SYSTEM_PROMPT
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_threshold_uses_fraction_when_output_is_small() {
|
||||
// 200k window, 1k output. fractional = 0.9*200000 = 180000;
|
||||
// output_reserved = 200000-1000 = 199000; min = 180000.
|
||||
assert_eq!(token_threshold(200_000, 1_000), 180_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_threshold_reserves_output_headroom() {
|
||||
// Large output relative to window: the output-reserve term dominates,
|
||||
// keeping input+output within the window.
|
||||
// 100k window, 40k output: fractional=90k, reserved=60k -> 60k.
|
||||
assert_eq!(token_threshold(100_000, 40_000), 60_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_threshold_saturates_when_output_exceeds_window() {
|
||||
// Degenerate (config validation forbids this, but math must not panic):
|
||||
// reserved saturates to 0, so threshold is 0 -> always hand off.
|
||||
assert_eq!(token_threshold(1000, 5000), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn byte_fallback_is_conservative_and_capped() {
|
||||
// Derived = token_threshold * 1 (1 byte/token upper bound). For
|
||||
// 200k/1k: 180000 bytes, well under a 16 MiB byte budget, so derived
|
||||
// wins (early handoff).
|
||||
let t = byte_fallback_threshold(200_000, 1_000, 16 * 1024 * 1024);
|
||||
assert_eq!(t, 180_000);
|
||||
// With a tiny byte budget the cap wins -> never exceeds it (window*90%).
|
||||
let capped = byte_fallback_threshold(200_000, 1_000, 8192);
|
||||
assert_eq!(capped, 8192 / 10 * 9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_tokens_is_upper_bound_on_tokens() {
|
||||
// 1 byte/token: a token is always >= 1 byte, so byte count is an
|
||||
// unconditional upper bound on the true token count.
|
||||
assert_eq!(estimate_tokens_from_bytes(0), 0);
|
||||
assert_eq!(estimate_tokens_from_bytes(1), 1);
|
||||
assert_eq!(estimate_tokens_from_bytes(4), 4);
|
||||
assert_eq!(estimate_tokens_from_bytes(5), 5);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,726 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use crate::mcp::truncate_at_boundary;
|
||||
|
||||
const MAX_HINTS_BYTES: usize = 128 * 1024;
|
||||
pub const MAX_SKILL_BODY_BYTES: usize = 32 * 1024;
|
||||
const SKILL_DIRS: &[&str] = &[".agents/skills", ".goose/skills", ".claude/skills"];
|
||||
|
||||
fn home_dir() -> Option<PathBuf> {
|
||||
std::env::var("HOME").ok().map(PathBuf::from)
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct SkillEntry {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
/// Absolute path to the SKILL.md file; used by `load_skill` to read on demand.
|
||||
pub path: PathBuf,
|
||||
/// Absolute paths to every non-SKILL.md file in the skill directory tree.
|
||||
/// Pre-enumerated at discovery time so `load_skill` can match by relative path
|
||||
/// without doing arbitrary filesystem lookups at call time.
|
||||
pub supporting_files: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
/// Handles both normal repos (`.git/` dir) and worktrees (`.git` file).
|
||||
fn find_git_root(start: &Path) -> Option<PathBuf> {
|
||||
let mut current = start.to_path_buf();
|
||||
loop {
|
||||
if current.join(".git").exists() {
|
||||
return Some(current);
|
||||
}
|
||||
match current.parent() {
|
||||
Some(parent) => current = parent.to_path_buf(),
|
||||
None => return None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn load_hint_files_impl(cwd: &Path, home: Option<&Path>) -> String {
|
||||
let mut chain = match find_git_root(cwd) {
|
||||
Some(root) => {
|
||||
let mut c: Vec<PathBuf> = cwd
|
||||
.ancestors()
|
||||
.take_while(|a| a.starts_with(&root))
|
||||
.map(|a| a.to_path_buf())
|
||||
.collect();
|
||||
// ancestors() yields cwd first, root last — reverse for root→cwd.
|
||||
c.reverse();
|
||||
c
|
||||
}
|
||||
None => vec![cwd.to_path_buf()],
|
||||
};
|
||||
|
||||
// Prepend ~/AGENTS.md as global layer, unless ~ is already in the chain.
|
||||
if let Some(home) = home {
|
||||
if !chain.iter().any(|d| d == home) {
|
||||
chain.insert(0, home.to_path_buf());
|
||||
}
|
||||
}
|
||||
|
||||
let mut result = String::new();
|
||||
for dir in &chain {
|
||||
let path = dir.join("AGENTS.md");
|
||||
let Ok(content) = std::fs::read_to_string(&path) else {
|
||||
continue;
|
||||
};
|
||||
if !result.is_empty() {
|
||||
result.push_str("\n\n");
|
||||
}
|
||||
let remaining = MAX_HINTS_BYTES.saturating_sub(result.len());
|
||||
if remaining == 0 {
|
||||
break;
|
||||
}
|
||||
if content.len() <= remaining {
|
||||
result.push_str(&content);
|
||||
} else {
|
||||
let truncated = truncate_at_boundary(&content, remaining);
|
||||
result.push_str(truncated);
|
||||
break;
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn parse_skill_frontmatter(content: &str) -> Option<(String, String)> {
|
||||
// Must start with `---`
|
||||
let rest = content.strip_prefix("---\n")?;
|
||||
// Find the closing `---`
|
||||
let close_pos = rest.find("\n---")?;
|
||||
let yaml_block = &rest[..close_pos];
|
||||
|
||||
let map: HashMap<String, serde_yaml::Value> = serde_yaml::from_str(yaml_block).ok()?;
|
||||
let name = map
|
||||
.get("name")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(str::to_string)?;
|
||||
let description = map
|
||||
.get("description")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(str::trim)
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
Some((name, description))
|
||||
}
|
||||
|
||||
fn scan_skill_dir(dir: &Path, seen: &mut HashSet<String>, skills: &mut Vec<SkillEntry>) {
|
||||
let Ok(entries) = std::fs::read_dir(dir) else {
|
||||
return;
|
||||
};
|
||||
let mut subdirs: Vec<PathBuf> = entries
|
||||
.filter_map(|e| e.ok())
|
||||
// Use std::fs::metadata (follows symlinks) rather than DirEntry::file_type
|
||||
// (which returns FileType::Symlink for symlinks, causing is_dir() to return
|
||||
// false even when the symlink target is a directory).
|
||||
.filter(|e| {
|
||||
std::fs::metadata(e.path())
|
||||
.map(|m| m.is_dir())
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.map(|e| e.path())
|
||||
.collect();
|
||||
subdirs.sort();
|
||||
|
||||
for subdir in subdirs {
|
||||
let skill_md = subdir.join("SKILL.md");
|
||||
let Ok(content) = std::fs::read_to_string(&skill_md) else {
|
||||
continue;
|
||||
};
|
||||
let Some((name, description)) = parse_skill_frontmatter(&content) else {
|
||||
continue;
|
||||
};
|
||||
if seen.contains(&name) {
|
||||
continue;
|
||||
}
|
||||
seen.insert(name.clone());
|
||||
|
||||
// Collect supporting files: every non-SKILL.md file in the skill dir tree.
|
||||
// Don't descend into subdirs that themselves have a SKILL.md — those are
|
||||
// separate skills with their own entries.
|
||||
let supporting_files = collect_supporting_files(&subdir);
|
||||
|
||||
skills.push(SkillEntry {
|
||||
name,
|
||||
description,
|
||||
path: skill_md,
|
||||
supporting_files,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Walk `skill_dir` recursively and return the absolute path of every file
|
||||
/// that is not `SKILL.md`. Subdirectories that contain their own `SKILL.md`
|
||||
/// are treated as separate skills and are not descended into.
|
||||
fn collect_supporting_files(skill_dir: &Path) -> Vec<PathBuf> {
|
||||
let mut result = Vec::new();
|
||||
let mut visited_dirs = HashSet::new();
|
||||
collect_supporting_files_impl(skill_dir, &mut result, &mut visited_dirs);
|
||||
result.sort();
|
||||
result
|
||||
}
|
||||
|
||||
fn collect_supporting_files_impl(
|
||||
current: &Path,
|
||||
out: &mut Vec<PathBuf>,
|
||||
visited_dirs: &mut HashSet<PathBuf>,
|
||||
) {
|
||||
let Ok(canonical_current) = current.canonicalize() else {
|
||||
return;
|
||||
};
|
||||
if !visited_dirs.insert(canonical_current) {
|
||||
return;
|
||||
}
|
||||
|
||||
let Ok(entries) = std::fs::read_dir(current) else {
|
||||
return;
|
||||
};
|
||||
let mut items: Vec<_> = entries.filter_map(|e| e.ok()).collect();
|
||||
items.sort_by_key(|e| e.path());
|
||||
|
||||
for entry in items {
|
||||
let path = entry.path();
|
||||
// Use std::fs::metadata (follows symlinks) so symlinked subdirs and files
|
||||
// inside a skill directory are handled correctly.
|
||||
let ft = match std::fs::metadata(&path) {
|
||||
Ok(m) => m,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if ft.is_dir() {
|
||||
// Don't descend into subdirs that are themselves skills.
|
||||
if path.join("SKILL.md").is_file() {
|
||||
continue;
|
||||
}
|
||||
collect_supporting_files_impl(&path, out, visited_dirs);
|
||||
} else if ft.is_file() && path.file_name().and_then(|n| n.to_str()) != Some("SKILL.md") {
|
||||
out.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn discover_skills_impl(cwd: &Path, home: Option<&Path>) -> Vec<SkillEntry> {
|
||||
let mut seen = HashSet::new();
|
||||
let mut skills = Vec::new();
|
||||
|
||||
for dir_suffix in SKILL_DIRS {
|
||||
scan_skill_dir(&cwd.join(dir_suffix), &mut seen, &mut skills);
|
||||
}
|
||||
|
||||
if let Some(home) = home {
|
||||
scan_skill_dir(&home.join(".agents/skills"), &mut seen, &mut skills);
|
||||
}
|
||||
|
||||
skills
|
||||
}
|
||||
|
||||
pub fn build_hints_section(cwd: &Path) -> (String, Vec<SkillEntry>) {
|
||||
build_hints_section_impl(cwd, home_dir().as_deref())
|
||||
}
|
||||
|
||||
fn build_hints_section_impl(cwd: &Path, home: Option<&Path>) -> (String, Vec<SkillEntry>) {
|
||||
let hints_text = load_hint_files_impl(cwd, home);
|
||||
let skills = discover_skills_impl(cwd, home);
|
||||
|
||||
if hints_text.is_empty() && skills.is_empty() {
|
||||
return (String::new(), skills);
|
||||
}
|
||||
|
||||
let mut out = String::from("# Additional Instructions\n");
|
||||
|
||||
if !hints_text.is_empty() {
|
||||
out.push_str("\n## Project Hints\n");
|
||||
out.push_str(&hints_text);
|
||||
out.push('\n');
|
||||
}
|
||||
|
||||
if !skills.is_empty() {
|
||||
out.push_str("\n## Available Skills\n");
|
||||
for skill in &skills {
|
||||
out.push_str(&format!("- {}: {}\n", skill.name, skill.description));
|
||||
}
|
||||
out.push_str(
|
||||
"\nUse the `load_skill` tool to read the full content of a skill before using it.\n",
|
||||
);
|
||||
}
|
||||
|
||||
(out, skills)
|
||||
}
|
||||
|
||||
/// Strip the YAML frontmatter block from a skill file's content and return
|
||||
/// the body. If no valid frontmatter is found, returns the content unchanged.
|
||||
pub(crate) fn strip_frontmatter(content: &str) -> &str {
|
||||
let Some(rest) = content.strip_prefix("---\n") else {
|
||||
return content;
|
||||
};
|
||||
let Some(close_pos) = rest.find("\n---") else {
|
||||
return content;
|
||||
};
|
||||
let after = &rest[close_pos + 4..]; // skip "\n---"
|
||||
after.strip_prefix('\n').unwrap_or(after)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn find_git_root_normal_repo() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let root = tmp.path();
|
||||
std::fs::create_dir(root.join(".git")).unwrap();
|
||||
assert_eq!(find_git_root(root), Some(root.to_path_buf()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn find_git_root_worktree() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let root = tmp.path();
|
||||
// .git as a file (worktree)
|
||||
std::fs::write(root.join(".git"), "gitdir: ../main/.git/worktrees/wt").unwrap();
|
||||
assert_eq!(find_git_root(root), Some(root.to_path_buf()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn find_git_root_none() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
// No .git anywhere under tmp
|
||||
let result = find_git_root(tmp.path());
|
||||
// In a CI environment the test itself may live inside a real git repo,
|
||||
// so only assert None when tmp is truly isolated (not a subpath of a git repo).
|
||||
// We verify by checking that any found root is NOT inside tmp.
|
||||
if let Some(found) = result {
|
||||
assert!(!found.starts_with(tmp.path()));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn find_git_root_from_subdirectory() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let root = tmp.path();
|
||||
std::fs::create_dir(root.join(".git")).unwrap();
|
||||
let deep = root.join("sub").join("deep");
|
||||
std::fs::create_dir_all(&deep).unwrap();
|
||||
assert_eq!(find_git_root(&deep), Some(root.to_path_buf()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_hint_files_single_at_cwd() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
// No .git → no git root discovery; only cwd is checked.
|
||||
std::fs::write(cwd.join("AGENTS.md"), "cwd hints").unwrap();
|
||||
let result = load_hint_files_impl(cwd, None);
|
||||
assert_eq!(result, "cwd hints");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_hint_files_git_root_and_cwd() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let root = tmp.path();
|
||||
std::fs::create_dir(root.join(".git")).unwrap();
|
||||
std::fs::write(root.join("AGENTS.md"), "root hints").unwrap();
|
||||
let sub = root.join("sub");
|
||||
std::fs::create_dir(&sub).unwrap();
|
||||
std::fs::write(sub.join("AGENTS.md"), "sub hints").unwrap();
|
||||
let result = load_hint_files_impl(&sub, None);
|
||||
// Root hints must come first.
|
||||
assert!(
|
||||
result.starts_with("root hints"),
|
||||
"expected root hints first, got: {result:?}"
|
||||
);
|
||||
assert!(result.contains("sub hints"), "missing sub hints");
|
||||
let root_pos = result.find("root hints").unwrap();
|
||||
let sub_pos = result.find("sub hints").unwrap();
|
||||
assert!(root_pos < sub_pos, "root hints should precede sub hints");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_hint_files_missing_files() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let result = load_hint_files_impl(tmp.path(), None);
|
||||
assert_eq!(result, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_skills_finds_across_dirs() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
|
||||
// Skill in .agents/skills/
|
||||
let agents_skill = cwd.join(".agents/skills/my-skill");
|
||||
std::fs::create_dir_all(&agents_skill).unwrap();
|
||||
std::fs::write(
|
||||
agents_skill.join("SKILL.md"),
|
||||
"---\nname: my-skill\ndescription: A skill\n---\nSkill body here.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Skill in .goose/skills/
|
||||
let goose_skill = cwd.join(".goose/skills/other-skill");
|
||||
std::fs::create_dir_all(&goose_skill).unwrap();
|
||||
std::fs::write(
|
||||
goose_skill.join("SKILL.md"),
|
||||
"---\nname: other-skill\ndescription: Another skill\n---\nOther body.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let skills = discover_skills_impl(cwd, None);
|
||||
assert_eq!(skills.len(), 2);
|
||||
let names: Vec<&str> = skills.iter().map(|s| s.name.as_str()).collect();
|
||||
assert!(names.contains(&"my-skill"), "missing my-skill");
|
||||
assert!(names.contains(&"other-skill"), "missing other-skill");
|
||||
// Paths should point to the SKILL.md files.
|
||||
for skill in &skills {
|
||||
assert!(skill.path.exists(), "path does not exist: {:?}", skill.path);
|
||||
assert!(skill.path.ends_with("SKILL.md"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_skills_dedup_by_name() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
|
||||
// Same name in .agents/skills/ (first) and .goose/skills/ (second)
|
||||
let agents_skill = cwd.join(".agents/skills/shared");
|
||||
std::fs::create_dir_all(&agents_skill).unwrap();
|
||||
std::fs::write(
|
||||
agents_skill.join("SKILL.md"),
|
||||
"---\nname: shared\ndescription: from agents\n---\nAgents body.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let goose_skill = cwd.join(".goose/skills/shared");
|
||||
std::fs::create_dir_all(&goose_skill).unwrap();
|
||||
std::fs::write(
|
||||
goose_skill.join("SKILL.md"),
|
||||
"---\nname: shared\ndescription: from goose\n---\nGoose body.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let skills = discover_skills_impl(cwd, None);
|
||||
assert_eq!(skills.len(), 1, "duplicate name should be deduplicated");
|
||||
assert_eq!(
|
||||
skills[0].description, "from agents",
|
||||
"first wins (.agents/)"
|
||||
);
|
||||
// Path should point to the .agents/ version (first wins).
|
||||
assert!(skills[0]
|
||||
.path
|
||||
.to_str()
|
||||
.unwrap()
|
||||
.contains(".agents/skills/shared"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_skills_skips_missing_name() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
|
||||
let skill_dir = cwd.join(".agents/skills/no-name");
|
||||
std::fs::create_dir_all(&skill_dir).unwrap();
|
||||
std::fs::write(
|
||||
skill_dir.join("SKILL.md"),
|
||||
"---\ndescription: No name here\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let skills = discover_skills_impl(cwd, None);
|
||||
assert!(skills.is_empty(), "entry without name should be skipped");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_hints_section_empty() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let (result, skills) = build_hints_section_impl(tmp.path(), None);
|
||||
assert_eq!(result, "");
|
||||
assert!(skills.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_hints_section_combined() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
|
||||
std::fs::write(cwd.join("AGENTS.md"), "Project-level hints.").unwrap();
|
||||
|
||||
let skill_dir = cwd.join(".agents/skills/buzz-cli");
|
||||
std::fs::create_dir_all(&skill_dir).unwrap();
|
||||
std::fs::write(
|
||||
skill_dir.join("SKILL.md"),
|
||||
"---\nname: buzz-cli\ndescription: CLI reference for Buzz managed agents\n---\nUse `buzz` to manage agents.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let (result, skills) = build_hints_section_impl(cwd, None);
|
||||
|
||||
assert!(
|
||||
result.contains("# Additional Instructions"),
|
||||
"missing header"
|
||||
);
|
||||
assert!(result.contains("## Project Hints"), "missing Project Hints");
|
||||
assert!(
|
||||
result.contains("Project-level hints."),
|
||||
"missing hints content"
|
||||
);
|
||||
assert!(
|
||||
result.contains("## Available Skills"),
|
||||
"missing Available Skills"
|
||||
);
|
||||
assert!(
|
||||
result.contains("buzz-cli: CLI reference for Buzz managed agents"),
|
||||
"missing skill bullet"
|
||||
);
|
||||
// Body must NOT be inlined — lazy loading only.
|
||||
assert!(
|
||||
!result.contains("Use `buzz` to manage agents."),
|
||||
"skill body must not be inlined in system prompt"
|
||||
);
|
||||
// The load_skill instruction must be present.
|
||||
assert!(
|
||||
result.contains("load_skill"),
|
||||
"missing load_skill instruction"
|
||||
);
|
||||
// The old ### heading format must not appear.
|
||||
assert!(
|
||||
!result.contains("### buzz-cli"),
|
||||
"skill body heading must not be inlined"
|
||||
);
|
||||
// The returned skills list should contain the discovered skill.
|
||||
assert_eq!(skills.len(), 1);
|
||||
assert_eq!(skills[0].name, "buzz-cli");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_hint_files_global_loaded_first() {
|
||||
let home = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
std::fs::write(home.path().join("AGENTS.md"), "global hints").unwrap();
|
||||
std::fs::write(cwd.path().join("AGENTS.md"), "local hints").unwrap();
|
||||
let result = load_hint_files_impl(cwd.path(), Some(home.path()));
|
||||
let global_pos = result.find("global hints").unwrap();
|
||||
let local_pos = result.find("local hints").unwrap();
|
||||
assert!(
|
||||
global_pos < local_pos,
|
||||
"global hints should precede local hints"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_hint_files_home_missing_agents_md() {
|
||||
let home = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
std::fs::write(cwd.path().join("AGENTS.md"), "local only").unwrap();
|
||||
let result = load_hint_files_impl(cwd.path(), Some(home.path()));
|
||||
assert_eq!(result, "local only");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_hint_files_no_home_dir() {
|
||||
let cwd = TempDir::new().unwrap();
|
||||
std::fs::write(cwd.path().join("AGENTS.md"), "local only").unwrap();
|
||||
let result = load_hint_files_impl(cwd.path(), None);
|
||||
assert_eq!(result, "local only");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_hint_files_dedup_when_home_in_chain() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let home = tmp.path();
|
||||
std::fs::write(home.join("AGENTS.md"), "single load").unwrap();
|
||||
let result = load_hint_files_impl(home, Some(home));
|
||||
assert_eq!(
|
||||
result.matches("single load").count(),
|
||||
1,
|
||||
"AGENTS.md should be loaded exactly once when CWD is home"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_hint_files_dedup_when_home_is_git_root() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let home = tmp.path();
|
||||
std::fs::create_dir(home.join(".git")).unwrap();
|
||||
std::fs::write(home.join("AGENTS.md"), "root+home hints").unwrap();
|
||||
let sub = home.join("sub");
|
||||
std::fs::create_dir(&sub).unwrap();
|
||||
let result = load_hint_files_impl(&sub, Some(home));
|
||||
assert_eq!(
|
||||
result.matches("root+home hints").count(),
|
||||
1,
|
||||
"AGENTS.md should be loaded once when home is git root"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_skills_global_skills_loaded() {
|
||||
let home = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let skill_dir = home.path().join(".agents/skills/global-skill");
|
||||
std::fs::create_dir_all(&skill_dir).unwrap();
|
||||
std::fs::write(
|
||||
skill_dir.join("SKILL.md"),
|
||||
"---\nname: global-skill\ndescription: A global skill\n---\nGlobal body.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let skills = discover_skills_impl(cwd.path(), Some(home.path()));
|
||||
assert_eq!(skills.len(), 1);
|
||||
assert_eq!(skills[0].name, "global-skill");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_skills_project_wins_over_global() {
|
||||
let home = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
|
||||
let project_skill = cwd.path().join(".agents/skills/shared");
|
||||
std::fs::create_dir_all(&project_skill).unwrap();
|
||||
std::fs::write(
|
||||
project_skill.join("SKILL.md"),
|
||||
"---\nname: shared\ndescription: from project\n---\nProject body.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let global_skill = home.path().join(".agents/skills/shared");
|
||||
std::fs::create_dir_all(&global_skill).unwrap();
|
||||
std::fs::write(
|
||||
global_skill.join("SKILL.md"),
|
||||
"---\nname: shared\ndescription: from global\n---\nGlobal body.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let skills = discover_skills_impl(cwd.path(), Some(home.path()));
|
||||
assert_eq!(skills.len(), 1, "duplicate name should be deduplicated");
|
||||
assert_eq!(
|
||||
skills[0].description, "from project",
|
||||
"project-level should win over global"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_skills_no_home_dir() {
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let skill_dir = cwd.path().join(".agents/skills/local");
|
||||
std::fs::create_dir_all(&skill_dir).unwrap();
|
||||
std::fs::write(
|
||||
skill_dir.join("SKILL.md"),
|
||||
"---\nname: local\ndescription: Local skill\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let skills = discover_skills_impl(cwd.path(), None);
|
||||
assert_eq!(skills.len(), 1);
|
||||
assert_eq!(skills[0].name, "local");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collect_supporting_files_finds_non_skill_md_files() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path();
|
||||
// SKILL.md should be excluded.
|
||||
std::fs::write(skill_dir.join("SKILL.md"), "---\nname: x\n---\n").unwrap();
|
||||
// A references subdir with files.
|
||||
let refs = skill_dir.join("references");
|
||||
std::fs::create_dir_all(&refs).unwrap();
|
||||
std::fs::write(refs.join("foo.md"), "foo").unwrap();
|
||||
std::fs::write(refs.join("bar.md"), "bar").unwrap();
|
||||
// A script at the top level.
|
||||
std::fs::write(skill_dir.join("setup.sh"), "#!/bin/sh").unwrap();
|
||||
|
||||
let files = collect_supporting_files(skill_dir);
|
||||
let names: Vec<String> = files
|
||||
.iter()
|
||||
.map(|p| p.file_name().unwrap().to_string_lossy().into_owned())
|
||||
.collect();
|
||||
assert!(
|
||||
names.contains(&"foo.md".to_owned()),
|
||||
"missing foo.md: {names:?}"
|
||||
);
|
||||
assert!(
|
||||
names.contains(&"bar.md".to_owned()),
|
||||
"missing bar.md: {names:?}"
|
||||
);
|
||||
assert!(
|
||||
names.contains(&"setup.sh".to_owned()),
|
||||
"missing setup.sh: {names:?}"
|
||||
);
|
||||
assert!(
|
||||
!names.contains(&"SKILL.md".to_owned()),
|
||||
"SKILL.md should be excluded: {names:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collect_supporting_files_does_not_descend_into_nested_skills() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path();
|
||||
std::fs::write(skill_dir.join("SKILL.md"), "---\nname: x\n---\n").unwrap();
|
||||
std::fs::write(skill_dir.join("helper.sh"), "#!/bin/sh").unwrap();
|
||||
|
||||
// A nested subdir that is itself a skill — should not be descended into.
|
||||
let nested = skill_dir.join("nested-skill");
|
||||
std::fs::create_dir_all(&nested).unwrap();
|
||||
std::fs::write(nested.join("SKILL.md"), "---\nname: nested\n---\n").unwrap();
|
||||
std::fs::write(nested.join("secret.md"), "should not appear").unwrap();
|
||||
|
||||
let files = collect_supporting_files(skill_dir);
|
||||
let names: Vec<String> = files
|
||||
.iter()
|
||||
.map(|p| p.file_name().unwrap().to_string_lossy().into_owned())
|
||||
.collect();
|
||||
assert!(
|
||||
names.contains(&"helper.sh".to_owned()),
|
||||
"missing helper.sh: {names:?}"
|
||||
);
|
||||
assert!(
|
||||
!names.contains(&"secret.md".to_owned()),
|
||||
"nested skill's files should not appear: {names:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn collect_supporting_files_skips_symlink_cycles() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let skill_dir = tmp.path();
|
||||
std::fs::write(skill_dir.join("SKILL.md"), "---\nname: x\n---\n").unwrap();
|
||||
|
||||
let refs = skill_dir.join("references");
|
||||
std::fs::create_dir_all(&refs).unwrap();
|
||||
let guide = refs.join("guide.md");
|
||||
std::fs::write(&guide, "guide").unwrap();
|
||||
std::os::unix::fs::symlink(&refs, refs.join("loop")).unwrap();
|
||||
|
||||
let files = collect_supporting_files(skill_dir);
|
||||
assert_eq!(files, vec![guide]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_skills_populates_supporting_files() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
let skill_dir = cwd.join(".agents/skills/with-refs");
|
||||
std::fs::create_dir_all(&skill_dir).unwrap();
|
||||
std::fs::write(
|
||||
skill_dir.join("SKILL.md"),
|
||||
"---\nname: with-refs\ndescription: Has refs\n---\nBody.\n",
|
||||
)
|
||||
.unwrap();
|
||||
let refs = skill_dir.join("references");
|
||||
std::fs::create_dir_all(&refs).unwrap();
|
||||
std::fs::write(refs.join("guide.md"), "guide content").unwrap();
|
||||
|
||||
let skills = discover_skills_impl(cwd, None);
|
||||
assert_eq!(skills.len(), 1);
|
||||
assert_eq!(skills[0].name, "with-refs");
|
||||
assert_eq!(skills[0].supporting_files.len(), 1);
|
||||
assert!(
|
||||
skills[0].supporting_files[0].ends_with("references/guide.md"),
|
||||
"unexpected path: {:?}",
|
||||
skills[0].supporting_files[0]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,974 @@
|
||||
#![forbid(unsafe_code)]
|
||||
mod agent;
|
||||
pub mod auth;
|
||||
mod builtin;
|
||||
pub mod catalog;
|
||||
pub mod config;
|
||||
mod handoff;
|
||||
mod hints;
|
||||
mod llm;
|
||||
mod mcp;
|
||||
pub mod types;
|
||||
mod wire;
|
||||
|
||||
pub use catalog::{discover_databricks_models, ModelEntry, DATABRICKS_V2_KNOWN_MODELS};
|
||||
pub use config::Provider;
|
||||
pub use types::AgentError;
|
||||
|
||||
/// Environment keys the Windows Git Bash resolver may inspect. `spawn_one()`
|
||||
/// forwards every key in this list into its otherwise-cleared MCP child; Doctor
|
||||
/// uses the same contract so a ready agent can always start its shell tool.
|
||||
#[cfg(windows)]
|
||||
pub const WINDOWS_SHELL_RESOLUTION_ENV: &[&str] = &[
|
||||
"PATH",
|
||||
"BUZZ_SHELL",
|
||||
"GIT_BASH",
|
||||
"SystemRoot",
|
||||
"ProgramFiles",
|
||||
"ProgramFiles(x86)",
|
||||
"LOCALAPPDATA",
|
||||
];
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::BufReader;
|
||||
use tokio::sync::{mpsc, watch, Mutex};
|
||||
|
||||
use crate::agent::RunCtx;
|
||||
use crate::config::{Config, MAX_SYSTEM_PROMPT_BYTES, PROTOCOL_VERSION};
|
||||
use crate::hints::SkillEntry;
|
||||
use crate::llm::Llm;
|
||||
use crate::mcp::McpRegistry;
|
||||
use crate::types::{ContentBlock, HistoryItem};
|
||||
use crate::wire::{
|
||||
classify, goose_session_update, Inbound, InitializeParams, SessionCancelParams,
|
||||
SessionNewParams, SessionPromptParams, SessionSetModelParams, SessionSteerParams, WireMsg,
|
||||
WireSender, INVALID_PARAMS, METHOD_NOT_FOUND, PARSE_ERROR,
|
||||
};
|
||||
|
||||
struct App {
|
||||
cfg: Config,
|
||||
llm: Arc<Llm>,
|
||||
sessions: Mutex<HashMap<String, Session>>,
|
||||
/// Cached model catalog for Databricks providers. Populated lazily on the
|
||||
/// first successful `session/new` discovery call. Failed discovery is never
|
||||
/// cached: static-token authentication errors reject session creation, while
|
||||
/// OAuth authentication and non-auth errors use the configured model for that
|
||||
/// response and retry on the next session.
|
||||
models_cache: tokio::sync::OnceCell<Vec<ModelEntry>>,
|
||||
}
|
||||
|
||||
struct Session {
|
||||
id: String,
|
||||
mcp: Arc<McpRegistry>,
|
||||
/// Skills discovered at session creation; used by the built-in `load_skill` tool.
|
||||
skills: Vec<SkillEntry>,
|
||||
history: Vec<HistoryItem>,
|
||||
cancel_tx: watch::Sender<bool>,
|
||||
busy: bool,
|
||||
/// Run id of the in-flight prompt, set when a prompt starts and cleared
|
||||
/// when it ends. `None` means no active run — a steer request targeting
|
||||
/// this session is rejected. Steer-capable clients learn this value from
|
||||
/// the `params.update._meta.goose.activeRunId` field on `session/update`.
|
||||
active_run_id: Option<String>,
|
||||
/// Sender for mid-turn steer messages. Created fresh per prompt (like
|
||||
/// `cancel_tx`); the running prompt loop holds the matching receiver and
|
||||
/// drains queued steers at round boundaries. `None` when no prompt is in
|
||||
/// flight.
|
||||
steer_tx: Option<mpsc::UnboundedSender<Vec<ContentBlock>>>,
|
||||
original_task: Option<String>,
|
||||
handoff_count: usize,
|
||||
/// Cache-summed input tokens the provider reported for this session's most
|
||||
/// recent request, or `None` before the first response (or after a handoff
|
||||
/// resets the context). Drives the token-based handoff gate; see
|
||||
/// [`RunCtx::should_handoff`].
|
||||
last_request_input_tokens: Option<u64>,
|
||||
/// History byte size when `last_request_input_tokens` was measured, paired
|
||||
/// with it so the gate can account for history appended since.
|
||||
last_request_history_bytes: Option<usize>,
|
||||
effective_system_prompt: Arc<str>,
|
||||
/// Per-session model override set by `session/set_model`. When `Some`,
|
||||
/// overrides `App::cfg.model` for all LLM calls on this session. Persists
|
||||
/// across `session/prompt` calls until changed.
|
||||
effective_model: Option<String>,
|
||||
/// Session-cumulative input tokens across all turns. Sent in the
|
||||
/// `_goose/unstable/session/update` usage notification so buzz-acp's
|
||||
/// `UsageTracker` can compute per-turn deltas symmetrically with goose.
|
||||
accumulated_input_tokens: u64,
|
||||
/// Session-cumulative output tokens across all turns.
|
||||
accumulated_output_tokens: u64,
|
||||
/// Session-cumulative cache-served input tokens across all turns — a subset
|
||||
/// of `accumulated_input_tokens`, not an addition to it. Emitted alongside
|
||||
/// it so a consumer can price the cached slice at the provider's discounted
|
||||
/// rate instead of assuming every input token cost full price.
|
||||
accumulated_cached_input_tokens: u64,
|
||||
/// Session-cumulative total-token state across all turns.
|
||||
///
|
||||
/// Mirrors the per-turn `TurnTotalState` tri-state: starts `Unseen`,
|
||||
/// becomes `Exact(n)` as turns with genuine provider totals complete,
|
||||
/// transitions permanently to `Unknown` when any turn lacks a total or
|
||||
/// when the cumulative would otherwise decrease. Only emitted in the
|
||||
/// `usage_update` notification when `Exact`.
|
||||
accumulated_total_state: crate::types::TurnTotalState,
|
||||
}
|
||||
|
||||
fn die(msg: String) -> ! {
|
||||
tracing::error!("{msg}");
|
||||
std::process::exit(2);
|
||||
}
|
||||
|
||||
pub fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
if matches!(args.get(1).map(String::as_str), Some("auth")) {
|
||||
return tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()?
|
||||
.block_on(auth_subcommand(&args[2..]));
|
||||
}
|
||||
tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()?
|
||||
.block_on(async_main());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn authenticate_databricks(host: &str) -> Result<(), AgentError> {
|
||||
auth::PkceOAuthTokenSource::new(llm::databricks_pkce_config(host))?
|
||||
.interactive_login()
|
||||
.await
|
||||
}
|
||||
|
||||
/// `buzz-agent auth <provider>` — run the interactive auth flow for a
|
||||
/// provider and persist the result, then exit. Today this supports Databricks
|
||||
/// OAuth 2.0 PKCE. Reads `DATABRICKS_HOST` from env; needs a browser on the
|
||||
/// machine.
|
||||
async fn auth_subcommand(args: &[String]) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let provider = args.first().map(String::as_str);
|
||||
match provider {
|
||||
Some("databricks" | "databricks_v2" | "databricks-v2") => {
|
||||
let host = std::env::var("DATABRICKS_HOST")
|
||||
.map_err(|_| "auth databricks: DATABRICKS_HOST required")?;
|
||||
authenticate_databricks(&host).await?;
|
||||
eprintln!("Authenticated. Token cached under ~/.config/buzz-agent/oauth/databricks/.");
|
||||
Ok(())
|
||||
}
|
||||
Some(other) => Err(format!("auth: unknown provider {other:?}").into()),
|
||||
None => Err("auth: provider required (try: buzz-agent auth databricks)".into()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn async_main() {
|
||||
tracing_subscriber::fmt()
|
||||
.with_writer(std::io::stderr)
|
||||
.with_ansi(false)
|
||||
.init();
|
||||
let cfg = Config::from_env().unwrap_or_else(|e| die(e));
|
||||
let llm = Arc::new(Llm::new(&cfg).unwrap_or_else(|e| die(e.to_string())));
|
||||
let max_line = cfg.max_line_bytes;
|
||||
let app = Arc::new(App {
|
||||
cfg,
|
||||
llm,
|
||||
sessions: Mutex::new(HashMap::new()),
|
||||
models_cache: tokio::sync::OnceCell::new(),
|
||||
});
|
||||
let (wire_tx, wire_rx) = mpsc::channel::<WireMsg>(64);
|
||||
let writer = tokio::spawn(wire::writer_task(wire_rx));
|
||||
if let Err(e) = read_loop(
|
||||
BufReader::new(tokio::io::stdin()),
|
||||
app.clone(),
|
||||
wire_tx,
|
||||
max_line,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!("io: reader: {e}");
|
||||
}
|
||||
for session in app.sessions.lock().await.values() {
|
||||
let _ = session.cancel_tx.send(true);
|
||||
}
|
||||
let _ = writer.await;
|
||||
}
|
||||
|
||||
async fn read_loop<R: tokio::io::AsyncBufRead + Unpin>(
|
||||
mut stdin: R,
|
||||
app: Arc<App>,
|
||||
wire_tx: WireSender,
|
||||
max_line: usize,
|
||||
) -> std::io::Result<()> {
|
||||
while let Some(line) = wire::read_bounded_line(&mut stdin, max_line).await? {
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
match serde_json::from_str::<Value>(&line) {
|
||||
Ok(msg) => dispatch(&app, msg, &wire_tx).await,
|
||||
Err(e) => {
|
||||
wire::send(
|
||||
&wire_tx,
|
||||
wire::err(Value::Null, PARSE_ERROR, &format!("jsonrpc: parse: {e}")),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn dispatch(app: &Arc<App>, msg: Value, wire_tx: &WireSender) {
|
||||
match classify(&msg) {
|
||||
Inbound::Request { id, method, params } => {
|
||||
handle_request(app, id, method, params, wire_tx).await
|
||||
}
|
||||
Inbound::Notification { method, params } => handle_notification(app, &method, params).await,
|
||||
Inbound::Ignored => {}
|
||||
Inbound::Invalid { id, code, message } => {
|
||||
wire::send(wire_tx, wire::err(id, code, &message)).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_request(
|
||||
app: &Arc<App>,
|
||||
id: Value,
|
||||
method: String,
|
||||
params: Value,
|
||||
wire_tx: &WireSender,
|
||||
) {
|
||||
match method.as_str() {
|
||||
"initialize" => initialize(id, params, wire_tx).await,
|
||||
"session/new" => {
|
||||
let app = app.clone();
|
||||
let wire_tx = wire_tx.clone();
|
||||
tokio::spawn(async move { session_new(&app, id, params, &wire_tx).await });
|
||||
}
|
||||
"session/prompt" => spawn_prompt(app.clone(), id, params, wire_tx.clone()),
|
||||
"session/set_model" => {
|
||||
set_model_session(app, id, params, wire_tx).await;
|
||||
}
|
||||
"session/cancel" => {
|
||||
cancel_session(app, params).await;
|
||||
wire::send(wire_tx, wire::ok(id, Value::Null)).await;
|
||||
}
|
||||
// goose-compatible non-standard extension: inject user input into the
|
||||
// currently active prompt without starting a new one. Mirrors goose's
|
||||
// `_goose/unstable/session/steer` wire contract so a single client-side
|
||||
// delivery path serves both agents.
|
||||
"_goose/unstable/session/steer" => {
|
||||
steer_session(app, id, params, wire_tx).await;
|
||||
}
|
||||
_ => {
|
||||
wire::send(
|
||||
wire_tx,
|
||||
wire::err(
|
||||
id,
|
||||
METHOD_NOT_FOUND,
|
||||
&format!("jsonrpc: method not found: {method}"),
|
||||
),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_notification(app: &Arc<App>, method: &str, params: Value) {
|
||||
if method == "session/cancel" {
|
||||
cancel_session(app, params).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn initialize(id: Value, params: Value, wire_tx: &WireSender) {
|
||||
let p: InitializeParams = match decode(params, "initialize") {
|
||||
Ok(p) => p,
|
||||
Err(m) => return reject(wire_tx, id, INVALID_PARAMS, &m).await,
|
||||
};
|
||||
// Honest negotiation: respond with the minimum of what the client
|
||||
// requested and what we support.
|
||||
// NOTE: gating `[Base]` injection on `protocol_version < 2` is a deliberate
|
||||
// temporary measure — we are squatting on ACP v2 ahead of the upstream ACP
|
||||
// RFD. Revisit when that RFD merges; otherwise a genuine upstream-v2 agent
|
||||
// would silently lose `[Base]`.
|
||||
let negotiated_version = p.protocol_version.min(PROTOCOL_VERSION);
|
||||
wire::send(
|
||||
wire_tx,
|
||||
wire::ok(
|
||||
id,
|
||||
json!({
|
||||
"protocolVersion": negotiated_version,
|
||||
"agentCapabilities": {
|
||||
"loadSession": false,
|
||||
"promptCapabilities": { "image": false, "audio": false, "embeddedContext": false },
|
||||
"mcpCapabilities": { "http": false, "sse": false },
|
||||
},
|
||||
"agentInfo": { "name": "buzz-agent", "version": env!("CARGO_PKG_VERSION") },
|
||||
}),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Resolve the Databricks model catalog for one `session/new` call.
|
||||
///
|
||||
/// Tries to use a previously-cached successful discovery result. If the cache is empty,
|
||||
/// runs `discover` and — on success — populates the cache for future calls. On failure
|
||||
/// the error is returned and the cell is intentionally left empty so the next session retries.
|
||||
///
|
||||
/// Extracted from `session_new` so that tests can drive this path with an injected
|
||||
/// discovery future without requiring a full `App` / transport stack.
|
||||
async fn resolve_models_catalog(
|
||||
cache: &tokio::sync::OnceCell<Vec<ModelEntry>>,
|
||||
discover: impl std::future::Future<Output = Result<Vec<ModelEntry>, AgentError>>,
|
||||
) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
cache.get_or_try_init(|| discover).await.cloned()
|
||||
}
|
||||
|
||||
/// Return the configured model as a one-entry catalog for this response.
|
||||
///
|
||||
/// This value is never written to `models_cache`; failed discovery must be retried by
|
||||
/// the next session rather than pinning degraded state for the process lifetime.
|
||||
fn configured_model_fallback(model: &str) -> Vec<ModelEntry> {
|
||||
let model = model.trim().to_string();
|
||||
vec![ModelEntry {
|
||||
id: model.clone(),
|
||||
name: model,
|
||||
}]
|
||||
}
|
||||
|
||||
async fn session_new(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSender) {
|
||||
let p: SessionNewParams = match decode(params, "session/new") {
|
||||
Ok(p) => p,
|
||||
Err(m) => return reject(wire_tx, id, INVALID_PARAMS, &m).await,
|
||||
};
|
||||
if p.cwd.is_empty() || !Path::new(&p.cwd).is_absolute() {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
"session/new: cwd must be an absolute path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
// Check cap without holding lock across MCP spawn (which may be slow).
|
||||
{
|
||||
let sessions = app.sessions.lock().await;
|
||||
if sessions.len() >= app.cfg.max_sessions {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
"session/new: max sessions reached",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
let (hints_text, skills) = if app.cfg.hints_enabled {
|
||||
hints::build_hints_section(std::path::Path::new(&p.cwd))
|
||||
} else {
|
||||
(String::new(), Vec::new())
|
||||
};
|
||||
let effective_system_prompt: Arc<str> = {
|
||||
// When the harness provides a systemPrompt (base_prompt + persona), use
|
||||
// it as the primary content and suppress the default. The default is only
|
||||
// a fallback for legacy harnesses that don't send systemPrompt.
|
||||
let base = match p.system_prompt.as_deref() {
|
||||
Some(client_prompt) if !client_prompt.trim().is_empty() => client_prompt.to_owned(),
|
||||
_ => app.cfg.system_prompt.clone(),
|
||||
};
|
||||
let prompt = if hints_text.is_empty() {
|
||||
base
|
||||
} else {
|
||||
format!("{base}\n\n{hints_text}")
|
||||
};
|
||||
// Reject combined prompts exceeding 512KB.
|
||||
if prompt.len() > MAX_SYSTEM_PROMPT_BYTES {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
&format!(
|
||||
"session/new: combined system prompt exceeds {}KB limit ({} bytes)",
|
||||
MAX_SYSTEM_PROMPT_BYTES / 1024,
|
||||
prompt.len()
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Arc::from(prompt)
|
||||
};
|
||||
// Resolve the model catalog before spawning MCP servers or registering a
|
||||
// session. A configured static credential cannot recover interactively, so
|
||||
// its authentication failure rejects before allocation. OAuth authentication
|
||||
// failures and other catalog failures use only the configured model for this
|
||||
// response, without caching, so session/prompt can run the existing PKCE flow.
|
||||
let available_models: Vec<Value> = {
|
||||
use crate::config::Provider;
|
||||
match app.cfg.provider {
|
||||
Provider::Databricks | Provider::DatabricksV2 => {
|
||||
let models = match resolve_models_catalog(
|
||||
&app.models_cache,
|
||||
discover_databricks_models(&app.cfg),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(models) => models,
|
||||
Err(error @ AgentError::LlmAuth(_)) if !app.cfg.api_key.is_empty() => {
|
||||
return reject(wire_tx, id, error.json_rpc_code(), &error.to_string())
|
||||
.await;
|
||||
}
|
||||
Err(error @ AgentError::LlmAuth(_)) => {
|
||||
tracing::warn!(
|
||||
error = %error,
|
||||
"Databricks OAuth model catalog unavailable; using configured model"
|
||||
);
|
||||
configured_model_fallback(&app.cfg.model)
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
error = %error,
|
||||
"Databricks model catalog unavailable; using configured model"
|
||||
);
|
||||
configured_model_fallback(&app.cfg.model)
|
||||
}
|
||||
};
|
||||
models
|
||||
.iter()
|
||||
.map(|m| json!({ "modelId": m.id, "name": m.name }))
|
||||
.collect()
|
||||
}
|
||||
_ => vec![json!({ "modelId": app.cfg.model, "name": app.cfg.model })],
|
||||
}
|
||||
};
|
||||
|
||||
let mcp = match McpRegistry::spawn_all(&app.cfg, &p.mcp_servers, &p.cwd).await {
|
||||
Ok(m) => Arc::new(m),
|
||||
Err(e) => return reject(wire_tx, id, e.json_rpc_code(), &e.to_string()).await,
|
||||
};
|
||||
let session_id = match session_token() {
|
||||
Ok(t) => format!("ses_{t}"),
|
||||
Err(e) => return reject(wire_tx, id, -32000, &e).await,
|
||||
};
|
||||
let (cancel_tx, _) = watch::channel(false);
|
||||
let mut sessions = app.sessions.lock().await;
|
||||
// Re-check cap (another session may have been created while we spawned MCP).
|
||||
if sessions.len() >= app.cfg.max_sessions {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
"session/new: max sessions reached",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
sessions.insert(
|
||||
session_id.clone(),
|
||||
Session {
|
||||
id: session_id.clone(),
|
||||
mcp,
|
||||
skills,
|
||||
history: Vec::new(),
|
||||
cancel_tx,
|
||||
busy: false,
|
||||
active_run_id: None,
|
||||
steer_tx: None,
|
||||
original_task: None,
|
||||
handoff_count: 0,
|
||||
last_request_input_tokens: None,
|
||||
last_request_history_bytes: None,
|
||||
effective_system_prompt,
|
||||
effective_model: None,
|
||||
accumulated_input_tokens: 0,
|
||||
accumulated_output_tokens: 0,
|
||||
accumulated_cached_input_tokens: 0,
|
||||
accumulated_total_state: crate::types::TurnTotalState::Unseen,
|
||||
},
|
||||
);
|
||||
drop(sessions);
|
||||
|
||||
wire::send(
|
||||
wire_tx,
|
||||
wire::ok(
|
||||
id,
|
||||
json!({
|
||||
"sessionId": session_id,
|
||||
"models": {
|
||||
"currentModelId": app.cfg.model,
|
||||
"availableModels": available_models,
|
||||
},
|
||||
}),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
fn decode<T: serde::de::DeserializeOwned>(params: Value, stage: &str) -> Result<T, String> {
|
||||
serde_json::from_value(params).map_err(|e| format!("{stage}: {e}"))
|
||||
}
|
||||
|
||||
async fn reject(wire_tx: &WireSender, id: Value, code: i32, message: &str) {
|
||||
wire::send(wire_tx, wire::err(id, code, message)).await;
|
||||
}
|
||||
|
||||
async fn cancel_session(app: &Arc<App>, params: Value) {
|
||||
if let Ok(p) = serde_json::from_value::<SessionCancelParams>(params) {
|
||||
if let Some(s) = app.sessions.lock().await.get(&p.session_id) {
|
||||
let _ = s.cancel_tx.send(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle `session/set_model`: apply a per-session model override immediately.
|
||||
///
|
||||
/// Validation:
|
||||
/// - Unknown `sessionId` → `invalid_params`.
|
||||
/// - Empty `modelId` → `invalid_params`.
|
||||
///
|
||||
/// On success: stores `model_id` on the session and responds `{ sessionId, modelId }`.
|
||||
/// The override is picked up by the next `session/prompt` call on this session.
|
||||
async fn set_model_session(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSender) {
|
||||
let p: SessionSetModelParams = match decode(params, "session/set_model") {
|
||||
Ok(p) => p,
|
||||
Err(m) => return reject(wire_tx, id, INVALID_PARAMS, &m).await,
|
||||
};
|
||||
if p.model_id.trim().is_empty() {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
"session/set_model: modelId must not be empty",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let mut sessions = app.sessions.lock().await;
|
||||
let Some(s) = sessions.get_mut(&p.session_id) else {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
"session/set_model: unknown session",
|
||||
)
|
||||
.await;
|
||||
};
|
||||
s.effective_model = Some(p.model_id.clone());
|
||||
tracing::info!(
|
||||
session_id = %p.session_id,
|
||||
model_id = %p.model_id,
|
||||
"session/set_model: model overridden"
|
||||
);
|
||||
drop(sessions);
|
||||
wire::send(
|
||||
wire_tx,
|
||||
wire::ok(
|
||||
id,
|
||||
json!({ "sessionId": p.session_id, "modelId": p.model_id }),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Handle `_goose/unstable/session/steer`: queue user input into the in-flight
|
||||
/// prompt. Validation mirrors goose's `on_steer_session`:
|
||||
/// - empty prompt → `invalid_params`
|
||||
/// - no active run (no prompt in flight) → `invalid_params`
|
||||
/// - `expectedRunId` mismatch → `invalid_params` (caller is steering a turn
|
||||
/// that already ended or rotated; it must fall back to cancel+merge)
|
||||
///
|
||||
/// On success the message is queued for pickup at the next round boundary and
|
||||
/// we reply `{ runId, messageId }`, then emit a `queuedSteer` session/update so
|
||||
/// the client can correlate the accepted steer with its eventual pickup.
|
||||
async fn steer_session(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSender) {
|
||||
let p: SessionSteerParams = match decode(params, "_goose/unstable/session/steer") {
|
||||
Ok(p) => p,
|
||||
Err(m) => return reject(wire_tx, id, INVALID_PARAMS, &m).await,
|
||||
};
|
||||
if p.prompt.is_empty() {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
"steer: prompt must not be empty",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if p.expected_run_id.is_empty() {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
"steer: expectedRunId must not be empty",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let message_id = format!("steer_{}", session_token().unwrap_or_else(|_| "x".into()));
|
||||
let run_id = {
|
||||
let sessions = app.sessions.lock().await;
|
||||
let Some(s) = sessions.get(&p.session_id) else {
|
||||
return reject(wire_tx, id, INVALID_PARAMS, "steer: unknown session").await;
|
||||
};
|
||||
let Some(active) = s.active_run_id.as_deref() else {
|
||||
return reject(wire_tx, id, INVALID_PARAMS, "steer: no active run to steer").await;
|
||||
};
|
||||
if active != p.expected_run_id {
|
||||
return reject(
|
||||
wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
&format!(
|
||||
"steer: expected active run id `{}` but found `{active}`",
|
||||
p.expected_run_id
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
// A live run always has a steer_tx; if the channel is gone the run is
|
||||
// tearing down — treat as no active run rather than queue into the void.
|
||||
match &s.steer_tx {
|
||||
Some(tx) if tx.send(p.prompt).is_ok() => active.to_owned(),
|
||||
_ => return reject(wire_tx, id, INVALID_PARAMS, "steer: no active run to steer").await,
|
||||
}
|
||||
};
|
||||
wire::send(
|
||||
wire_tx,
|
||||
wire::ok(id, json!({ "runId": run_id, "messageId": message_id })),
|
||||
)
|
||||
.await;
|
||||
// Best-effort correlation hint for the client; mirrors goose's
|
||||
// `send_queued_steer_update`. Not load-bearing for delivery.
|
||||
wire::send(
|
||||
wire_tx,
|
||||
wire::session_update_with_goose_meta(
|
||||
&p.session_id,
|
||||
json!({ "sessionUpdate": "session_info_update" }),
|
||||
json!({ "queuedSteer": { "messageId": message_id, "runId": run_id } }),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
fn spawn_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender) {
|
||||
tokio::spawn(async move { run_prompt(app, id, params, wire_tx).await });
|
||||
}
|
||||
|
||||
async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender) {
|
||||
let p: SessionPromptParams = match decode(params, "session/prompt") {
|
||||
Ok(p) => p,
|
||||
Err(m) => return reject(&wire_tx, id, INVALID_PARAMS, &m).await,
|
||||
};
|
||||
let (
|
||||
sid,
|
||||
mcp,
|
||||
skills,
|
||||
mut history,
|
||||
mut original_task,
|
||||
mut handoff_count,
|
||||
mut last_request_input_tokens,
|
||||
mut last_request_history_bytes,
|
||||
mut cancel_rx,
|
||||
effective_system_prompt,
|
||||
effective_model_override,
|
||||
run_id,
|
||||
mut steer_rx,
|
||||
usage_baseline,
|
||||
) = match acquire_session(&app, &p.session_id).await {
|
||||
Ok(v) => v,
|
||||
Err(reason) => {
|
||||
return reject(
|
||||
&wire_tx,
|
||||
id,
|
||||
INVALID_PARAMS,
|
||||
&format!("session/prompt: {reason}"),
|
||||
)
|
||||
.await
|
||||
}
|
||||
};
|
||||
// Advertise the active run id so steer-capable clients can target this turn
|
||||
// via `expectedRunId`. Mirrors goose's `send_active_run_update`.
|
||||
wire::send(
|
||||
&wire_tx,
|
||||
wire::session_update_with_goose_meta(
|
||||
&sid,
|
||||
json!({ "sessionUpdate": "session_info_update" }),
|
||||
json!({ "activeRunId": run_id }),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
// Resolve effective model: session override wins over config default.
|
||||
let effective_model_str = effective_model_override
|
||||
.as_deref()
|
||||
.unwrap_or(&app.cfg.model);
|
||||
let mut turn_input_tokens: Option<u64> = None;
|
||||
let mut turn_output_tokens: Option<u64> = None;
|
||||
let mut turn_cached_input_tokens: Option<u64> = None;
|
||||
let mut turn_total_state = crate::types::TurnTotalState::Unseen;
|
||||
let mut ctx = RunCtx {
|
||||
cfg: &app.cfg,
|
||||
effective_model: effective_model_str,
|
||||
session_id: &sid,
|
||||
system_prompt: &effective_system_prompt,
|
||||
llm: &app.llm,
|
||||
mcp: &mcp,
|
||||
skills: &skills,
|
||||
wire: &wire_tx,
|
||||
cancel: &mut cancel_rx,
|
||||
steer: &mut steer_rx,
|
||||
history: &mut history,
|
||||
original_task: &mut original_task,
|
||||
handoff_count: &mut handoff_count,
|
||||
run_id,
|
||||
last_request_input_tokens: &mut last_request_input_tokens,
|
||||
last_request_history_bytes: &mut last_request_history_bytes,
|
||||
turn_input_tokens: &mut turn_input_tokens,
|
||||
turn_output_tokens: &mut turn_output_tokens,
|
||||
turn_cached_input_tokens: &mut turn_cached_input_tokens,
|
||||
turn_total_state: &mut turn_total_state,
|
||||
usage_baseline,
|
||||
};
|
||||
let result = ctx.run(p.prompt).await;
|
||||
if let Some(s) = app.sessions.lock().await.get_mut(&sid) {
|
||||
s.busy = false;
|
||||
// Clear run state so a late steer can't queue into a finished turn.
|
||||
s.active_run_id = None;
|
||||
s.steer_tx = None;
|
||||
s.history = history;
|
||||
s.original_task = original_task;
|
||||
s.handoff_count = handoff_count;
|
||||
s.last_request_input_tokens = last_request_input_tokens;
|
||||
s.last_request_history_bytes = last_request_history_bytes;
|
||||
}
|
||||
// Update session-cumulative token counters and emit the usage notification
|
||||
// BEFORE sending the session/prompt response. buzz-acp's UsageTracker
|
||||
// processes the notification while the turn is still in-flight (i.e. before
|
||||
// the response triggers take_turn_usage()), which is required for the
|
||||
// begin_turn gate to recognise it as publishable.
|
||||
//
|
||||
// Only emit when at least one token count was observed — a turn with no
|
||||
// provider response (validation failure, pre-response cancellation) carries
|
||||
// no information and must not produce a kind 44200 record per NIP-AM.
|
||||
if turn_input_tokens.is_some() || turn_output_tokens.is_some() {
|
||||
let accumulated = {
|
||||
let mut sessions = app.sessions.lock().await;
|
||||
if let Some(s) = sessions.get_mut(&sid) {
|
||||
s.accumulated_input_tokens = s
|
||||
.accumulated_input_tokens
|
||||
.saturating_add(turn_input_tokens.unwrap_or(0));
|
||||
s.accumulated_output_tokens = s
|
||||
.accumulated_output_tokens
|
||||
.saturating_add(turn_output_tokens.unwrap_or(0));
|
||||
s.accumulated_cached_input_tokens = s
|
||||
.accumulated_cached_input_tokens
|
||||
.saturating_add(turn_cached_input_tokens.unwrap_or(0));
|
||||
// Fold the per-turn total state into the session cumulative.
|
||||
// Unknown poisons the session permanently; Exact adds to running sum;
|
||||
// Unseen (turn emitted no usage) leaves the cumulative unchanged.
|
||||
// Uses TurnTotalState::merge_session, which applies the same
|
||||
// checked-add / overflow-poisons contract as the per-response fold.
|
||||
s.accumulated_total_state =
|
||||
s.accumulated_total_state.merge_session(turn_total_state);
|
||||
Some((
|
||||
s.accumulated_input_tokens,
|
||||
s.accumulated_output_tokens,
|
||||
s.accumulated_cached_input_tokens,
|
||||
s.accumulated_total_state,
|
||||
))
|
||||
} else {
|
||||
// Session is gone — the accumulated baseline no longer exists, so
|
||||
// there is nothing correct to emit. Skip the usage notification.
|
||||
None
|
||||
}
|
||||
};
|
||||
if let Some((accumulated_in, accumulated_out, accumulated_cached, accumulated_total)) =
|
||||
accumulated
|
||||
{
|
||||
// Same builder the run loop uses for its per-round reports, so the
|
||||
// final notification is shape-identical to the ones that preceded
|
||||
// it and a consumer taking the high-water mark lands on this one.
|
||||
let update = wire::usage_update_payload(
|
||||
accumulated_in,
|
||||
accumulated_out,
|
||||
accumulated_cached,
|
||||
accumulated_total,
|
||||
effective_model_str,
|
||||
);
|
||||
wire::send(&wire_tx, goose_session_update(&sid, update)).await;
|
||||
}
|
||||
}
|
||||
match result {
|
||||
Ok(stop) => {
|
||||
wire::send(
|
||||
&wire_tx,
|
||||
wire::ok(id, json!({ "stopReason": stop.as_wire() })),
|
||||
)
|
||||
.await
|
||||
}
|
||||
Err(e) => wire::send(&wire_tx, wire::err(id, e.json_rpc_code(), &e.to_string())).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn acquire_session(
|
||||
app: &Arc<App>,
|
||||
session_id: &str,
|
||||
) -> Result<
|
||||
(
|
||||
String,
|
||||
Arc<McpRegistry>,
|
||||
Vec<SkillEntry>,
|
||||
Vec<HistoryItem>,
|
||||
Option<String>,
|
||||
usize,
|
||||
Option<u64>,
|
||||
Option<usize>,
|
||||
watch::Receiver<bool>,
|
||||
Arc<str>,
|
||||
Option<String>,
|
||||
String,
|
||||
mpsc::UnboundedReceiver<Vec<ContentBlock>>,
|
||||
crate::types::SessionUsageBaseline,
|
||||
),
|
||||
&'static str,
|
||||
> {
|
||||
let mut sessions = app.sessions.lock().await;
|
||||
let s = sessions.get_mut(session_id).ok_or("unknown session")?;
|
||||
if s.busy {
|
||||
return Err("prompt already in flight");
|
||||
}
|
||||
// Generate the run id before mutating session state. On RNG failure we reject
|
||||
// the prompt cleanly: the session stays idle and the caller can retry. Generating
|
||||
// after `s.busy = true` with `?` would wedge the session permanently busy.
|
||||
let run_id = format!(
|
||||
"run_{}",
|
||||
session_token().map_err(|_| "rng failure; retry prompt")?
|
||||
);
|
||||
s.busy = true;
|
||||
let (tx, rx) = watch::channel(false);
|
||||
s.cancel_tx = tx;
|
||||
// Skills are read-only after session creation; clone the Vec so RunCtx
|
||||
// can hold a reference without holding the sessions lock.
|
||||
let skills = s.skills.clone();
|
||||
// Fresh run id + steer channel for this turn. The run id lets steer-capable
|
||||
// clients target *this* turn (rejecting steers aimed at a turn that already
|
||||
// ended); the channel carries mid-turn injections to the run loop.
|
||||
s.active_run_id = Some(run_id.clone());
|
||||
let (steer_tx, steer_rx) = mpsc::unbounded_channel();
|
||||
s.steer_tx = Some(steer_tx);
|
||||
let effective_model = s.effective_model.clone();
|
||||
Ok((
|
||||
s.id.clone(),
|
||||
s.mcp.clone(),
|
||||
skills,
|
||||
std::mem::take(&mut s.history),
|
||||
s.original_task.take(),
|
||||
s.handoff_count,
|
||||
s.last_request_input_tokens,
|
||||
s.last_request_history_bytes,
|
||||
rx,
|
||||
Arc::clone(&s.effective_system_prompt),
|
||||
effective_model,
|
||||
run_id,
|
||||
steer_rx,
|
||||
// Snapshot rather than a handle: the run loop reports cumulative usage
|
||||
// after every LLM round, and taking the sessions lock on each of those
|
||||
// would serialise concurrent sessions behind one another's provider
|
||||
// round-trips. Nothing else advances these counters while this turn
|
||||
// holds `busy`, so the snapshot cannot go stale under it.
|
||||
crate::types::SessionUsageBaseline {
|
||||
input_tokens: s.accumulated_input_tokens,
|
||||
output_tokens: s.accumulated_output_tokens,
|
||||
cached_input_tokens: s.accumulated_cached_input_tokens,
|
||||
total_state: s.accumulated_total_state,
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
fn session_token() -> Result<String, String> {
|
||||
let mut b = [0u8; 8];
|
||||
getrandom::fill(&mut b).map_err(|e| format!("rng: getrandom failed: {e}"))?;
|
||||
Ok(b.iter().map(|x| format!("{x:02x}")).collect())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::catalog::ModelEntry;
|
||||
use crate::types::AgentError;
|
||||
|
||||
/// Regression: a discovery error must not pin the models_cache for the process lifetime.
|
||||
///
|
||||
/// `resolve_models_catalog` uses `get_or_try_init` so an `Err` leaves the `OnceCell`
|
||||
/// empty and the next `session/new` retries discovery. This test calls
|
||||
/// `resolve_models_catalog` directly — the same function `session_new` calls — so
|
||||
/// reverting `session_new` to `get_or_init` (or any other cache-on-error variant) would
|
||||
/// break this test, not just the standalone `OnceCell` semantics.
|
||||
#[tokio::test]
|
||||
async fn models_cache_does_not_pin_on_discovery_error() {
|
||||
let cache: tokio::sync::OnceCell<Vec<ModelEntry>> = tokio::sync::OnceCell::new();
|
||||
|
||||
// First call — discovery failure is surfaced and leaves the cell empty.
|
||||
let error = crate::resolve_models_catalog(&cache, async {
|
||||
Err::<Vec<ModelEntry>, AgentError>(AgentError::Llm("transient failure".into()))
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(error, AgentError::Llm(_)));
|
||||
|
||||
// Second call — discovery succeeds. Cell is now populated and returned.
|
||||
let discovered = vec![ModelEntry {
|
||||
id: "databricks-meta-llama-3-1-70b-instruct".into(),
|
||||
name: "databricks-meta-llama-3-1-70b-instruct".into(),
|
||||
}];
|
||||
let discovered_clone = discovered.clone();
|
||||
let second = crate::resolve_models_catalog(&cache, async move {
|
||||
Ok::<Vec<ModelEntry>, AgentError>(discovered_clone)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
second, discovered,
|
||||
"second call must return the discovered catalog"
|
||||
);
|
||||
assert!(
|
||||
cache.get().is_some(),
|
||||
"cell must be populated after successful discovery"
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get().unwrap(),
|
||||
&discovered,
|
||||
"cache must hold the successful discovery result"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_catalog_does_not_cache_oauth_auth_fallback() {
|
||||
let cache: tokio::sync::OnceCell<Vec<ModelEntry>> = tokio::sync::OnceCell::new();
|
||||
let error = crate::resolve_models_catalog(&cache, async {
|
||||
Err::<Vec<ModelEntry>, AgentError>(AgentError::LlmAuth("sign in again".into()))
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(error, AgentError::LlmAuth(_)));
|
||||
assert!(cache.get().is_none());
|
||||
|
||||
let discovered = vec![ModelEntry {
|
||||
id: "authenticated-model".into(),
|
||||
name: "authenticated-model".into(),
|
||||
}];
|
||||
let result = crate::resolve_models_catalog(&cache, async {
|
||||
Ok::<Vec<ModelEntry>, AgentError>(discovered.clone())
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result, discovered);
|
||||
assert_eq!(cache.get(), Some(&discovered));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn configured_model_fallback_is_trimmed_and_singular() {
|
||||
assert_eq!(
|
||||
crate::configured_model_fallback(" configured-model "),
|
||||
vec![ModelEntry {
|
||||
id: "configured-model".into(),
|
||||
name: "configured-model".into(),
|
||||
}]
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,6 @@
|
||||
fn main() {
|
||||
if let Err(e) = buzz_agent::run() {
|
||||
eprintln!("Error: {e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,703 @@
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// Byte-equivalent charged to the handoff/context-pressure gate for a single
|
||||
/// image tool result. The gate maps bytes to tokens at 1 byte/token (see
|
||||
/// `handoff::CONSERVATIVE_BYTES_PER_TOKEN`), so this is also the per-image
|
||||
/// token budget. Providers bill an image as visual *tiles*, not its base64
|
||||
/// length: Anthropic caps at ~1600 tokens/image and OpenAI high-detail lands
|
||||
/// ~1.1K–1.5K. We charge 16 KiB — a generous ceiling that still over-counts
|
||||
/// the real ~2K cost, while being ~190× smaller than the base64 length of a
|
||||
/// typical multi-MiB screenshot. Charging `data.len()` to the gate instead
|
||||
/// made a single `view_image` (~3.1M base64 bytes) trip the handoff gate on a
|
||||
/// fresh context.
|
||||
const IMAGE_CONTEXT_TOKEN_EQUIV: usize = 16 * 1024;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ToolResultContent {
|
||||
Text(String),
|
||||
Image { data: String, mime_type: String },
|
||||
}
|
||||
|
||||
impl ToolResultContent {
|
||||
/// Real serialized size in bytes. Used by `truncate_history` to keep the
|
||||
/// outgoing request body under `max_history_bytes` — an image rides the
|
||||
/// wire as its full base64 string, so that string's length is what counts
|
||||
/// here. For context-window/handoff pressure use
|
||||
/// [`Self::context_pressure_bytes`] instead, which charges an image its
|
||||
/// (far smaller) visual-token equivalent.
|
||||
pub fn estimated_bytes(&self) -> usize {
|
||||
match self {
|
||||
Self::Text(s) => s.len(),
|
||||
Self::Image { data, mime_type } => data.len() + mime_type.len(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Token-equivalent context-window pressure, in bytes (the handoff gate
|
||||
/// maps bytes→tokens at 1:1). Identical to [`Self::estimated_bytes`] for
|
||||
/// text, but an image is charged a flat [`IMAGE_CONTEXT_TOKEN_EQUIV`]
|
||||
/// budget rather than its base64 length — providers bill it as visual
|
||||
/// tiles (~2K tokens), so counting `data.len()` over-counts by ~1500× and
|
||||
/// forces a handoff on a single image.
|
||||
pub fn context_pressure_bytes(&self) -> usize {
|
||||
match self {
|
||||
Self::Text(s) => s.len(),
|
||||
Self::Image { data: _, mime_type } => IMAGE_CONTEXT_TOKEN_EQUIV + mime_type.len(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_text_lossy(&self) -> String {
|
||||
match self {
|
||||
Self::Text(s) => s.clone(),
|
||||
Self::Image { data, mime_type } => {
|
||||
format!("[image: {mime_type}, {} base64 bytes]", data.len())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum HistoryItem {
|
||||
User(String),
|
||||
Assistant {
|
||||
text: String,
|
||||
tool_calls: Vec<ToolCall>,
|
||||
reasoning_details: Option<Value>,
|
||||
},
|
||||
ToolResult(ToolResult),
|
||||
}
|
||||
|
||||
impl HistoryItem {
|
||||
pub fn estimated_bytes(&self) -> usize {
|
||||
self.size_with(ToolResultContent::estimated_bytes)
|
||||
}
|
||||
|
||||
/// Token-equivalent context-window pressure, in bytes. Mirrors
|
||||
/// [`Self::estimated_bytes`] but charges image tool results their visual-
|
||||
/// token equivalent rather than their base64 length — see
|
||||
/// [`ToolResultContent::context_pressure_bytes`]. The handoff gate uses
|
||||
/// this; `truncate_history` (request-body sizing) uses `estimated_bytes`.
|
||||
pub fn context_pressure_bytes(&self) -> usize {
|
||||
self.size_with(ToolResultContent::context_pressure_bytes)
|
||||
}
|
||||
|
||||
fn size_with(&self, content_size: fn(&ToolResultContent) -> usize) -> usize {
|
||||
match self {
|
||||
Self::User(s) => s.len(),
|
||||
Self::Assistant {
|
||||
text,
|
||||
tool_calls,
|
||||
reasoning_details,
|
||||
} => {
|
||||
text.len()
|
||||
+ tool_calls
|
||||
.iter()
|
||||
.map(|c| {
|
||||
c.provider_id.len()
|
||||
+ c.name.len()
|
||||
+ serde_json::to_vec(&c.arguments)
|
||||
.map(|b| b.len())
|
||||
.unwrap_or(0)
|
||||
// `provider_extra` (e.g. a Gemini
|
||||
// `thoughtSignature`) is re-serialized into
|
||||
// every replayed call, so it counts toward the
|
||||
// request body and the context-pressure gate.
|
||||
+ serde_json::to_vec(&c.provider_extra)
|
||||
.map(|b| b.len())
|
||||
.unwrap_or(0)
|
||||
})
|
||||
.sum::<usize>()
|
||||
+ reasoning_details
|
||||
.as_ref()
|
||||
.and_then(|v| serde_json::to_vec(v).ok())
|
||||
.map(|b| b.len())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
Self::ToolResult(r) => {
|
||||
r.provider_id.len() + r.content.iter().map(content_size).sum::<usize>()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolCall {
|
||||
pub provider_id: String,
|
||||
pub name: String,
|
||||
pub arguments: Value,
|
||||
/// Fields the provider put on the tool call that we do not model, kept so
|
||||
/// the assistant turn can be replayed the way it arrived.
|
||||
///
|
||||
/// Gemini on the Databricks MLflow route returns a `thoughtSignature` per
|
||||
/// call and *requires* it echoed back: replaying without it fails the whole
|
||||
/// request with `Function call is missing a thought_signature in functionCall
|
||||
/// parts`. For an agent loop that lands on the very first tool call, so the
|
||||
/// model is unusable without this. Carrying whatever we did not model,
|
||||
/// rather than naming that one field, means the next provider with an opaque
|
||||
/// per-call token needs no change here.
|
||||
pub provider_extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolResult {
|
||||
pub provider_id: String,
|
||||
pub content: Vec<ToolResultContent>,
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
impl ToolResult {
|
||||
pub fn text(&self) -> String {
|
||||
self.content
|
||||
.iter()
|
||||
.map(ToolResultContent::as_text_lossy)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LlmResponse {
|
||||
pub text: String,
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
pub stop: ProviderStop,
|
||||
/// Total input tokens the provider reported for this request, or `None`
|
||||
/// if the response carried no usage. For Anthropic/Databricks this is the
|
||||
/// inclusive sum `input_tokens + cache_read_input_tokens +
|
||||
/// cache_creation_input_tokens` (plain `input_tokens` excludes cached
|
||||
/// tokens, so reading it alone would undercount). Used to gate handoff on
|
||||
/// the real token budget rather than a byte estimate.
|
||||
pub input_tokens: Option<u64>,
|
||||
/// The portion of `input_tokens` the provider served from its prompt cache,
|
||||
/// or `None` when the response reported no cache split. Providers bill this
|
||||
/// slice at a large discount (roughly 10x for both OpenAI and Anthropic),
|
||||
/// so a consumer that prices all of `input_tokens` at the full rate
|
||||
/// *overstates* cost — by a lot on an append-only agent loop, where most of
|
||||
/// each request is a prefix the provider already has.
|
||||
///
|
||||
/// This is a subset of `input_tokens`, never an addition to it: every
|
||||
/// provider we speak to reports an inclusive input total, so adding this
|
||||
/// would double-count.
|
||||
pub cached_input_tokens: Option<u64>,
|
||||
/// Output tokens the provider reported for this request, or `None` if the
|
||||
/// response carried no usage. Used to accumulate per-turn output counts
|
||||
/// for NIP-AM metric publishing.
|
||||
pub output_tokens: Option<u64>,
|
||||
/// Provider-reported total tokens for this request, or `None` when the
|
||||
/// provider does not report a genuine total. Present for OpenAI-shaped
|
||||
/// responses (`usage.total_tokens`). Always `None` for Anthropic, which
|
||||
/// reports only category counts; NIP-AM forbids summing categories into a
|
||||
/// total. Callers must not derive this by summing `input_tokens +
|
||||
/// output_tokens` — that is what the UI display approximation is for.
|
||||
pub total_tokens: Option<u64>,
|
||||
/// Reasoning/thinking content emitted by the model before its answer, if
|
||||
/// any. Non-empty when the provider returns extended-thinking tokens:
|
||||
///
|
||||
/// - Responses API: concatenated `summary[].text` from `type == "reasoning"` output items.
|
||||
/// - Anthropic: concatenated `thinking` from `type == "thinking"` content blocks.
|
||||
/// - OpenAI chat/completions: not exposed; always empty.
|
||||
///
|
||||
/// Empty string when the provider returned no reasoning content.
|
||||
pub reasoning: String,
|
||||
/// Raw `reasoning_details` array from an OpenRouter response, if present.
|
||||
/// Replayed on subsequent turns so the model can continue its chain-of-thought.
|
||||
/// `None` for all non-OpenRouter providers.
|
||||
pub reasoning_details: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum ProviderStop {
|
||||
EndTurn,
|
||||
ToolUse,
|
||||
MaxTokens,
|
||||
Refusal,
|
||||
Other,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolDef {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub input_schema: Value,
|
||||
}
|
||||
|
||||
/// Tri-state accumulator for provider-reported total tokens within one ACP turn.
|
||||
///
|
||||
/// Tracks whether every usage-bearing LLM response in the turn supplied a genuine
|
||||
/// provider total. Used to accumulate a reliable per-turn total and contribute to
|
||||
/// the session-cumulative total.
|
||||
///
|
||||
/// - `Unseen`: no usage-bearing response observed yet (initial state for each turn).
|
||||
/// - `Exact(n)`: every response so far reported a total; `n` is their sum.
|
||||
/// - `Unknown`: at least one response lacked a total — permanently poisoned for
|
||||
/// this turn. The session-cumulative also transitions to Unknown when any turn
|
||||
/// lands Unknown, and stays there until a new session resets it.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum TurnTotalState {
|
||||
#[default]
|
||||
Unseen,
|
||||
Exact(u64),
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl TurnTotalState {
|
||||
/// Add two exact token counts with overflow protection.
|
||||
///
|
||||
/// Returns `Exact(acc + n)` on success or `Unknown` on overflow.
|
||||
/// This is the single implementation of the checked-add / overflow-poisons
|
||||
/// contract; both `fold()` and `merge_session()` call this helper so a
|
||||
/// change to overflow semantics needs to be made in exactly one place.
|
||||
fn checked_exact_sum(acc: u64, n: u64) -> TurnTotalState {
|
||||
match acc.checked_add(n) {
|
||||
Some(sum) => TurnTotalState::Exact(sum),
|
||||
None => TurnTotalState::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
/// Fold one provider-reported total into the current state.
|
||||
///
|
||||
/// `total`: `Some(n)` when the provider included a genuine total on this
|
||||
/// response; `None` when it was absent (e.g. Anthropic, or an OpenAI
|
||||
/// response that omits usage). Absence of a total on any usage-bearing
|
||||
/// response poisons the whole turn.
|
||||
///
|
||||
/// Overflow is handled by `checked_exact_sum`: a saturated value would
|
||||
/// not be a genuine provider-reported total, so overflow → `Unknown`.
|
||||
pub fn fold(self, total: Option<u64>) -> TurnTotalState {
|
||||
match (self, total) {
|
||||
// Already poisoned — stays Unknown regardless.
|
||||
(TurnTotalState::Unknown, _) => TurnTotalState::Unknown,
|
||||
// No total from this response — poison the accumulator.
|
||||
(_, None) => TurnTotalState::Unknown,
|
||||
// First response with a total.
|
||||
(TurnTotalState::Unseen, Some(n)) => TurnTotalState::Exact(n),
|
||||
// Subsequent response — delegate to the shared checked-sum helper.
|
||||
(TurnTotalState::Exact(acc), Some(n)) => Self::checked_exact_sum(acc, n),
|
||||
}
|
||||
}
|
||||
|
||||
/// Merge a completed turn's total state into the session-cumulative state.
|
||||
///
|
||||
/// This is the turn→session boundary accumulation:
|
||||
/// - An `Unseen` turn (no usage-bearing responses) leaves the cumulative unchanged.
|
||||
/// - Any `Unknown` side poisons the session permanently.
|
||||
/// - Two `Exact` values are summed via `checked_exact_sum`; overflow → `Unknown`.
|
||||
///
|
||||
/// The checked-add logic lives in `checked_exact_sum`; both this function and
|
||||
/// `fold()` call that helper so overflow semantics are defined once.
|
||||
pub fn merge_session(self, turn: TurnTotalState) -> TurnTotalState {
|
||||
match (self, turn) {
|
||||
// Either side poisoned → session is poisoned.
|
||||
(TurnTotalState::Unknown, _) | (_, TurnTotalState::Unknown) => TurnTotalState::Unknown,
|
||||
// Turn had no usage-bearing responses → no change to cumulative.
|
||||
(acc, TurnTotalState::Unseen) => acc,
|
||||
// First exact turn — adopt its value.
|
||||
(TurnTotalState::Unseen, TurnTotalState::Exact(n)) => TurnTotalState::Exact(n),
|
||||
// Add to running exact sum — delegate to the shared checked-sum helper.
|
||||
(TurnTotalState::Exact(acc), TurnTotalState::Exact(n)) => {
|
||||
Self::checked_exact_sum(acc, n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Consume the exact value if present; `None` for `Unseen` or `Unknown`.
|
||||
pub fn exact_value(self) -> Option<u64> {
|
||||
match self {
|
||||
TurnTotalState::Exact(n) => Some(n),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The session-cumulative usage counters as of the START of a turn.
|
||||
///
|
||||
/// Copied out of the session under the lock when a turn begins and handed to
|
||||
/// `RunCtx` by value, so the run loop can emit a cumulative `usage_update`
|
||||
/// after every LLM round without reaching back into `App.sessions` (which it
|
||||
/// holds no handle to, and which is locked by the turn's own bookkeeping at
|
||||
/// both ends).
|
||||
///
|
||||
/// This exists so that usage is durable *during* a turn rather than only after
|
||||
/// it. The counters a turn accrues live in the prompt task's stack frame until
|
||||
/// the turn returns; a process killed mid-turn takes them with it and the
|
||||
/// tokens are billed by the provider but recorded nowhere. That is not
|
||||
/// hypothetical — it silently under-reported a long-horizon benchmark's cost by
|
||||
/// several-fold, because every phase of a `continue_until_timeout` run is
|
||||
/// terminated mid-turn by design.
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct SessionUsageBaseline {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
/// The cache-served subset of `input_tokens`, not an addition to it.
|
||||
pub cached_input_tokens: u64,
|
||||
pub total_state: TurnTotalState,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum StopReason {
|
||||
EndTurn,
|
||||
Cancelled,
|
||||
MaxTokens,
|
||||
MaxTurnRequests,
|
||||
Refusal,
|
||||
}
|
||||
|
||||
impl StopReason {
|
||||
pub fn as_wire(self) -> &'static str {
|
||||
match self {
|
||||
Self::EndTurn => "end_turn",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::MaxTokens => "max_tokens",
|
||||
Self::MaxTurnRequests => "max_turn_requests",
|
||||
Self::Refusal => "refusal",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct McpServerStdio {
|
||||
pub name: String,
|
||||
pub command: String,
|
||||
#[serde(default)]
|
||||
pub args: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub env: Vec<EnvVar>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct EnvVar {
|
||||
pub name: String,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
ResourceLink {
|
||||
uri: String,
|
||||
},
|
||||
#[serde(other)]
|
||||
Unsupported,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum AgentError {
|
||||
InvalidParams(String),
|
||||
Llm(String),
|
||||
LlmAuth(String),
|
||||
LlmModelNotFound(String),
|
||||
/// The provider rejected the request because the input exceeded the
|
||||
/// model's context window (an HTTP 400 whose body names a context-length
|
||||
/// overflow). Typed rather than folded into [`Self::Llm`] because the
|
||||
/// agent loop treats it as a *recovery* signal, not a terminal error: it
|
||||
/// is the only ground-truth indication that history must shrink, needing
|
||||
/// no window estimate that could itself be miscalibrated.
|
||||
///
|
||||
/// Classified where the HTTP status and body are still separate values, so
|
||||
/// the loop never has to sniff a formatted string — by the time an error
|
||||
/// leaves `Llm::complete` it has already been decorated with the model
|
||||
/// name.
|
||||
LlmContextExceeded(String),
|
||||
/// The provider explicitly rejected image content for the selected model.
|
||||
/// Kept distinct so the agent loop can remove the unsupported image from
|
||||
/// replayed history and give the model a recoverable tool error.
|
||||
UnsupportedImageInput(String),
|
||||
Mcp(String),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for AgentError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::InvalidParams(s) => write!(f, "invalid params: {s}"),
|
||||
Self::Llm(s) => write!(f, "llm: {s}"),
|
||||
Self::LlmAuth(s) => write!(f, "llm auth: {s}"),
|
||||
Self::LlmModelNotFound(s) => write!(f, "llm model not found: {s}"),
|
||||
Self::LlmContextExceeded(s) => write!(f, "llm context exceeded: {s}"),
|
||||
Self::UnsupportedImageInput(s) => write!(f, "llm image input unsupported: {s}"),
|
||||
Self::Mcp(s) => write!(f, "mcp: {s}"),
|
||||
Self::Cancelled => write!(f, "cancelled"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for AgentError {}
|
||||
|
||||
impl AgentError {
|
||||
pub fn json_rpc_code(&self) -> i32 {
|
||||
match self {
|
||||
Self::InvalidParams(_) => -32602,
|
||||
Self::LlmAuth(_) => -32001,
|
||||
Self::LlmModelNotFound(_) => -32002,
|
||||
_ => -32000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn clamp(mut s: String, max: usize) -> String {
|
||||
if s.len() <= max {
|
||||
return s;
|
||||
}
|
||||
const MARKER: &str = "\n[truncated]";
|
||||
let budget = max.saturating_sub(MARKER.len());
|
||||
let mut cut = budget;
|
||||
while cut > 0 && !s.is_char_boundary(cut) {
|
||||
cut -= 1;
|
||||
}
|
||||
s.truncate(cut);
|
||||
if max >= MARKER.len() {
|
||||
s.push_str(MARKER);
|
||||
}
|
||||
s
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn image_item(base64_len: usize) -> HistoryItem {
|
||||
HistoryItem::ToolResult(ToolResult {
|
||||
provider_id: "call_1".into(),
|
||||
content: vec![ToolResultContent::Image {
|
||||
data: "A".repeat(base64_len),
|
||||
mime_type: "image/png".into(),
|
||||
}],
|
||||
is_error: false,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_estimated_bytes_is_real_wire_size() {
|
||||
// `truncate_history` relies on this to keep the request body under
|
||||
// `max_history_bytes`, so an image must report its full base64 length.
|
||||
let img = ToolResultContent::Image {
|
||||
data: "A".repeat(3_000_000),
|
||||
mime_type: "image/png".into(),
|
||||
};
|
||||
assert_eq!(img.estimated_bytes(), 3_000_000 + "image/png".len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_context_pressure_is_token_equivalent_not_base64_len() {
|
||||
// The handoff gate must charge an image its visual-token equivalent,
|
||||
// not its base64 length — otherwise one screenshot trips the gate.
|
||||
let img = ToolResultContent::Image {
|
||||
data: "A".repeat(3_000_000),
|
||||
mime_type: "image/png".into(),
|
||||
};
|
||||
assert_eq!(
|
||||
img.context_pressure_bytes(),
|
||||
IMAGE_CONTEXT_TOKEN_EQUIV + "image/png".len()
|
||||
);
|
||||
// And it must be independent of the (huge) base64 payload length.
|
||||
let bigger = ToolResultContent::Image {
|
||||
data: "A".repeat(10_000_000),
|
||||
mime_type: "image/png".into(),
|
||||
};
|
||||
assert_eq!(
|
||||
img.context_pressure_bytes(),
|
||||
bigger.context_pressure_bytes()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_image_does_not_trip_default_handoff_threshold() {
|
||||
// Regression: a single ~3.1M-base64-byte `view_image` result on an
|
||||
// otherwise-empty history must NOT exceed the default pre-usage
|
||||
// handoff cap. The gate's byte-fallback threshold with the shipped
|
||||
// defaults (max_context_tokens=200_000, max_output_tokens=32_768) is
|
||||
// min(200_000*9/10, 200_000-32_768) = 167_232 "bytes". Before the fix
|
||||
// this item counted ~3.1M and tripped instantly.
|
||||
let item = image_item(3_118_884);
|
||||
const DEFAULT_PRE_USAGE_THRESHOLD: usize = 167_232;
|
||||
assert!(
|
||||
item.context_pressure_bytes() <= DEFAULT_PRE_USAGE_THRESHOLD,
|
||||
"one image charged {} bytes of context pressure, over the {} threshold",
|
||||
item.context_pressure_bytes(),
|
||||
DEFAULT_PRE_USAGE_THRESHOLD
|
||||
);
|
||||
// The real wire size, by contrast, is still the full base64 payload.
|
||||
assert!(item.estimated_bytes() >= 3_118_884);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assistant_size_counts_provider_extra() {
|
||||
// A Gemini `thoughtSignature` rides the wire on every replayed call, so
|
||||
// both size measures must see it — otherwise `truncate_history` and the
|
||||
// handoff gate under-count and let the real request exceed the budget.
|
||||
let mut extra = Map::new();
|
||||
extra.insert("thoughtSignature".into(), Value::String("S".repeat(500)));
|
||||
let with_extra = HistoryItem::Assistant {
|
||||
text: String::new(),
|
||||
tool_calls: vec![ToolCall {
|
||||
provider_id: "id".into(),
|
||||
name: "t".into(),
|
||||
arguments: Value::Null,
|
||||
provider_extra: extra,
|
||||
}],
|
||||
reasoning_details: None,
|
||||
};
|
||||
let without_extra = HistoryItem::Assistant {
|
||||
text: String::new(),
|
||||
tool_calls: vec![ToolCall {
|
||||
provider_id: "id".into(),
|
||||
name: "t".into(),
|
||||
arguments: Value::Null,
|
||||
provider_extra: Map::new(),
|
||||
}],
|
||||
reasoning_details: None,
|
||||
};
|
||||
assert!(with_extra.estimated_bytes() > without_extra.estimated_bytes() + 500);
|
||||
assert_eq!(
|
||||
with_extra.estimated_bytes(),
|
||||
with_extra.context_pressure_bytes(),
|
||||
"provider_extra is text, so both measures must agree"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_content_size_is_identical_for_both_measures() {
|
||||
// Only images diverge; text must size the same under both paths.
|
||||
let text = ToolResultContent::Text("hello world".into());
|
||||
assert_eq!(text.estimated_bytes(), text.context_pressure_bytes());
|
||||
let item = HistoryItem::User("a user message".into());
|
||||
assert_eq!(item.estimated_bytes(), item.context_pressure_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod turn_total_state_tests {
|
||||
use super::TurnTotalState;
|
||||
|
||||
// ── TurnTotalState::fold ───────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn fold_first_response_with_total_becomes_exact() {
|
||||
let state = TurnTotalState::Unseen;
|
||||
assert_eq!(state.fold(Some(100)), TurnTotalState::Exact(100));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fold_first_response_without_total_becomes_unknown() {
|
||||
// Missing total on any usage-bearing response poisons the turn.
|
||||
let state = TurnTotalState::Unseen;
|
||||
assert_eq!(state.fold(None), TurnTotalState::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_provider_rounds_all_with_totals_sum_correctly() {
|
||||
// Multiple rounds all reporting a genuine total → Exact with their sum.
|
||||
let state = TurnTotalState::Unseen;
|
||||
let state = state.fold(Some(100));
|
||||
let state = state.fold(Some(50));
|
||||
let state = state.fold(Some(75));
|
||||
assert_eq!(state, TurnTotalState::Exact(225));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mixed_present_and_missing_totals_within_one_turn_poisons_accumulator() {
|
||||
// First round has a total, second does not → Unknown (permanently poisoned).
|
||||
let state = TurnTotalState::Unseen;
|
||||
let state = state.fold(Some(100)); // Exact(100)
|
||||
let state = state.fold(None); // Missing → Unknown
|
||||
assert_eq!(state, TurnTotalState::Unknown);
|
||||
// Further rounds with totals don't un-poison.
|
||||
let state = state.fold(Some(50));
|
||||
assert_eq!(state, TurnTotalState::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_stays_unknown_regardless_of_subsequent_totals() {
|
||||
// Once poisoned, no subsequent total can recover the state.
|
||||
let state = TurnTotalState::Unknown;
|
||||
assert_eq!(state.fold(Some(999)), TurnTotalState::Unknown);
|
||||
assert_eq!(state.fold(None), TurnTotalState::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exact_value_returns_some_only_for_exact_variant() {
|
||||
assert_eq!(TurnTotalState::Unseen.exact_value(), None);
|
||||
assert_eq!(TurnTotalState::Unknown.exact_value(), None);
|
||||
assert_eq!(TurnTotalState::Exact(42).exact_value(), Some(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_is_unseen() {
|
||||
let state: TurnTotalState = Default::default();
|
||||
assert_eq!(state, TurnTotalState::Unseen);
|
||||
}
|
||||
|
||||
// ── overflow: fold ─────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn fold_overflow_poisons_turn_not_saturates() {
|
||||
// u64::MAX + 1 would saturate; checked_add must poison instead.
|
||||
let state = TurnTotalState::Exact(u64::MAX);
|
||||
assert_eq!(
|
||||
state.fold(Some(1)),
|
||||
TurnTotalState::Unknown,
|
||||
"overflow in fold() must produce Unknown, not Exact(u64::MAX)"
|
||||
);
|
||||
}
|
||||
|
||||
// ── TurnTotalState::merge_session ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn merge_session_unseen_turn_leaves_cumulative_unchanged() {
|
||||
// An Unseen turn (no usage-bearing responses) must not alter the cumulative.
|
||||
assert_eq!(
|
||||
TurnTotalState::Exact(100).merge_session(TurnTotalState::Unseen),
|
||||
TurnTotalState::Exact(100),
|
||||
);
|
||||
assert_eq!(
|
||||
TurnTotalState::Unseen.merge_session(TurnTotalState::Unseen),
|
||||
TurnTotalState::Unseen,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_session_exact_turn_adds_to_exact_cumulative() {
|
||||
assert_eq!(
|
||||
TurnTotalState::Exact(100).merge_session(TurnTotalState::Exact(50)),
|
||||
TurnTotalState::Exact(150),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_session_first_exact_turn_from_unseen_adopts_value() {
|
||||
assert_eq!(
|
||||
TurnTotalState::Unseen.merge_session(TurnTotalState::Exact(200)),
|
||||
TurnTotalState::Exact(200),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_session_unknown_turn_poisons_cumulative_permanently() {
|
||||
assert_eq!(
|
||||
TurnTotalState::Exact(100).merge_session(TurnTotalState::Unknown),
|
||||
TurnTotalState::Unknown,
|
||||
);
|
||||
// Poisoned session stays poisoned even with Unseen turn.
|
||||
assert_eq!(
|
||||
TurnTotalState::Unknown.merge_session(TurnTotalState::Unseen),
|
||||
TurnTotalState::Unknown,
|
||||
);
|
||||
// Poisoned session stays poisoned even with another Exact turn.
|
||||
assert_eq!(
|
||||
TurnTotalState::Unknown.merge_session(TurnTotalState::Exact(999)),
|
||||
TurnTotalState::Unknown,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_session_overflow_poisons_not_saturates() {
|
||||
// Overflow at the session boundary must also produce Unknown.
|
||||
assert_eq!(
|
||||
TurnTotalState::Exact(u64::MAX).merge_session(TurnTotalState::Exact(1)),
|
||||
TurnTotalState::Unknown,
|
||||
"overflow in merge_session() must produce Unknown, not Exact(u64::MAX)"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,335 @@
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::types::{ContentBlock, McpServerStdio};
|
||||
|
||||
pub const PARSE_ERROR: i32 = -32700;
|
||||
pub const INVALID_REQUEST: i32 = -32600;
|
||||
pub const METHOD_NOT_FOUND: i32 = -32601;
|
||||
pub const INVALID_PARAMS: i32 = -32602;
|
||||
|
||||
pub enum WireMsg {
|
||||
Notify(Value),
|
||||
}
|
||||
|
||||
pub type WireSender = mpsc::Sender<WireMsg>;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Inbound {
|
||||
Request {
|
||||
id: Value,
|
||||
method: String,
|
||||
params: Value,
|
||||
},
|
||||
Notification {
|
||||
method: String,
|
||||
params: Value,
|
||||
},
|
||||
Ignored,
|
||||
Invalid {
|
||||
id: Value,
|
||||
code: i32,
|
||||
message: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct InitializeParams {
|
||||
#[serde(rename = "protocolVersion")]
|
||||
pub protocol_version: u32,
|
||||
#[serde(default, rename = "clientCapabilities")]
|
||||
pub _client_capabilities: Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionNewParams {
|
||||
pub cwd: String,
|
||||
#[serde(default)]
|
||||
pub mcp_servers: Vec<McpServerStdio>,
|
||||
#[serde(default)]
|
||||
pub system_prompt: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionPromptParams {
|
||||
pub session_id: String,
|
||||
pub prompt: Vec<ContentBlock>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionCancelParams {
|
||||
pub session_id: String,
|
||||
}
|
||||
|
||||
/// Params for goose's non-standard `_goose/unstable/session/steer` request:
|
||||
/// inject user input into the *currently active* prompt without starting a new
|
||||
/// one. `expected_run_id` must match the run id buzz-agent advertised via
|
||||
/// `params.update._meta.goose.activeRunId` on a `session/update`, so a steer
|
||||
/// can't race a turn that already ended or hasn't started.
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionSteerParams {
|
||||
pub session_id: String,
|
||||
#[serde(default)]
|
||||
pub prompt: Vec<ContentBlock>,
|
||||
pub expected_run_id: String,
|
||||
}
|
||||
|
||||
/// Params for `session/set_model`: override the active model for an existing
|
||||
/// session without respawning. Applied immediately; subsequent prompts on this
|
||||
/// session use `model_id` instead of the configured `BUZZ_AGENT_MODEL`.
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionSetModelParams {
|
||||
pub session_id: String,
|
||||
pub model_id: String,
|
||||
}
|
||||
|
||||
pub fn classify(msg: &Value) -> Inbound {
|
||||
if !msg.is_object() || msg.get("jsonrpc").and_then(Value::as_str) != Some("2.0") {
|
||||
return Inbound::Invalid {
|
||||
id: msg.get("id").cloned().unwrap_or(Value::Null),
|
||||
code: INVALID_REQUEST,
|
||||
message: "jsonrpc: missing or invalid version".into(),
|
||||
};
|
||||
}
|
||||
let id = msg.get("id").cloned();
|
||||
let method = msg.get("method").and_then(Value::as_str).map(str::to_owned);
|
||||
let params = msg.get("params").cloned().unwrap_or(Value::Null);
|
||||
|
||||
match (method, id) {
|
||||
(Some(m), Some(id)) => Inbound::Request {
|
||||
id,
|
||||
method: m,
|
||||
params,
|
||||
},
|
||||
(Some(m), None) => Inbound::Notification { method: m, params },
|
||||
// Bare responses (id present, no method) are unexpected — buzz-agent
|
||||
// does not issue requests to the client. Ignore silently.
|
||||
(None, Some(_)) => Inbound::Ignored,
|
||||
(None, None) => Inbound::Invalid {
|
||||
id: Value::Null,
|
||||
code: INVALID_REQUEST,
|
||||
message: "jsonrpc: missing method and id".into(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ok(id: Value, result: Value) -> Value {
|
||||
json!({ "jsonrpc": "2.0", "id": id, "result": result })
|
||||
}
|
||||
|
||||
pub fn err(id: Value, code: i32, message: &str) -> Value {
|
||||
json!({ "jsonrpc": "2.0", "id": id, "error": { "code": code, "message": message } })
|
||||
}
|
||||
|
||||
pub fn session_update(sid: &str, update: Value) -> Value {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "session/update",
|
||||
"params": { "sessionId": sid, "update": update },
|
||||
})
|
||||
}
|
||||
|
||||
/// A `_goose/unstable/session/update` notification — the separate top-level
|
||||
/// method goose uses for custom usage and status events. Used by buzz-agent
|
||||
/// to emit the `usage_update` payload so buzz-acp's `UsageTracker` can treat
|
||||
/// buzz-agent and goose symmetrically.
|
||||
pub fn goose_session_update(sid: &str, update: Value) -> Value {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "_goose/unstable/session/update",
|
||||
"params": { "sessionId": sid, "update": update },
|
||||
})
|
||||
}
|
||||
|
||||
/// Build the `usage_update` payload for a `_goose/unstable/session/update`.
|
||||
///
|
||||
/// Shared by the two places that report usage — after each LLM round inside a
|
||||
/// turn, and once more when the turn completes — so the wire shape cannot drift
|
||||
/// between them. A consumer takes the high-water mark per session, so the
|
||||
/// mid-turn payloads are supersets of each other and the final one wins; a
|
||||
/// divergence in field names or units between the two call sites would instead
|
||||
/// show up as tokens silently vanishing, which is the failure this reporting
|
||||
/// exists to prevent.
|
||||
///
|
||||
/// All counts are SESSION-cumulative, matching goose, so buzz-acp's
|
||||
/// `UsageTracker` can compute per-turn deltas symmetrically for both agents.
|
||||
pub fn usage_update_payload(
|
||||
accumulated_input_tokens: u64,
|
||||
accumulated_output_tokens: u64,
|
||||
accumulated_cached_input_tokens: u64,
|
||||
accumulated_total: crate::types::TurnTotalState,
|
||||
model: &str,
|
||||
) -> Value {
|
||||
let mut update = json!({
|
||||
"sessionUpdate": "usage_update",
|
||||
// used: total tokens as a context-usage proxy;
|
||||
// contextLimit: 0 (buzz-agent has no context limit tracking).
|
||||
"used": accumulated_input_tokens.saturating_add(accumulated_output_tokens),
|
||||
"contextLimit": 0u64,
|
||||
"accumulatedInputTokens": accumulated_input_tokens,
|
||||
"accumulatedOutputTokens": accumulated_output_tokens,
|
||||
// A subset of accumulatedInputTokens, not an addition to it. Extends
|
||||
// goose's usage_update shape; a consumer that does not know the field
|
||||
// ignores it and prices exactly as it did before.
|
||||
"accumulatedCachedInputTokens": accumulated_cached_input_tokens,
|
||||
"model": model,
|
||||
});
|
||||
// Only when the cumulative is exactly known — never when Unseen (no total
|
||||
// ever observed) or Unknown (at least one turn lacked a total). A goose
|
||||
// consumer that doesn't recognise the field ignores it.
|
||||
if let Some(total) = accumulated_total.exact_value() {
|
||||
update["accumulatedTotalTokens"] = json!(total);
|
||||
}
|
||||
update
|
||||
}
|
||||
|
||||
/// A `session/update` notification carrying a `update._meta.goose.<key>` field.
|
||||
/// Used to advertise `activeRunId` (so steer-capable clients can target the
|
||||
/// in-flight run) and `queuedSteer` (so they can correlate an accepted steer
|
||||
/// with the chunk that later picks it up) — matching goose's wire layout where
|
||||
/// `_meta` is nested inside the `update` object (per the ACP `SessionInfoUpdate`
|
||||
/// schema), not alongside it at the params level.
|
||||
pub fn session_update_with_goose_meta(sid: &str, update: Value, goose_meta: Value) -> Value {
|
||||
let mut update = update;
|
||||
update["_meta"] = json!({ "goose": goose_meta });
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "session/update",
|
||||
"params": {
|
||||
"sessionId": sid,
|
||||
"update": update,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn send(wire: &WireSender, msg: Value) {
|
||||
let _ = wire.send(WireMsg::Notify(msg)).await;
|
||||
}
|
||||
|
||||
pub async fn read_bounded_line<R: AsyncBufRead + Unpin>(
|
||||
stdin: &mut R,
|
||||
max: usize,
|
||||
) -> std::io::Result<Option<String>> {
|
||||
let mut buf: Vec<u8> = Vec::new();
|
||||
loop {
|
||||
let chunk = stdin.fill_buf().await?;
|
||||
if chunk.is_empty() {
|
||||
if !buf.is_empty() {
|
||||
tracing::error!(
|
||||
"io: unterminated frame at EOF ({} bytes dropped)",
|
||||
buf.len()
|
||||
);
|
||||
}
|
||||
return Ok(None);
|
||||
}
|
||||
let take = chunk
|
||||
.iter()
|
||||
.position(|b| *b == b'\n')
|
||||
.map_or(chunk.len(), |i| i + 1);
|
||||
if buf.len().saturating_add(take) > max {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
format!("io: line exceeds max ({max} bytes)"),
|
||||
));
|
||||
}
|
||||
buf.extend_from_slice(&chunk[..take]);
|
||||
stdin.consume(take);
|
||||
if buf.ends_with(b"\n") {
|
||||
buf.pop();
|
||||
if buf.ends_with(b"\r") {
|
||||
buf.pop();
|
||||
}
|
||||
match String::from_utf8(buf) {
|
||||
Ok(s) => return Ok(Some(s)),
|
||||
Err(_) => {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"io: frame contains invalid UTF-8",
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn writer_task(mut rx: mpsc::Receiver<WireMsg>) {
|
||||
let mut stdout = tokio::io::stdout();
|
||||
while let Some(msg) = rx.recv().await {
|
||||
let WireMsg::Notify(v) = msg;
|
||||
let mut s = match serde_json::to_string(&v) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
tracing::error!("io: serialize: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
s.push('\n');
|
||||
if stdout.write_all(s.as_bytes()).await.is_err() {
|
||||
return;
|
||||
}
|
||||
let _ = stdout.flush().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn session_new_params_deserializes_system_prompt() {
|
||||
let json = serde_json::json!({
|
||||
"cwd": "/tmp/test",
|
||||
"mcpServers": [],
|
||||
"systemPrompt": "You are a helpful agent."
|
||||
});
|
||||
let params: SessionNewParams = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(params.cwd, "/tmp/test");
|
||||
assert_eq!(
|
||||
params.system_prompt.as_deref(),
|
||||
Some("You are a helpful agent.")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_new_params_system_prompt_defaults_to_none() {
|
||||
let json = serde_json::json!({
|
||||
"cwd": "/tmp/test",
|
||||
"mcpServers": []
|
||||
});
|
||||
let params: SessionNewParams = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(params.cwd, "/tmp/test");
|
||||
assert!(params.system_prompt.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_new_params_ignores_unknown_fields() {
|
||||
// Backward compat: old agents with new harness — unknown fields are ignored.
|
||||
let json = serde_json::json!({
|
||||
"cwd": "/tmp/test",
|
||||
"mcpServers": [],
|
||||
"unknownField": "should be ignored"
|
||||
});
|
||||
let params: SessionNewParams = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(params.cwd, "/tmp/test");
|
||||
assert!(params.system_prompt.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_new_params_empty_string_system_prompt() {
|
||||
// An explicit empty string is distinct from absent — deserializes to Some("").
|
||||
let json = serde_json::json!({
|
||||
"cwd": "/tmp/test",
|
||||
"mcpServers": [],
|
||||
"systemPrompt": ""
|
||||
});
|
||||
let params: SessionNewParams = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(params.system_prompt, Some(String::new()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,334 @@
|
||||
//! Tiny fake MCP server for integration tests.
|
||||
//!
|
||||
//! Reads JSON-RPC line frames on stdin and replies on stdout. Driven by
|
||||
//! environment variables so tests can simulate misbehavior:
|
||||
//!
|
||||
//! FAKE_MCP_HANG_INIT=1 — never reply to `initialize` (init timeout)
|
||||
//! FAKE_MCP_HANG_TOOLS=1 — never reply to `tools/list` (list timeout)
|
||||
//! FAKE_MCP_TOOL_COUNT=N — return N tools (default: 1)
|
||||
//! FAKE_MCP_HUGE_DESC=1 — every tool description is 100 KB
|
||||
//! FAKE_MCP_DESC_SIZE=N — every tool description is N bytes (overrides HUGE_DESC)
|
||||
//! FAKE_MCP_TOOL_DELAY=N — `tools/call` sleeps N seconds before replying
|
||||
//! (use a large value, e.g. 999, to simulate hang)
|
||||
//! FAKE_MCP_RESULT_SIZE=N — `tools/call` returns an N-byte text result
|
||||
//! (default: the literal "ok"); grows history
|
||||
//! FAKE_MCP_IMAGE_RESULT=1 — `tools/call` returns text plus a PNG image block
|
||||
//! FAKE_MCP_PID_FILE=path — write the child PID to `path` on startup
|
||||
//! (for tests that want to verify the child died)
|
||||
//! FAKE_MCP_SPAWN_GRANDCHILD=1
|
||||
//! — on `tools/call`, spawn a `sleep 999`
|
||||
//! grandchild before hanging. Its PID is
|
||||
//! written to FAKE_MCP_GRANDCHILD_PID_FILE
|
||||
//! so a test can verify the entire process
|
||||
//! tree dies on timeout.
|
||||
//! FAKE_MCP_GRANDCHILD_PID_FILE=path
|
||||
//! — path to write the grandchild PID to.
|
||||
//! FAKE_MCP_STOP_HOOK=1 — expose a `_Stop` hook tool
|
||||
//! FAKE_MCP_STOP_TEXT=text — `_Stop` returns this text (default: "keep going")
|
||||
//! FAKE_MCP_STOP_DELAY=N — `_Stop` sleeps N seconds before replying
|
||||
//! (use a large value to simulate hang/timeout)
|
||||
//! FAKE_MCP_STOP_COUNT=N — `_Stop` returns STOP_TEXT for the first N
|
||||
//! invocations; empty string thereafter. If
|
||||
//! unset, every call returns STOP_TEXT.
|
||||
//! FAKE_MCP_POSTCOMPACT_HOOK=1
|
||||
//! — expose a `_PostCompact` hook tool
|
||||
//! FAKE_MCP_POSTCOMPACT_TEXT=text
|
||||
//! — `_PostCompact` returns this (default: "")
|
||||
//! FAKE_MCP_SHELL_TOOL=1 — expose a tool whose bare name is `shell`
|
||||
//! (registered as `<server>__shell`), taking a
|
||||
//! `command` string. Lets a test drive the
|
||||
//! reply guard's recognition of a real,
|
||||
//! registered shell tool.
|
||||
|
||||
use std::io::{BufRead, Write};
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn env_flag(k: &str) -> bool {
|
||||
std::env::var(k).map(|v| v != "0").unwrap_or(false)
|
||||
}
|
||||
|
||||
fn env_usize(k: &str, default: usize) -> usize {
|
||||
std::env::var(k)
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn env_u64(k: &str, default: u64) -> u64 {
|
||||
std::env::var(k)
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn write_response(id: Value, result: Value) {
|
||||
let msg = json!({ "jsonrpc": "2.0", "id": id, "result": result });
|
||||
let mut s = serde_json::to_string(&msg).expect("serialize");
|
||||
s.push('\n');
|
||||
let mut out = std::io::stdout().lock();
|
||||
out.write_all(s.as_bytes()).expect("write");
|
||||
out.flush().expect("flush");
|
||||
}
|
||||
|
||||
fn hang_forever() -> ! {
|
||||
loop {
|
||||
std::thread::sleep(std::time::Duration::from_secs(60));
|
||||
}
|
||||
}
|
||||
|
||||
fn make_tools(
|
||||
count: usize,
|
||||
desc: &str,
|
||||
include_stop_hook: bool,
|
||||
include_post_compact_hook: bool,
|
||||
include_shell_tool: bool,
|
||||
) -> Vec<Value> {
|
||||
let mut tools: Vec<Value> = (0..count)
|
||||
.map(|i| {
|
||||
json!({
|
||||
"name": format!("tool_{i}"),
|
||||
"description": desc,
|
||||
"inputSchema": { "type": "object", "properties": {} },
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
if include_stop_hook {
|
||||
tools.push(json!({
|
||||
"name": "_Stop",
|
||||
"description": "stop hook",
|
||||
"inputSchema": { "type": "object", "properties": {} },
|
||||
}));
|
||||
}
|
||||
if include_post_compact_hook {
|
||||
tools.push(json!({
|
||||
"name": "_PostCompact",
|
||||
"description": "post compact hook",
|
||||
"inputSchema": { "type": "object", "properties": {} },
|
||||
}));
|
||||
}
|
||||
if include_shell_tool {
|
||||
tools.push(json!({
|
||||
"name": "shell",
|
||||
"description": "run a shell command",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": { "command": { "type": "string" } },
|
||||
"required": ["command"],
|
||||
},
|
||||
}));
|
||||
}
|
||||
tools
|
||||
}
|
||||
|
||||
fn main() {
|
||||
// Optional: write our own PID so a test can later check the process is gone.
|
||||
if let Ok(path) = std::env::var("FAKE_MCP_PID_FILE") {
|
||||
let pid = std::process::id().to_string();
|
||||
let _ = std::fs::write(&path, pid);
|
||||
}
|
||||
|
||||
let hang_init = env_flag("FAKE_MCP_HANG_INIT");
|
||||
let hang_tools = env_flag("FAKE_MCP_HANG_TOOLS");
|
||||
let tool_count = env_usize("FAKE_MCP_TOOL_COUNT", 1);
|
||||
// FAKE_MCP_DESC_SIZE wins over FAKE_MCP_HUGE_DESC when set.
|
||||
let desc: String = if let Some(n) = std::env::var("FAKE_MCP_DESC_SIZE")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<usize>().ok())
|
||||
{
|
||||
"x".repeat(n)
|
||||
} else if env_flag("FAKE_MCP_HUGE_DESC") {
|
||||
"x".repeat(100_000)
|
||||
} else {
|
||||
"fake tool".to_owned()
|
||||
};
|
||||
let tool_delay_secs = env_u64("FAKE_MCP_TOOL_DELAY", 0);
|
||||
// Tool-call result text size in bytes (default: the literal "ok"). Lets a
|
||||
// test grow session history by a controlled amount via a tool result.
|
||||
let result_size = env_u64("FAKE_MCP_RESULT_SIZE", 0) as usize;
|
||||
let stop_hook = env_flag("FAKE_MCP_STOP_HOOK");
|
||||
let stop_text = std::env::var("FAKE_MCP_STOP_TEXT").unwrap_or_else(|_| "keep going".to_owned());
|
||||
let stop_delay_secs = env_u64("FAKE_MCP_STOP_DELAY", 0);
|
||||
// 0 means "unset" → unlimited; any positive value caps the number of
|
||||
// calls that return STOP_TEXT before flipping to empty string.
|
||||
let stop_count_limit: usize = env_usize("FAKE_MCP_STOP_COUNT", usize::MAX);
|
||||
let mut stop_calls_seen: usize = 0;
|
||||
let post_compact_hook = env_flag("FAKE_MCP_POSTCOMPACT_HOOK");
|
||||
let shell_tool = env_flag("FAKE_MCP_SHELL_TOOL");
|
||||
let post_compact_text = std::env::var("FAKE_MCP_POSTCOMPACT_TEXT").unwrap_or_default();
|
||||
|
||||
// Use a channel-based stdin reader so notifications (which carry no id)
|
||||
// are captured even while the main thread is sleeping during a tool call.
|
||||
let cancel_log_path = std::env::var("FAKE_MCP_CANCEL_LOG").ok();
|
||||
let (tx, rx) = std::sync::mpsc::channel::<(Value, Option<Value>)>();
|
||||
let cancel_log_for_thread = cancel_log_path.clone();
|
||||
std::thread::spawn(move || {
|
||||
let stdin = std::io::stdin();
|
||||
let lines = stdin.lock().lines();
|
||||
for line in lines {
|
||||
let line = match line {
|
||||
Ok(l) => l,
|
||||
Err(_) => return,
|
||||
};
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
let msg: Value = match serde_json::from_str(&line) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let method = msg.get("method").and_then(Value::as_str).unwrap_or("");
|
||||
let id = msg.get("id").cloned();
|
||||
// Notifications carry no id. Log cancellations if configured.
|
||||
if id.is_none() || id == Some(Value::Null) {
|
||||
if method == "notifications/cancelled" {
|
||||
if let Some(ref path) = cancel_log_for_thread {
|
||||
use std::io::Write as _;
|
||||
if let Ok(mut f) = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(path)
|
||||
{
|
||||
let _ = writeln!(f, "{}", line.trim());
|
||||
}
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
// Send requests (with id) to the main processing loop.
|
||||
let _ = tx.send((msg, id));
|
||||
}
|
||||
});
|
||||
|
||||
while let Ok((msg, id_opt)) = rx.recv() {
|
||||
let method = msg.get("method").and_then(Value::as_str).unwrap_or("");
|
||||
let id = id_opt.unwrap_or(Value::Null);
|
||||
|
||||
match method {
|
||||
"initialize" => {
|
||||
if hang_init {
|
||||
hang_forever();
|
||||
}
|
||||
write_response(
|
||||
id,
|
||||
json!({
|
||||
"protocolVersion": "2025-06-18",
|
||||
"capabilities": { "tools": {} },
|
||||
"serverInfo": { "name": "fake-mcp", "version": "0.0.0" },
|
||||
}),
|
||||
);
|
||||
}
|
||||
"tools/list" => {
|
||||
if hang_tools {
|
||||
hang_forever();
|
||||
}
|
||||
write_response(
|
||||
id,
|
||||
json!({
|
||||
"tools": make_tools(
|
||||
tool_count,
|
||||
&desc,
|
||||
stop_hook,
|
||||
post_compact_hook,
|
||||
shell_tool,
|
||||
)
|
||||
}),
|
||||
);
|
||||
}
|
||||
"tools/call" => {
|
||||
// Signal that the request was received (for tests that
|
||||
// need to wait until the call is in-flight before cancelling).
|
||||
// Write the request id so tests can correlate with cancel.
|
||||
if let Ok(path) = std::env::var("FAKE_MCP_CALL_RECEIVED") {
|
||||
let id_str = serde_json::to_string(&id).unwrap_or_else(|_| "?".into());
|
||||
let _ = std::fs::write(&path, id_str);
|
||||
}
|
||||
let called_name = msg
|
||||
.get("params")
|
||||
.and_then(|p| p.get("name"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("");
|
||||
// Optionally spawn a long-sleeping grandchild so the test
|
||||
// can verify process-group killing reaches the whole tree.
|
||||
if env_flag("FAKE_MCP_SPAWN_GRANDCHILD") {
|
||||
let child = std::process::Command::new("sleep")
|
||||
.arg("999")
|
||||
.spawn()
|
||||
.expect("spawn grandchild");
|
||||
if let Ok(path) = std::env::var("FAKE_MCP_GRANDCHILD_PID_FILE") {
|
||||
let _ = std::fs::write(&path, child.id().to_string());
|
||||
}
|
||||
std::mem::forget(child);
|
||||
}
|
||||
if called_name == "_Stop" {
|
||||
if stop_delay_secs > 0 {
|
||||
std::thread::sleep(std::time::Duration::from_secs(stop_delay_secs));
|
||||
}
|
||||
// Once we exceed the configured count, return empty
|
||||
// text so the agent treats it as no objection. This
|
||||
// lets a test exercise the "objected then cleared"
|
||||
// path without relying on the rejection budget.
|
||||
let payload = if stop_calls_seen < stop_count_limit {
|
||||
stop_text.clone()
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
stop_calls_seen = stop_calls_seen.saturating_add(1);
|
||||
write_response(
|
||||
id,
|
||||
json!({
|
||||
"content": [{ "type": "text", "text": payload }],
|
||||
"isError": false,
|
||||
}),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if called_name == "_PostCompact" {
|
||||
write_response(
|
||||
id,
|
||||
json!({
|
||||
"content": [{ "type": "text", "text": post_compact_text }],
|
||||
"isError": false,
|
||||
}),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if tool_delay_secs > 0 {
|
||||
std::thread::sleep(std::time::Duration::from_secs(tool_delay_secs));
|
||||
}
|
||||
let result_text = if result_size > 0 {
|
||||
"x".repeat(result_size)
|
||||
} else {
|
||||
"ok".to_owned()
|
||||
};
|
||||
let content = if env_flag("FAKE_MCP_IMAGE_RESULT") {
|
||||
json!([
|
||||
{ "type": "text", "text": result_text },
|
||||
{ "type": "image", "data": "aW1n", "mimeType": "image/png" },
|
||||
])
|
||||
} else {
|
||||
json!([{ "type": "text", "text": result_text }])
|
||||
};
|
||||
write_response(
|
||||
id,
|
||||
json!({
|
||||
"content": content,
|
||||
"isError": false,
|
||||
}),
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
// Unknown method: respond with an error so rmcp doesn't hang.
|
||||
let err = json!({
|
||||
"jsonrpc": "2.0", "id": id,
|
||||
"error": { "code": -32601, "message": format!("method not found: {method}") },
|
||||
});
|
||||
let mut s = serde_json::to_string(&err).unwrap();
|
||||
s.push('\n');
|
||||
let mut out = std::io::stdout().lock();
|
||||
let _ = out.write_all(s.as_bytes());
|
||||
let _ = out.flush();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,944 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
struct Harness {
|
||||
child: tokio::process::Child,
|
||||
stdin: tokio::process::ChildStdin,
|
||||
stdout: BufReader<tokio::process::ChildStdout>,
|
||||
next_id: i64,
|
||||
}
|
||||
|
||||
impl Harness {
|
||||
async fn spawn(extra: &[(&str, &str)]) -> Self {
|
||||
let bin = env!("CARGO_BIN_EXE_buzz-agent");
|
||||
let mut cmd = tokio::process::Command::new(bin);
|
||||
cmd.env("BUZZ_AGENT_PROVIDER", "openai")
|
||||
.env("OPENAI_COMPAT_API_KEY", "test")
|
||||
.env("OPENAI_COMPAT_MODEL", "fake-model")
|
||||
.env("BUZZ_AGENT_LLM_TIMEOUT_SECS", "5")
|
||||
.env("BUZZ_AGENT_TOOL_TIMEOUT_SECS", "5")
|
||||
.env("BUZZ_AGENT_MAX_ROUNDS", "4")
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::null())
|
||||
.kill_on_drop(true);
|
||||
for (k, v) in extra {
|
||||
cmd.env(k, v);
|
||||
}
|
||||
let mut child = cmd.spawn().expect("spawn buzz-agent");
|
||||
let stdin = child.stdin.take().unwrap();
|
||||
let stdout = BufReader::new(child.stdout.take().unwrap());
|
||||
Self {
|
||||
child,
|
||||
stdin,
|
||||
stdout,
|
||||
next_id: 1,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(&mut self, method: &str, params: Value) -> i64 {
|
||||
let id = self.next_id;
|
||||
self.next_id += 1;
|
||||
self.write_json(json!({
|
||||
"jsonrpc": "2.0", "id": id, "method": method, "params": params
|
||||
}))
|
||||
.await;
|
||||
id
|
||||
}
|
||||
|
||||
async fn notify(&mut self, method: &str, params: Value) {
|
||||
self.write_json(json!({
|
||||
"jsonrpc": "2.0", "method": method, "params": params
|
||||
}))
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn write_json(&mut self, msg: Value) {
|
||||
let mut s = serde_json::to_string(&msg).unwrap();
|
||||
s.push('\n');
|
||||
self.stdin.write_all(s.as_bytes()).await.unwrap();
|
||||
self.stdin.flush().await.unwrap();
|
||||
}
|
||||
|
||||
async fn write_raw(&mut self, raw: &[u8]) {
|
||||
let _ = self.stdin.write_all(raw).await;
|
||||
let _ = self.stdin.flush().await;
|
||||
}
|
||||
|
||||
async fn recv(&mut self) -> Value {
|
||||
let mut line = String::new();
|
||||
let n = tokio::time::timeout(Duration::from_secs(10), self.stdout.read_line(&mut line))
|
||||
.await
|
||||
.expect("recv timeout")
|
||||
.expect("read line");
|
||||
assert!(n > 0, "agent EOF");
|
||||
serde_json::from_str(&line).expect("non-JSON line")
|
||||
}
|
||||
|
||||
async fn recv_for_id(&mut self, id: i64) -> Value {
|
||||
loop {
|
||||
let v = self.recv().await;
|
||||
if v["id"] == json!(id) {
|
||||
return v;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn recv_until<F: FnMut(&Value) -> bool>(&mut self, mut pred: F) -> Value {
|
||||
loop {
|
||||
let v = self.recv().await;
|
||||
if pred(&v) {
|
||||
return v;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn shutdown(mut self) {
|
||||
drop(self.stdin);
|
||||
let _ = tokio::time::timeout(Duration::from_secs(2), self.child.wait()).await;
|
||||
let _ = self.child.start_kill();
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_fake_llm(responses: Vec<Value>) -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let queue = Arc::new(Mutex::new(VecDeque::from(responses)));
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (mut sock, _) = match listener.accept().await {
|
||||
Ok(p) => p,
|
||||
Err(_) => return,
|
||||
};
|
||||
let queue = queue.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut buf = Vec::new();
|
||||
let mut tmp = [0u8; 4096];
|
||||
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||||
match sock.read(&mut tmp).await {
|
||||
Ok(0) | Err(_) => return,
|
||||
Ok(n) => buf.extend_from_slice(&tmp[..n]),
|
||||
}
|
||||
if buf.len() > 1_000_000 {
|
||||
return;
|
||||
}
|
||||
}
|
||||
let body = queue
|
||||
.lock()
|
||||
.await
|
||||
.pop_front()
|
||||
.unwrap_or_else(|| json!({ "error": "no canned response" }));
|
||||
let body_s = serde_json::to_string(&body).unwrap();
|
||||
let resp = format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||||
body_s.len(),
|
||||
body_s,
|
||||
);
|
||||
let _ = sock.write_all(resp.as_bytes()).await;
|
||||
let _ = sock.shutdown().await;
|
||||
});
|
||||
}
|
||||
});
|
||||
url
|
||||
}
|
||||
|
||||
fn openai_text(content: &str) -> Value {
|
||||
json!({
|
||||
"id": "cc-1", "object": "chat.completion", "model": "fake-model",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": { "role": "assistant", "content": content },
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
})
|
||||
}
|
||||
|
||||
fn openai_tool_call(id: &str, name: &str, args: Value) -> Value {
|
||||
json!({
|
||||
"id": "cc-2", "object": "chat.completion", "model": "fake-model",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant", "content": null,
|
||||
"tool_calls": [{
|
||||
"id": id, "type": "function",
|
||||
"function": { "name": name, "arguments": args.to_string() },
|
||||
}],
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
}],
|
||||
})
|
||||
}
|
||||
|
||||
async fn handshake(h: &mut Harness) -> String {
|
||||
let init_id = h
|
||||
.send(
|
||||
"initialize",
|
||||
json!({ "protocolVersion": 2, "clientCapabilities": {} }),
|
||||
)
|
||||
.await;
|
||||
let init = h.recv_for_id(init_id).await;
|
||||
assert_eq!(init["result"]["protocolVersion"], 2);
|
||||
assert_eq!(init["result"]["agentInfo"]["name"], "buzz-agent");
|
||||
assert_eq!(
|
||||
init["result"]["agentCapabilities"]["promptCapabilities"]["image"],
|
||||
false
|
||||
);
|
||||
|
||||
let new_id = h
|
||||
.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
|
||||
.await;
|
||||
let new = h.recv_for_id(new_id).await;
|
||||
let sid = new["result"]["sessionId"].as_str().unwrap().to_owned();
|
||||
assert!(sid.starts_with("ses_"));
|
||||
sid
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_text_only_response() {
|
||||
let url = spawn_fake_llm(vec![openai_text("hello back")]).await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
let sid = handshake(&mut h).await;
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "hi" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let result = h.recv_for_id(p).await;
|
||||
assert_eq!(result["result"]["stopReason"], "end_turn");
|
||||
assert!(result.get("error").is_none());
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_full_tool_call_transcript() {
|
||||
let url = spawn_fake_llm(vec![
|
||||
openai_tool_call("call_xyz", "fake__do_thing", json!({ "foo": "bar" })),
|
||||
openai_text("done"),
|
||||
])
|
||||
.await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
let sid = handshake(&mut h).await;
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "use the tool" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let failed = h
|
||||
.recv_until(|v| {
|
||||
v.get("method") == Some(&json!("session/update"))
|
||||
&& v["params"]["update"]["sessionUpdate"] == "tool_call_update"
|
||||
&& v["params"]["update"]["status"] == "failed"
|
||||
})
|
||||
.await;
|
||||
assert_eq!(failed["params"]["sessionId"], sid);
|
||||
assert_eq!(failed["params"]["update"]["toolCallId"], "call_xyz");
|
||||
assert_eq!(
|
||||
failed["params"]["update"]["rawOutput"]["error"],
|
||||
"unknown tool: fake__do_thing"
|
||||
);
|
||||
|
||||
let final_resp = h.recv_for_id(p).await;
|
||||
assert_eq!(final_resp["result"]["stopReason"], "end_turn");
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_permission_denied_continues() {
|
||||
let url = spawn_fake_llm(vec![openai_text("ok with no tool")]).await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
let sid = handshake(&mut h).await;
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "hi" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let final_resp = h.recv_for_id(p).await;
|
||||
assert_eq!(final_resp["result"]["stopReason"], "end_turn");
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_initialize_version_check() {
|
||||
let url = spawn_fake_llm(vec![]).await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
let id = h
|
||||
.send(
|
||||
"initialize",
|
||||
json!({ "protocolVersion": 99, "clientCapabilities": {} }),
|
||||
)
|
||||
.await;
|
||||
let resp = h.recv_for_id(id).await;
|
||||
assert_eq!(resp["result"]["protocolVersion"], 2);
|
||||
|
||||
let id2 = h
|
||||
.send(
|
||||
"initialize",
|
||||
json!({ "protocolVersion": 1, "clientCapabilities": {} }),
|
||||
)
|
||||
.await;
|
||||
let ok = h.recv_for_id(id2).await;
|
||||
assert_eq!(ok["result"]["protocolVersion"], 1);
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_session_new_rejects_relative_cwd() {
|
||||
let url = spawn_fake_llm(vec![]).await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
let _ = h
|
||||
.send(
|
||||
"initialize",
|
||||
json!({ "protocolVersion": 1, "clientCapabilities": {} }),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv().await;
|
||||
|
||||
let id = h
|
||||
.send(
|
||||
"session/new",
|
||||
json!({ "cwd": "relative/path", "mcpServers": [] }),
|
||||
)
|
||||
.await;
|
||||
let resp = h.recv_for_id(id).await;
|
||||
assert_eq!(resp["error"]["code"], -32602);
|
||||
assert!(resp["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("cwd must be an absolute path"));
|
||||
|
||||
let id_empty = h
|
||||
.send("session/new", json!({ "cwd": "", "mcpServers": [] }))
|
||||
.await;
|
||||
let resp = h.recv_for_id(id_empty).await;
|
||||
assert_eq!(resp["error"]["code"], -32602);
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_malformed_json_rpc() {
|
||||
let url = spawn_fake_llm(vec![]).await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
h.write_raw(b"this is not json\n").await;
|
||||
let v = h.recv().await;
|
||||
assert_eq!(v["error"]["code"], -32700);
|
||||
assert_eq!(v["id"], Value::Null);
|
||||
|
||||
h.write_json(json!({ "jsonrpc": "1.0", "method": "initialize", "id": 1 }))
|
||||
.await;
|
||||
let v = h.recv().await;
|
||||
assert_eq!(v["error"]["code"], -32600);
|
||||
|
||||
h.write_json(json!({ "jsonrpc": "2.0" })).await;
|
||||
let v = h.recv().await;
|
||||
assert_eq!(v["error"]["code"], -32600);
|
||||
|
||||
let init_id = h
|
||||
.send(
|
||||
"initialize",
|
||||
json!({ "protocolVersion": 1, "clientCapabilities": {} }),
|
||||
)
|
||||
.await;
|
||||
let ok = h.recv_for_id(init_id).await;
|
||||
assert_eq!(ok["result"]["protocolVersion"], 1);
|
||||
|
||||
let bad_id = h.send("nonsense/method", json!({})).await;
|
||||
let v = h.recv_for_id(bad_id).await;
|
||||
assert_eq!(v["error"]["code"], -32601);
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_unsupported_content_block() {
|
||||
let url = spawn_fake_llm(vec![openai_text("ok")]).await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
let sid = handshake(&mut h).await;
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "image", "data": "..." }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let resp = h.recv_for_id(p).await;
|
||||
assert_eq!(resp["error"]["code"], -32602);
|
||||
assert!(resp["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("unsupported content block"));
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_concurrent_prompt_rejected() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
tokio::spawn(async move {
|
||||
let (mut sock, _) = listener.accept().await.unwrap();
|
||||
let mut buf = Vec::new();
|
||||
let mut tmp = [0u8; 4096];
|
||||
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||||
let n = sock.read(&mut tmp).await.unwrap_or(0);
|
||||
if n == 0 {
|
||||
return;
|
||||
}
|
||||
buf.extend_from_slice(&tmp[..n]);
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
let body = openai_text("done").to_string();
|
||||
let resp = format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
);
|
||||
let _ = sock.write_all(resp.as_bytes()).await;
|
||||
let _ = sock.shutdown().await;
|
||||
});
|
||||
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
let sid = handshake(&mut h).await;
|
||||
|
||||
let p1 = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({ "sessionId": sid, "prompt": [{"type":"text","text":"go"}] }),
|
||||
)
|
||||
.await;
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
let p2 = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({ "sessionId": sid, "prompt": [{"type":"text","text":"again"}] }),
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut p1_ok = false;
|
||||
let mut p2_err = false;
|
||||
for _ in 0..10 {
|
||||
let v = h.recv().await;
|
||||
if v["id"] == json!(p1) {
|
||||
assert_eq!(v["result"]["stopReason"], "end_turn");
|
||||
p1_ok = true;
|
||||
} else if v["id"] == json!(p2) {
|
||||
assert_eq!(v["error"]["code"], -32602);
|
||||
p2_err = true;
|
||||
}
|
||||
if p1_ok && p2_err {
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(p1_ok && p2_err, "expected p1=ok, p2=busy");
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_oversized_line_kills_agent() {
|
||||
let url = spawn_fake_llm(vec![]).await;
|
||||
let bin = env!("CARGO_BIN_EXE_buzz-agent");
|
||||
let mut cmd = tokio::process::Command::new(bin);
|
||||
cmd.env("BUZZ_AGENT_PROVIDER", "openai")
|
||||
.env("OPENAI_COMPAT_API_KEY", "test")
|
||||
.env("OPENAI_COMPAT_MODEL", "fake-model")
|
||||
.env("OPENAI_COMPAT_BASE_URL", &url)
|
||||
.env("BUZZ_AGENT_MAX_LINE_BYTES", "256")
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::null())
|
||||
.kill_on_drop(true);
|
||||
let mut child = cmd.spawn().unwrap();
|
||||
let mut stdin = child.stdin.take().unwrap();
|
||||
let big = "x".repeat(1024);
|
||||
let _ = stdin.write_all(big.as_bytes()).await;
|
||||
let _ = stdin.write_all(b"\n").await;
|
||||
drop(stdin);
|
||||
let _ = tokio::time::timeout(Duration::from_secs(5), child.wait())
|
||||
.await
|
||||
.expect("agent did not exit on oversized line");
|
||||
}
|
||||
|
||||
/// Build an Anthropic Messages API response with an optional `thinking` block
|
||||
/// followed by a `text` block. The `thinking` field is omitted when `None`.
|
||||
fn anthropic_thinking_response(thinking: Option<&str>, text: &str) -> Value {
|
||||
let mut content: Vec<Value> = Vec::new();
|
||||
if let Some(t) = thinking {
|
||||
content.push(json!({ "type": "thinking", "thinking": t }));
|
||||
}
|
||||
content.push(json!({ "type": "text", "text": text }));
|
||||
json!({
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-fake",
|
||||
"stop_reason": "end_turn",
|
||||
"content": content,
|
||||
"usage": { "input_tokens": 10, "output_tokens": 5 },
|
||||
})
|
||||
}
|
||||
|
||||
/// Build an OpenAI Responses API response with a `reasoning` output item
|
||||
/// (containing a single `summary_text` entry) followed by a message item.
|
||||
fn responses_reasoning_response(reasoning: &str, text: &str) -> Value {
|
||||
json!({
|
||||
"id": "resp_1",
|
||||
"status": "completed",
|
||||
"output": [
|
||||
{
|
||||
"type": "reasoning",
|
||||
"id": "rs_1",
|
||||
"summary": [{ "type": "summary_text", "text": reasoning }],
|
||||
},
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"content": [{ "type": "output_text", "text": text }],
|
||||
},
|
||||
],
|
||||
"usage": { "input_tokens": 10 },
|
||||
})
|
||||
}
|
||||
|
||||
/// Drain all `session/update` notifications until the `session/prompt` reply
|
||||
/// arrives for `prompt_id`, collecting notification payloads in order.
|
||||
async fn collect_updates_until_done(h: &mut Harness, prompt_id: i64) -> Vec<Value> {
|
||||
let mut updates = Vec::new();
|
||||
loop {
|
||||
let v = h.recv().await;
|
||||
if v.get("id") == Some(&json!(prompt_id)) {
|
||||
return updates;
|
||||
}
|
||||
if v.get("method") == Some(&json!("session/update")) {
|
||||
if let Some(u) = v["params"].get("update") {
|
||||
updates.push(u.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Asserts that `agent_thought_chunk` appears in `updates` BEFORE
|
||||
/// `agent_message_chunk`, and that both are present.
|
||||
fn assert_thought_before_message(updates: &[Value]) {
|
||||
let thought_pos = updates
|
||||
.iter()
|
||||
.position(|u| u["sessionUpdate"] == "agent_thought_chunk");
|
||||
let message_pos = updates
|
||||
.iter()
|
||||
.position(|u| u["sessionUpdate"] == "agent_message_chunk");
|
||||
assert!(
|
||||
thought_pos.is_some(),
|
||||
"expected agent_thought_chunk in updates: {updates:?}"
|
||||
);
|
||||
assert!(
|
||||
message_pos.is_some(),
|
||||
"expected agent_message_chunk in updates: {updates:?}"
|
||||
);
|
||||
assert!(
|
||||
thought_pos.unwrap() < message_pos.unwrap(),
|
||||
"agent_thought_chunk must precede agent_message_chunk, got updates: {updates:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_thought_chunk_emitted_before_message_chunk_anthropic() {
|
||||
// Anthropic extended-thinking: the response contains a `thinking` block
|
||||
// followed by a `text` block. We expect agent_thought_chunk to be emitted
|
||||
// before agent_message_chunk on the wire.
|
||||
let url = spawn_fake_llm(vec![anthropic_thinking_response(
|
||||
Some("Let me reason about this carefully."),
|
||||
"Here is my answer.",
|
||||
)])
|
||||
.await;
|
||||
let mut h = Harness::spawn(&[
|
||||
("BUZZ_AGENT_PROVIDER", "anthropic"),
|
||||
("ANTHROPIC_API_KEY", "test"),
|
||||
("ANTHROPIC_MODEL", "claude-fake"),
|
||||
("ANTHROPIC_BASE_URL", &url),
|
||||
("OPENAI_COMPAT_BASE_URL", ""),
|
||||
])
|
||||
.await;
|
||||
|
||||
let sid = handshake(&mut h).await;
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "think hard" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let updates = collect_updates_until_done(&mut h, p).await;
|
||||
assert_thought_before_message(&updates);
|
||||
|
||||
let thought = updates
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_thought_chunk")
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
thought["content"]["text"],
|
||||
"Let me reason about this carefully."
|
||||
);
|
||||
|
||||
let message = updates
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_message_chunk")
|
||||
.unwrap();
|
||||
assert_eq!(message["content"]["text"], "Here is my answer.");
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_thought_chunk_emitted_before_message_chunk_responses_api() {
|
||||
// OpenAI Responses API: reasoning item followed by message item.
|
||||
// Setting OPENAI_COMPAT_API=responses forces the Responses API parse path.
|
||||
let url = spawn_fake_llm(vec![responses_reasoning_response(
|
||||
"Thinking step by step.",
|
||||
"Final answer.",
|
||||
)])
|
||||
.await;
|
||||
let mut h = Harness::spawn(&[
|
||||
("BUZZ_AGENT_PROVIDER", "openai"),
|
||||
("OPENAI_COMPAT_API_KEY", "test"),
|
||||
("OPENAI_COMPAT_MODEL", "fake-model"),
|
||||
("OPENAI_COMPAT_API", "responses"),
|
||||
("OPENAI_COMPAT_BASE_URL", &url),
|
||||
])
|
||||
.await;
|
||||
|
||||
let sid = handshake(&mut h).await;
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "reason it out" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let updates = collect_updates_until_done(&mut h, p).await;
|
||||
assert_thought_before_message(&updates);
|
||||
|
||||
let thought = updates
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_thought_chunk")
|
||||
.unwrap();
|
||||
assert_eq!(thought["content"]["text"], "Thinking step by step.");
|
||||
|
||||
let message = updates
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_message_chunk")
|
||||
.unwrap();
|
||||
assert_eq!(message["content"]["text"], "Final answer.");
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_thought_chunk_emitted_before_message_chunk_chat_completions_reasoning_content() {
|
||||
// OpenAI chat/completions path with DeepSeek-style `reasoning_content` field
|
||||
// on the message object. OPENAI_COMPAT_API defaults to Auto, which routes
|
||||
// non-openai.com hosts to chat/completions — this is the live path for
|
||||
// self-hosted reasoning models (DeepSeek, vLLM, etc.).
|
||||
let response = json!({
|
||||
"id": "cc-r1", "object": "chat.completion", "model": "fake-model",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Here is the answer.",
|
||||
"reasoning_content": "Let me think through this step by step.",
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
});
|
||||
let url = spawn_fake_llm(vec![response]).await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
let sid = handshake(&mut h).await;
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "solve it" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let updates = collect_updates_until_done(&mut h, p).await;
|
||||
assert_thought_before_message(&updates);
|
||||
|
||||
let thought = updates
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_thought_chunk")
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
thought["content"]["text"],
|
||||
"Let me think through this step by step."
|
||||
);
|
||||
|
||||
let message = updates
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_message_chunk")
|
||||
.unwrap();
|
||||
assert_eq!(message["content"]["text"], "Here is the answer.");
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_no_reasoning_no_thought_chunk() {
|
||||
// Plain text response with no reasoning content — no agent_thought_chunk
|
||||
// should appear on the wire. This guards against empty thought emissions.
|
||||
let url = spawn_fake_llm(vec![openai_text("just text, no thinking")]).await;
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
|
||||
let sid = handshake(&mut h).await;
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "hi" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let updates = collect_updates_until_done(&mut h, p).await;
|
||||
|
||||
let has_thought = updates
|
||||
.iter()
|
||||
.any(|u| u["sessionUpdate"] == "agent_thought_chunk");
|
||||
assert!(
|
||||
!has_thought,
|
||||
"expected no agent_thought_chunk for a plain text response, got: {updates:?}"
|
||||
);
|
||||
|
||||
let has_message = updates
|
||||
.iter()
|
||||
.any(|u| u["sessionUpdate"] == "agent_message_chunk");
|
||||
assert!(has_message, "expected agent_message_chunk in updates");
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_cancel_notification_no_reply() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
tokio::spawn(async move {
|
||||
let (mut sock, _) = listener.accept().await.unwrap();
|
||||
let mut buf = Vec::new();
|
||||
let mut tmp = [0u8; 4096];
|
||||
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||||
let n = sock.read(&mut tmp).await.unwrap_or(0);
|
||||
if n == 0 {
|
||||
return;
|
||||
}
|
||||
buf.extend_from_slice(&tmp[..n]);
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(800)).await;
|
||||
let body = openai_text("done").to_string();
|
||||
let resp = format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
);
|
||||
let _ = sock.write_all(resp.as_bytes()).await;
|
||||
let _ = sock.shutdown().await;
|
||||
});
|
||||
|
||||
let mut h = Harness::spawn(&[("OPENAI_COMPAT_BASE_URL", &url)]).await;
|
||||
let sid = handshake(&mut h).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({ "sessionId": sid, "prompt": [{"type":"text","text":"go"}] }),
|
||||
)
|
||||
.await;
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
h.notify("session/cancel", json!({ "sessionId": sid }))
|
||||
.await;
|
||||
|
||||
let final_resp = h.recv_for_id(p).await;
|
||||
let stop = final_resp["result"]["stopReason"].as_str().unwrap_or("");
|
||||
assert!(
|
||||
stop == "cancelled" || stop == "end_turn",
|
||||
"unexpected stopReason {stop}"
|
||||
);
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// ACP v2 ContentChunk compliance: both `agent_thought_chunk` and
|
||||
/// `agent_message_chunk` must carry `messageId` and `content` when the
|
||||
/// client negotiates protocol version 2.
|
||||
///
|
||||
/// ACP v2 requires `ContentChunk.messageId` (required in v2 schema at
|
||||
/// agentclientprotocol/agent-client-protocol schema/v2/schema.json @d13d1baa).
|
||||
/// ACP v1 allows the field, so adding it is backwards-safe.
|
||||
///
|
||||
/// Additional invariants verified here:
|
||||
/// - The thought and assistant message IDs are **distinct** (two logical messages).
|
||||
/// - IDs do **not** recur across two consecutive `session/prompt` calls in the same
|
||||
/// ACP session (`run_id` is fresh per prompt, so no cross-turn collision).
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_acp_v2_chunks_carry_message_id() {
|
||||
// OpenAI Responses API: reasoning item + text item. Both emitted chunks
|
||||
// must have messageId + content on a v2 connection.
|
||||
// Two responses so we can send two session/prompt calls and verify no ID reuse.
|
||||
let url = spawn_fake_llm(vec![
|
||||
responses_reasoning_response("Thinking about it.", "Here is my response."),
|
||||
responses_reasoning_response("Thinking again.", "Second response."),
|
||||
])
|
||||
.await;
|
||||
let mut h = Harness::spawn(&[
|
||||
("BUZZ_AGENT_PROVIDER", "openai"),
|
||||
("OPENAI_COMPAT_API_KEY", "test"),
|
||||
("OPENAI_COMPAT_MODEL", "fake-model"),
|
||||
("OPENAI_COMPAT_API", "responses"),
|
||||
("OPENAI_COMPAT_BASE_URL", &url),
|
||||
])
|
||||
.await;
|
||||
|
||||
let sid = handshake(&mut h).await; // negotiates protocolVersion: 2
|
||||
|
||||
// ── First prompt ──────────────────────────────────────────────────────────
|
||||
let p1 = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "think and respond" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let updates1 = collect_updates_until_done(&mut h, p1).await;
|
||||
|
||||
let thought1 = updates1
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_thought_chunk")
|
||||
.expect("agent_thought_chunk must be emitted on prompt 1");
|
||||
let message1 = updates1
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_message_chunk")
|
||||
.expect("agent_message_chunk must be emitted on prompt 1");
|
||||
|
||||
// ACP v2 ContentChunk compliance: messageId must be present and non-empty.
|
||||
let thought_id1 = thought1["messageId"]
|
||||
.as_str()
|
||||
.expect("agent_thought_chunk must carry messageId (ACP v2 required field)");
|
||||
assert!(
|
||||
!thought_id1.is_empty(),
|
||||
"agent_thought_chunk messageId must not be empty"
|
||||
);
|
||||
|
||||
let message_id1 = message1["messageId"]
|
||||
.as_str()
|
||||
.expect("agent_message_chunk must carry messageId (ACP v2 required field)");
|
||||
assert!(
|
||||
!message_id1.is_empty(),
|
||||
"agent_message_chunk messageId must not be empty"
|
||||
);
|
||||
|
||||
// Thought and assistant message are two distinct logical messages — their IDs must differ.
|
||||
assert_ne!(
|
||||
thought_id1, message_id1,
|
||||
"agent_thought_chunk and agent_message_chunk are distinct logical messages; their messageIds must differ"
|
||||
);
|
||||
|
||||
// content must be present and correct.
|
||||
assert_eq!(
|
||||
thought1["content"]["text"], "Thinking about it.",
|
||||
"thought content mismatch"
|
||||
);
|
||||
assert_eq!(
|
||||
message1["content"]["text"], "Here is my response.",
|
||||
"message content mismatch"
|
||||
);
|
||||
|
||||
// ── Second prompt (same ACP session) ─────────────────────────────────────
|
||||
let p2 = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({
|
||||
"sessionId": sid,
|
||||
"prompt": [{ "type": "text", "text": "think again" }],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let updates2 = collect_updates_until_done(&mut h, p2).await;
|
||||
|
||||
let thought2 = updates2
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_thought_chunk")
|
||||
.expect("agent_thought_chunk must be emitted on prompt 2");
|
||||
let message2 = updates2
|
||||
.iter()
|
||||
.find(|u| u["sessionUpdate"] == "agent_message_chunk")
|
||||
.expect("agent_message_chunk must be emitted on prompt 2");
|
||||
|
||||
let thought_id2 = thought2["messageId"]
|
||||
.as_str()
|
||||
.expect("agent_thought_chunk must carry messageId on prompt 2");
|
||||
let message_id2 = message2["messageId"]
|
||||
.as_str()
|
||||
.expect("agent_message_chunk must carry messageId on prompt 2");
|
||||
|
||||
// IDs from prompt 2 must be distinct from each other.
|
||||
assert_ne!(
|
||||
thought_id2, message_id2,
|
||||
"prompt 2: thought and message IDs must differ"
|
||||
);
|
||||
|
||||
// IDs must NOT recur across prompts — ACP requires session-unique messageIds.
|
||||
assert_ne!(
|
||||
thought_id1, thought_id2,
|
||||
"thought messageId must not recur across session/prompt calls (run_id must differ)"
|
||||
);
|
||||
assert_ne!(
|
||||
message_id1, message_id2,
|
||||
"message messageId must not recur across session/prompt calls (run_id must differ)"
|
||||
);
|
||||
|
||||
h.shutdown().await;
|
||||
}
|
||||
@@ -0,0 +1,574 @@
|
||||
//! Integration tests for AGENTS.md / SKILL.md hint loading.
|
||||
//!
|
||||
//! Uses the same subprocess + capturing-LLM pattern as `regressions.rs`.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
struct CapturingLlm {
|
||||
url: String,
|
||||
captured: Arc<Mutex<Vec<Value>>>,
|
||||
}
|
||||
|
||||
async fn spawn_capturing_llm(responses: Vec<Value>) -> CapturingLlm {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let queue = Arc::new(Mutex::new(VecDeque::from(responses)));
|
||||
let captured: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
|
||||
let cap2 = captured.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (mut sock, _) = match listener.accept().await {
|
||||
Ok(p) => p,
|
||||
Err(_) => return,
|
||||
};
|
||||
let queue = queue.clone();
|
||||
let captured = cap2.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut buf = Vec::new();
|
||||
let mut tmp = [0u8; 8192];
|
||||
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||||
match sock.read(&mut tmp).await {
|
||||
Ok(0) | Err(_) => return,
|
||||
Ok(n) => buf.extend_from_slice(&tmp[..n]),
|
||||
}
|
||||
if buf.len() > 4_000_000 {
|
||||
return;
|
||||
}
|
||||
}
|
||||
let header_end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4;
|
||||
let headers = &buf[..header_end];
|
||||
let mut body_len = 0usize;
|
||||
for line in headers.split(|b| *b == b'\n') {
|
||||
let line = std::str::from_utf8(line).unwrap_or("");
|
||||
if let Some(rest) = line.to_ascii_lowercase().strip_prefix("content-length:") {
|
||||
body_len = rest.trim().trim_end_matches('\r').parse().unwrap_or(0);
|
||||
}
|
||||
}
|
||||
while buf.len() < header_end + body_len {
|
||||
match sock.read(&mut tmp).await {
|
||||
Ok(0) | Err(_) => return,
|
||||
Ok(n) => buf.extend_from_slice(&tmp[..n]),
|
||||
}
|
||||
}
|
||||
if let Ok(req) = serde_json::from_slice::<Value>(&buf[header_end..]) {
|
||||
captured.lock().await.push(req);
|
||||
}
|
||||
let body = queue
|
||||
.lock()
|
||||
.await
|
||||
.pop_front()
|
||||
.unwrap_or_else(|| json!({ "error": "no canned response" }));
|
||||
let body_s = serde_json::to_string(&body).unwrap();
|
||||
let resp = format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
|
||||
Content-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||||
body_s.len(),
|
||||
body_s,
|
||||
);
|
||||
let _ = sock.write_all(resp.as_bytes()).await;
|
||||
let _ = sock.shutdown().await;
|
||||
});
|
||||
}
|
||||
});
|
||||
CapturingLlm { url, captured }
|
||||
}
|
||||
|
||||
struct Harness {
|
||||
child: tokio::process::Child,
|
||||
stdin: tokio::process::ChildStdin,
|
||||
stdout: BufReader<tokio::process::ChildStdout>,
|
||||
next_id: i64,
|
||||
}
|
||||
|
||||
impl Harness {
|
||||
async fn spawn_with_env(base_url: &str, extra: &[(&str, &str)]) -> Self {
|
||||
let bin = env!("CARGO_BIN_EXE_buzz-agent");
|
||||
let mut cmd = tokio::process::Command::new(bin);
|
||||
cmd.env("BUZZ_AGENT_PROVIDER", "openai")
|
||||
.env("OPENAI_COMPAT_API_KEY", "test")
|
||||
.env("OPENAI_COMPAT_MODEL", "fake-model")
|
||||
.env("OPENAI_COMPAT_BASE_URL", base_url)
|
||||
.env("BUZZ_AGENT_LLM_TIMEOUT_SECS", "5")
|
||||
.env("BUZZ_AGENT_TOOL_TIMEOUT_SECS", "5")
|
||||
.env("BUZZ_AGENT_MAX_ROUNDS", "8")
|
||||
.env("BUZZ_AGENT_MCP_INIT_TIMEOUT_SECS", "2");
|
||||
for (k, v) in extra {
|
||||
cmd.env(k, v);
|
||||
}
|
||||
cmd.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::inherit())
|
||||
.kill_on_drop(true);
|
||||
let mut child = cmd.spawn().expect("spawn buzz-agent");
|
||||
let stdin = child.stdin.take().unwrap();
|
||||
let stdout = BufReader::new(child.stdout.take().unwrap());
|
||||
Self {
|
||||
child,
|
||||
stdin,
|
||||
stdout,
|
||||
next_id: 1,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(&mut self, method: &str, params: Value) -> i64 {
|
||||
let id = self.next_id;
|
||||
self.next_id += 1;
|
||||
self.write(json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }))
|
||||
.await;
|
||||
id
|
||||
}
|
||||
|
||||
async fn write(&mut self, msg: Value) {
|
||||
let mut s = serde_json::to_string(&msg).unwrap();
|
||||
s.push('\n');
|
||||
self.stdin.write_all(s.as_bytes()).await.unwrap();
|
||||
self.stdin.flush().await.unwrap();
|
||||
}
|
||||
|
||||
async fn recv(&mut self) -> Value {
|
||||
let mut line = String::new();
|
||||
let n = tokio::time::timeout(Duration::from_secs(15), self.stdout.read_line(&mut line))
|
||||
.await
|
||||
.expect("recv timeout")
|
||||
.expect("read line");
|
||||
assert!(n > 0, "agent EOF");
|
||||
serde_json::from_str(&line).expect("non-JSON line")
|
||||
}
|
||||
|
||||
async fn recv_until<F: FnMut(&Value) -> bool>(&mut self, mut pred: F) -> Value {
|
||||
loop {
|
||||
let v = self.recv().await;
|
||||
if pred(&v) {
|
||||
return v;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn shutdown(mut self) {
|
||||
drop(self.stdin);
|
||||
let _ = tokio::time::timeout(Duration::from_secs(2), self.child.wait()).await;
|
||||
let _ = self.child.start_kill();
|
||||
}
|
||||
}
|
||||
|
||||
fn openai_text(content: &str) -> Value {
|
||||
json!({
|
||||
"id": "cc-1", "object": "chat.completion", "model": "fake-model",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": { "role": "assistant", "content": content },
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
})
|
||||
}
|
||||
|
||||
async fn init_session(h: &mut Harness, cwd: &str) -> String {
|
||||
h.send(
|
||||
"initialize",
|
||||
json!({"protocolVersion": 1, "clientCapabilities": {}}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv().await;
|
||||
h.send("session/new", json!({"cwd": cwd, "mcpServers": []}))
|
||||
.await;
|
||||
let r = h
|
||||
.recv_until(|v| v.get("result").is_some() || v.get("error").is_some())
|
||||
.await;
|
||||
r["result"]["sessionId"]
|
||||
.as_str()
|
||||
.expect("sessionId")
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
/// AGENTS.md in cwd is loaded into the system prompt.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn hints_loaded_from_cwd_agents_md() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
let marker = "BUZZ_HINTS_MARKER_42";
|
||||
std::fs::write(cwd.join("AGENTS.md"), marker).unwrap();
|
||||
|
||||
let llm = spawn_capturing_llm(vec![openai_text("done")]).await;
|
||||
let mut h = Harness::spawn_with_env(&llm.url, &[]).await;
|
||||
let sid = init_session(&mut h, cwd.to_str().unwrap()).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
|
||||
let captured = llm.captured.lock().await;
|
||||
assert!(!captured.is_empty(), "no LLM request captured");
|
||||
let system = captured[0]["messages"][0]["content"].as_str().unwrap_or("");
|
||||
assert!(
|
||||
system.contains(marker),
|
||||
"system prompt does not contain AGENTS.md marker: {system}"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// BUZZ_AGENT_NO_HINTS=1 suppresses hint loading.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn hints_suppressed_with_env_var() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
let marker = "SUPPRESS_CHECK_MARKER_99";
|
||||
std::fs::write(cwd.join("AGENTS.md"), marker).unwrap();
|
||||
|
||||
let llm = spawn_capturing_llm(vec![openai_text("done")]).await;
|
||||
let mut h = Harness::spawn_with_env(&llm.url, &[("BUZZ_AGENT_NO_HINTS", "1")]).await;
|
||||
let sid = init_session(&mut h, cwd.to_str().unwrap()).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
|
||||
let captured = llm.captured.lock().await;
|
||||
assert!(!captured.is_empty(), "no LLM request captured");
|
||||
let system = captured[0]["messages"][0]["content"].as_str().unwrap_or("");
|
||||
assert!(
|
||||
!system.contains(marker),
|
||||
"system prompt should NOT contain marker when hints disabled: {system}"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// SKILL.md files in .agents/skills/ are loaded into the system prompt as metadata only.
|
||||
/// The body is NOT inlined; the agent uses `load_skill` to fetch it on demand.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn skills_loaded_from_agents_skills_dir() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
let skill_dir = cwd.join(".agents/skills/test-skill");
|
||||
std::fs::create_dir_all(&skill_dir).unwrap();
|
||||
std::fs::write(
|
||||
skill_dir.join("SKILL.md"),
|
||||
"---\nname: test-skill\ndescription: A test skill\n---\nSKILL_BODY_MARKER_77\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let llm = spawn_capturing_llm(vec![openai_text("done")]).await;
|
||||
let mut h = Harness::spawn_with_env(&llm.url, &[]).await;
|
||||
let sid = init_session(&mut h, cwd.to_str().unwrap()).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
|
||||
let captured = llm.captured.lock().await;
|
||||
assert!(!captured.is_empty(), "no LLM request captured");
|
||||
let system = captured[0]["messages"][0]["content"].as_str().unwrap_or("");
|
||||
// Skill name must appear in the metadata listing.
|
||||
assert!(
|
||||
system.contains("test-skill"),
|
||||
"system prompt missing skill name: {system}"
|
||||
);
|
||||
// Body must NOT be inlined — lazy loading only.
|
||||
assert!(
|
||||
!system.contains("SKILL_BODY_MARKER_77"),
|
||||
"skill body must not be inlined in system prompt: {system}"
|
||||
);
|
||||
// The load_skill instruction must be present.
|
||||
assert!(
|
||||
system.contains("load_skill"),
|
||||
"system prompt missing load_skill instruction: {system}"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// AGENTS.md files at git root and subdirectory are both loaded, root first.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn git_root_hints_included() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let root = tmp.path();
|
||||
std::fs::create_dir(root.join(".git")).unwrap();
|
||||
std::fs::write(root.join("AGENTS.md"), "ROOT_HINT_MARKER_11").unwrap();
|
||||
let sub = root.join("sub");
|
||||
std::fs::create_dir(&sub).unwrap();
|
||||
std::fs::write(sub.join("AGENTS.md"), "SUB_HINT_MARKER_22").unwrap();
|
||||
|
||||
let llm = spawn_capturing_llm(vec![openai_text("done")]).await;
|
||||
let mut h = Harness::spawn_with_env(&llm.url, &[]).await;
|
||||
let sid = init_session(&mut h, sub.to_str().unwrap()).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
|
||||
let captured = llm.captured.lock().await;
|
||||
assert!(!captured.is_empty(), "no LLM request captured");
|
||||
let system = captured[0]["messages"][0]["content"].as_str().unwrap_or("");
|
||||
assert!(
|
||||
system.contains("ROOT_HINT_MARKER_11"),
|
||||
"system prompt missing root hint: {system}"
|
||||
);
|
||||
assert!(
|
||||
system.contains("SUB_HINT_MARKER_22"),
|
||||
"system prompt missing sub hint: {system}"
|
||||
);
|
||||
let root_pos = system.find("ROOT_HINT_MARKER_11").unwrap();
|
||||
let sub_pos = system.find("SUB_HINT_MARKER_22").unwrap();
|
||||
assert!(
|
||||
root_pos < sub_pos,
|
||||
"root hint should appear before sub hint in system prompt"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// ~/AGENTS.md (global) is loaded before CWD AGENTS.md when HOME is set.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn global_agents_md_loaded() {
|
||||
let home_tmp = tempfile::TempDir::new().unwrap();
|
||||
let cwd_tmp = tempfile::TempDir::new().unwrap();
|
||||
std::fs::write(home_tmp.path().join("AGENTS.md"), "GLOBAL_HINT_MARKER_55").unwrap();
|
||||
std::fs::write(cwd_tmp.path().join("AGENTS.md"), "LOCAL_HINT_MARKER_66").unwrap();
|
||||
|
||||
let llm = spawn_capturing_llm(vec![openai_text("done")]).await;
|
||||
let mut h =
|
||||
Harness::spawn_with_env(&llm.url, &[("HOME", home_tmp.path().to_str().unwrap())]).await;
|
||||
let sid = init_session(&mut h, cwd_tmp.path().to_str().unwrap()).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
|
||||
let captured = llm.captured.lock().await;
|
||||
assert!(!captured.is_empty(), "no LLM request captured");
|
||||
let system = captured[0]["messages"][0]["content"].as_str().unwrap_or("");
|
||||
assert!(
|
||||
system.contains("GLOBAL_HINT_MARKER_55"),
|
||||
"system prompt missing global hint: {system}"
|
||||
);
|
||||
assert!(
|
||||
system.contains("LOCAL_HINT_MARKER_66"),
|
||||
"system prompt missing local hint: {system}"
|
||||
);
|
||||
let global_pos = system.find("GLOBAL_HINT_MARKER_55").unwrap();
|
||||
let local_pos = system.find("LOCAL_HINT_MARKER_66").unwrap();
|
||||
assert!(
|
||||
global_pos < local_pos,
|
||||
"global hint should appear before local hint in system prompt"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// Global skills from ~/.agents/skills/ are loaded; project-level wins on name conflict.
|
||||
/// Bodies are NOT inlined — only metadata (name + description) appears in the system prompt.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn global_skills_loaded_and_project_wins() {
|
||||
let home_tmp = tempfile::TempDir::new().unwrap();
|
||||
let cwd_tmp = tempfile::TempDir::new().unwrap();
|
||||
|
||||
let global_only_dir = home_tmp.path().join(".agents/skills/global-only");
|
||||
std::fs::create_dir_all(&global_only_dir).unwrap();
|
||||
std::fs::write(
|
||||
global_only_dir.join("SKILL.md"),
|
||||
"---\nname: global-only\ndescription: A global skill\n---\nGLOBAL_SKILL_BODY_88\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let global_shared_dir = home_tmp.path().join(".agents/skills/shared-name");
|
||||
std::fs::create_dir_all(&global_shared_dir).unwrap();
|
||||
std::fs::write(
|
||||
global_shared_dir.join("SKILL.md"),
|
||||
"---\nname: shared-name\ndescription: Global version\n---\nGLOBAL_SHARED_BODY_LOSE\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let project_shared_dir = cwd_tmp.path().join(".agents/skills/shared-name");
|
||||
std::fs::create_dir_all(&project_shared_dir).unwrap();
|
||||
std::fs::write(
|
||||
project_shared_dir.join("SKILL.md"),
|
||||
"---\nname: shared-name\ndescription: Project version\n---\nPROJECT_SHARED_BODY_WIN\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let llm = spawn_capturing_llm(vec![openai_text("done")]).await;
|
||||
let mut h =
|
||||
Harness::spawn_with_env(&llm.url, &[("HOME", home_tmp.path().to_str().unwrap())]).await;
|
||||
let sid = init_session(&mut h, cwd_tmp.path().to_str().unwrap()).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
|
||||
let captured = llm.captured.lock().await;
|
||||
assert!(!captured.is_empty(), "no LLM request captured");
|
||||
let system = captured[0]["messages"][0]["content"].as_str().unwrap_or("");
|
||||
// Both skill names must appear in the metadata listing.
|
||||
assert!(
|
||||
system.contains("global-only"),
|
||||
"system prompt missing global-only skill name: {system}"
|
||||
);
|
||||
assert!(
|
||||
system.contains("shared-name"),
|
||||
"system prompt missing shared-name skill: {system}"
|
||||
);
|
||||
// Project description wins over global for the shared name.
|
||||
assert!(
|
||||
system.contains("Project version"),
|
||||
"system prompt should show project description for shared-name: {system}"
|
||||
);
|
||||
assert!(
|
||||
!system.contains("Global version"),
|
||||
"system prompt should NOT show global description for shared-name: {system}"
|
||||
);
|
||||
// Bodies must NOT be inlined.
|
||||
assert!(
|
||||
!system.contains("GLOBAL_SKILL_BODY_88"),
|
||||
"skill body must not be inlined: {system}"
|
||||
);
|
||||
assert!(
|
||||
!system.contains("PROJECT_SHARED_BODY_WIN"),
|
||||
"skill body must not be inlined: {system}"
|
||||
);
|
||||
assert!(
|
||||
!system.contains("GLOBAL_SHARED_BODY_LOSE"),
|
||||
"shadowed skill body must not be inlined: {system}"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// Skill directories that are symlinks (e.g. managed by ai-rules) are discovered
|
||||
/// correctly — `DirEntry::file_type()` returns `FileType::Symlink` for symlinks,
|
||||
/// so the old `is_dir()` check silently dropped them. We now use
|
||||
/// `std::fs::metadata()` which follows the symlink.
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn symlinked_skill_dir_is_discovered() {
|
||||
let real_skill_root = tempfile::TempDir::new().unwrap();
|
||||
let real_skill_dir = real_skill_root.path().join("symlinked-skill");
|
||||
std::fs::create_dir_all(&real_skill_dir).unwrap();
|
||||
std::fs::write(
|
||||
real_skill_dir.join("SKILL.md"),
|
||||
"---\nname: symlinked-skill\ndescription: A symlinked skill\n---\nSYMLINK_SKILL_BODY_42\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
let skills_dir = cwd.join(".agents/skills");
|
||||
std::fs::create_dir_all(&skills_dir).unwrap();
|
||||
|
||||
// Create a symlink: .agents/skills/symlinked-skill -> real_skill_dir
|
||||
std::os::unix::fs::symlink(&real_skill_dir, skills_dir.join("symlinked-skill")).unwrap();
|
||||
|
||||
let llm = spawn_capturing_llm(vec![openai_text("done")]).await;
|
||||
let mut h = Harness::spawn_with_env(&llm.url, &[]).await;
|
||||
let sid = init_session(&mut h, cwd.to_str().unwrap()).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
|
||||
let captured = llm.captured.lock().await;
|
||||
assert!(!captured.is_empty(), "no LLM request captured");
|
||||
let system = captured[0]["messages"][0]["content"].as_str().unwrap_or("");
|
||||
// The symlinked skill name must appear in the metadata listing.
|
||||
assert!(
|
||||
system.contains("symlinked-skill"),
|
||||
"system prompt missing symlinked skill name: {system}"
|
||||
);
|
||||
// Body must NOT be inlined.
|
||||
assert!(
|
||||
!system.contains("SYMLINK_SKILL_BODY_42"),
|
||||
"symlinked skill body must not be inlined in system prompt: {system}"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// `load_skill` tool is advertised when skills exist, and returns the skill body.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn load_skill_tool_returns_body() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let cwd = tmp.path();
|
||||
let skill_dir = cwd.join(".agents/skills/my-skill");
|
||||
std::fs::create_dir_all(&skill_dir).unwrap();
|
||||
std::fs::write(
|
||||
skill_dir.join("SKILL.md"),
|
||||
"---\nname: my-skill\ndescription: A skill\n---\nSKILL_BODY_CONTENT_99\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Round 1: LLM calls load_skill("my-skill").
|
||||
// Round 2: LLM returns end_turn after seeing the body.
|
||||
let load_skill_call = json!({
|
||||
"id": "cc-ls", "object": "chat.completion", "model": "fake-model",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant", "content": null,
|
||||
"tool_calls": [{
|
||||
"id": "tc-1", "type": "function",
|
||||
"function": {
|
||||
"name": "load_skill",
|
||||
"arguments": "{\"name\":\"my-skill\"}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": "tool_calls"
|
||||
}]
|
||||
});
|
||||
let end_turn = openai_text("done");
|
||||
|
||||
let llm = spawn_capturing_llm(vec![load_skill_call, end_turn]).await;
|
||||
let mut h = Harness::spawn_with_env(&llm.url, &[]).await;
|
||||
let sid = init_session(&mut h, cwd.to_str().unwrap()).await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"use my-skill"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
|
||||
// The second LLM request (round 2) should contain the skill body in tool results.
|
||||
let reqs = llm.captured.lock().await;
|
||||
assert!(
|
||||
reqs.len() >= 2,
|
||||
"expected at least 2 LLM requests, got {}",
|
||||
reqs.len()
|
||||
);
|
||||
let round2_str = serde_json::to_string(&reqs[1]).unwrap();
|
||||
assert!(
|
||||
round2_str.contains("SKILL_BODY_CONTENT_99"),
|
||||
"load_skill result must contain skill body in round 2 request.\nGot: {round2_str}"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
//! Integration test for OpenAI auto-upgrade chat→responses.
|
||||
//!
|
||||
//! Starts a tiny HTTP server that:
|
||||
//! 1. accepts a POST to /chat/completions, replies 400 with a body that
|
||||
//! mentions `/v1/responses` (mirrors the Databricks GPT-5.5 signal);
|
||||
//! 2. accepts a POST to /responses, replies 200 with a Responses-shaped
|
||||
//! JSON envelope.
|
||||
//!
|
||||
//! Spawns `buzz-agent` with `provider=openai` + `OPENAI_COMPAT_API=auto`
|
||||
//! pointed at the fake server, drives one prompt through the ACP wire
|
||||
//! protocol, and verifies the prompt completes with `stopReason=end_turn`
|
||||
//! — which can only happen if the second (Responses) request succeeded.
|
||||
|
||||
use std::io::{Read, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::process::Stdio;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::json;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::process::Command;
|
||||
use tokio::time::timeout;
|
||||
|
||||
/// Spawns a single-shot fake provider. Returns the base URL (e.g.
|
||||
/// `http://127.0.0.1:54321`). The server stays up for the lifetime of
|
||||
/// the process — we don't need to clean it up explicitly.
|
||||
fn spawn_fake_provider() -> (String, Arc<AtomicUsize>, Arc<AtomicUsize>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
listener.set_nonblocking(false).unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let chat_hits = Arc::new(AtomicUsize::new(0));
|
||||
let responses_hits = Arc::new(AtomicUsize::new(0));
|
||||
let chat = chat_hits.clone();
|
||||
let resp = responses_hits.clone();
|
||||
|
||||
std::thread::spawn(move || {
|
||||
loop {
|
||||
let (mut sock, _) = match listener.accept() {
|
||||
Ok(p) => p,
|
||||
Err(_) => return,
|
||||
};
|
||||
let chat = chat.clone();
|
||||
let resp = resp.clone();
|
||||
std::thread::spawn(move || {
|
||||
sock.set_read_timeout(Some(Duration::from_secs(5))).ok();
|
||||
// Read request head + body. Naive: read until we have the
|
||||
// request line + headers, then read Content-Length bytes.
|
||||
let mut buf = Vec::with_capacity(4096);
|
||||
let mut tmp = [0u8; 4096];
|
||||
loop {
|
||||
if buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
match sock.read(&mut tmp) {
|
||||
Ok(0) | Err(_) => return,
|
||||
Ok(n) => buf.extend_from_slice(&tmp[..n]),
|
||||
}
|
||||
if buf.len() > 256 * 1024 {
|
||||
return;
|
||||
}
|
||||
}
|
||||
let head_end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4;
|
||||
let head = String::from_utf8_lossy(&buf[..head_end]).to_string();
|
||||
// Drain the body to satisfy keep-alive; we don't actually
|
||||
// need it.
|
||||
let cl = head
|
||||
.lines()
|
||||
.find_map(|l| {
|
||||
l.strip_prefix("content-length:")
|
||||
.or_else(|| l.strip_prefix("Content-Length:"))
|
||||
})
|
||||
.and_then(|s| s.trim().parse::<usize>().ok())
|
||||
.unwrap_or(0);
|
||||
while buf.len() < head_end + cl {
|
||||
match sock.read(&mut tmp) {
|
||||
Ok(0) | Err(_) => break,
|
||||
Ok(n) => buf.extend_from_slice(&tmp[..n]),
|
||||
}
|
||||
}
|
||||
|
||||
let (status, body) = if head.contains("POST /chat/completions") {
|
||||
chat.fetch_add(1, Ordering::SeqCst);
|
||||
let body = json!({
|
||||
"error": {
|
||||
"code": "BAD_REQUEST",
|
||||
"message": "Function tools with reasoning_effort are not supported for gpt-5.5 in /v1/chat/completions. Please use /v1/responses instead."
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
(400u16, body)
|
||||
} else if head.contains("POST /responses") {
|
||||
resp.fetch_add(1, Ordering::SeqCst);
|
||||
let body = json!({
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "ok from responses"}]
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
(200u16, body)
|
||||
} else {
|
||||
(404u16, "{}".to_string())
|
||||
};
|
||||
let reason = match status {
|
||||
200 => "OK",
|
||||
400 => "Bad Request",
|
||||
_ => "Not Found",
|
||||
};
|
||||
let resp_text = format!(
|
||||
"HTTP/1.1 {status} {reason}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||||
body.len(), body
|
||||
);
|
||||
let _ = sock.write_all(resp_text.as_bytes());
|
||||
let _ = sock.shutdown(std::net::Shutdown::Write);
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
(url, chat_hits, responses_hits)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_auto_upgrades_chat_to_responses_on_databricks_signal() {
|
||||
let (base_url, chat_hits, resp_hits) = spawn_fake_provider();
|
||||
|
||||
let bin = env!("CARGO_BIN_EXE_buzz-agent");
|
||||
let mut cmd = Command::new(bin);
|
||||
cmd.env("BUZZ_AGENT_PROVIDER", "openai")
|
||||
.env("OPENAI_COMPAT_API_KEY", "test")
|
||||
.env("OPENAI_COMPAT_MODEL", "gpt-5.5")
|
||||
.env("OPENAI_COMPAT_BASE_URL", &base_url)
|
||||
// No OPENAI_COMPAT_API — must default to "auto" so the upgrade
|
||||
// path is enabled.
|
||||
.env_remove("OPENAI_COMPAT_API")
|
||||
.env("BUZZ_AGENT_LLM_TIMEOUT_SECS", "5")
|
||||
.env("BUZZ_AGENT_MAX_ROUNDS", "4")
|
||||
.env("BUZZ_AGENT_MCP_INIT_TIMEOUT_SECS", "2")
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::inherit())
|
||||
.kill_on_drop(true);
|
||||
|
||||
let mut child = cmd.spawn().expect("spawn buzz-agent");
|
||||
let mut stdin = child.stdin.take().unwrap();
|
||||
let mut stdout = BufReader::new(child.stdout.take().unwrap());
|
||||
|
||||
async fn send(stdin: &mut tokio::process::ChildStdin, v: serde_json::Value) {
|
||||
let line = format!("{v}\n");
|
||||
stdin.write_all(line.as_bytes()).await.unwrap();
|
||||
stdin.flush().await.unwrap();
|
||||
}
|
||||
async fn recv(stdout: &mut BufReader<tokio::process::ChildStdout>) -> serde_json::Value {
|
||||
let mut line = String::new();
|
||||
timeout(Duration::from_secs(8), stdout.read_line(&mut line))
|
||||
.await
|
||||
.expect("recv timed out")
|
||||
.expect("recv io");
|
||||
serde_json::from_str(&line).expect("recv json")
|
||||
}
|
||||
|
||||
send(
|
||||
&mut stdin,
|
||||
json!({
|
||||
"jsonrpc": "2.0", "id": 1, "method": "initialize",
|
||||
"params": {"protocolVersion": 1, "clientCapabilities": {},
|
||||
"clientInfo": {"name": "auto-upgrade-test"}}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let init = recv(&mut stdout).await;
|
||||
assert!(init.get("result").is_some(), "initialize: {init}");
|
||||
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
send(
|
||||
&mut stdin,
|
||||
json!({
|
||||
"jsonrpc": "2.0", "id": 2, "method": "session/new",
|
||||
"params": {"cwd": cwd.to_string_lossy(), "mcpServers": []}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let sess = recv(&mut stdout).await;
|
||||
let sid = sess["result"]["sessionId"]
|
||||
.as_str()
|
||||
.unwrap_or_else(|| panic!("session/new failed: {sess}"))
|
||||
.to_string();
|
||||
|
||||
send(
|
||||
&mut stdin,
|
||||
json!({
|
||||
"jsonrpc": "2.0", "id": 3, "method": "session/prompt",
|
||||
"params": {"sessionId": sid,
|
||||
"prompt": [{"type": "text", "text": "hi"}]}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Drain notifications until we see the response for id=3.
|
||||
let mut stop_reason: Option<String> = None;
|
||||
for _ in 0..40 {
|
||||
let msg = recv(&mut stdout).await;
|
||||
if msg.get("id") == Some(&json!(3)) {
|
||||
if let Some(r) = msg.get("result") {
|
||||
stop_reason = r
|
||||
.get("stopReason")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert_eq!(stop_reason.as_deref(), Some("end_turn"));
|
||||
assert_eq!(
|
||||
chat_hits.load(Ordering::SeqCst),
|
||||
1,
|
||||
"must have tried chat first"
|
||||
);
|
||||
assert!(
|
||||
resp_hits.load(Ordering::SeqCst) >= 1,
|
||||
"must have upgraded to responses"
|
||||
);
|
||||
|
||||
drop(stdin);
|
||||
let _ = child.wait().await;
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user