Compare commits

..
Author SHA1 Message Date
Dark-Alex-17 f44722df04 lint: Removed accidental session file
CI / All (ubuntu-latest) (push) Failing after 32s
CI / All (macos-latest) (push) Canceled after 0s
CI / All (windows-latest) (push) Canceled after 0s
2026-08-19 12:59:50 -06:00
Dark-Alex-17 3eaae0e652 feat: Dynamically detect RAG embedding model dimension for any given model 2026-08-18 19:57:11 -06:00
Dark-Alex-17 b12829db39 fix: drain crossterm characters in zellij in kitty contexts to prevent DA1 responses from entering prompt 2026-08-18 15:30:07 -06:00
Dark-Alex-17 dffaf6b9db fix: drain tty input when displaying inquire prompts to prevent unintentional escapes 2026-08-18 15:10:30 -06:00
Dark-Alex-17 a9a4ccca88 fix: breaking iwe MCP server changes with newest version 2026-08-18 11:29:36 -06:00
Dark-Alex-17 7cb7d66575 fix: improved handling of non service-specific secrets using sbx custom-secrets
CI / All (ubuntu-latest) (push) Failing after 30s
CI / All (macos-latest) (push) Canceled after 0s
CI / All (windows-latest) (push) Canceled after 0s
2026-08-17 18:01:21 -06:00
Dark-Alex-17 400b50fbd0 feat: Support auto confirmation for gatekeeper agents 2026-08-17 15:48:36 -06:00
Dark-Alex-17 7eeff2a226 fix: include tool output to LLM_OUTPUT in errors as well as stderr 2026-08-17 11:38:21 -06:00
Dark-Alex-17 374e8c0bf7 fix: prevent tmp-overwriting
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Canceled after 0s
CI / All (windows-latest) (push) Canceled after 0s
2026-08-15 19:35:12 -06:00
Dark-Alex-17 ede87df960 fix: Corrected a rare edge case on how tool files are generated during parallel executions 2026-08-15 18:41:13 -06:00
Dark-Alex-17 95a8c3df44 fix: cosmetic fix after improved bash tool handling 2026-08-15 17:31:32 -06:00
Dark-Alex-17 f57bd21ee4 fix: Improved subagent escalation handling 2026-08-14 16:53:37 -06:00
Dark-Alex-17 e1b6e3f8c6 fix: latent parsing bugs in fs_patch and argc 2026-08-14 16:12:41 -06:00
Dark-Alex-17 644d899f78 fix: Prevent infinite hangs in coder agent and implement timeouts for LLM API calls and interactive tools 2026-08-14 15:22:25 -06:00
Dark-Alex-17 2596194417 feat: retry LLM API calls once after 401 by force-refreshing the OAuth token
The Client trait's default chat_completions, chat_completions_streaming,
and embeddings methods now classify failures via ApiStatusError: on a
401 with a cached OAuth token, the token is distrusted (identity-aware
marker) and the call retried exactly once — the retry's prepare step
sees the marker and force-refreshes. Streaming retries only while the
SSE handler has received no content, preventing duplicate rendering.
A second 401 propagates the original error; other retry errors
propagate as-is. API-key clients never retry. No backoff by design:
cost is bounded to one refresh + one retry per failing request.
2026-08-14 13:17:56 -06:00
Dark-Alex-17 684f19250a feat: identity-aware rejected-token marker for LLM OAuth cache
distrust_access_token compare-and-invalidates the in-memory entry only
when the cached token equals the rejected one, so a concurrent refresh
is never clobbered. is_valid_access_token and both expiry checks in
prepare_oauth_access_token treat marked tokens as expired, forcing a
refresh of provider-rejected tokens that are still locally unexpired.
The marker is cleared after every completed refresh, including ones
that return the same token.
2026-08-14 13:07:51 -06:00
Dark-Alex-17 f3d59ade11 feat: typed ApiStatusError carrying HTTP status through LLM client errors
catch_error and sse_stream now bail with ApiStatusError{status, message}
instead of bare anyhow strings, preserving every existing Display output
byte-for-byte. Enables structural status classification (e.g. 401
detection) via downcast through anyhow context chains.
2026-08-14 13:01:08 -06:00
Dark-Alex-17 e1604c58ea feat: per-request OAuth token injection with mid-session refresh for HTTP MCP servers
Replace the spawn-time static Authorization header for OAuth-managed HTTP
MCP servers with McpOAuthClient, a custom implementation of rmcp's
StreamableHttpClient trait that resolves the bearer token on every
request via load_or_refresh_mcp_token. Tokens that expire mid-session
now refresh transparently instead of failing tool calls until restart.

On a 401 for an injected token, the wrapper force-refreshes (identity-
aware: a still-unexpired copy of the rejected token is not trusted) and
retries exactly once, matching Claude Code / official SDK semantics.
Both *_with_max_sse_event_size trait methods are overridden to preserve
the inner client's SSE size enforcement, and the inner reqwest client
mirrors rmcp's default (pool_max_idle_per_host(0), no redirects).

SSE, stdio, and static-header HTTP paths are unchanged; startup
warning semantics (McpAuthRequired reasons) are preserved. Verified
live: mid-session backdated token refreshed transparently during an
active atlassian session.
2026-08-14 12:33:46 -06:00
Dark-Alex-17 dcacb3a962 chore: upgrade rmcp 1.8.0 -> 3.1.2
Compile-clean upgrade verified: zero source changes needed, full test
suite green, clippy clean. Coyote's rmcp API surface (14 items) dodges
all 2.0/3.0 breaking changes; the "Auth required" error string matched
by is_auth_required_error is intact in 3.1.2.
2026-08-14 11:10:45 -06:00
Dark-Alex-17 d791098e51 feat: reason-specific warnings for MCP servers that fail OAuth at startup
Distinguish why an OAuth MCP server was not started: never authenticated
(no stored credentials), stored token expired and refresh failed, or the
server rejected a token that looked valid. McpTokenStatus replaces the
Option<String> return of load_or_refresh_mcp_token, and McpAuthRequired
carries the reason across the error boundary via anyhow context.
2026-08-14 11:07:56 -06:00
Dark-Alex-17 d31110cd67 fix: allow nested italics inside bold spans in markdown renderer
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-13 16:18:50 -06:00
Dark-Alex-17 0f35e03a85 fix: properly handle OAuth refreshes
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-13 16:00:42 -06:00
Dark-Alex-17 c2b0c120d7 chore: Added grok4.6 to models.yaml
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-13 14:20:35 -06:00
Dark-Alex-17 65c9be36b2 fix: correct newline removal from fs_write and fs_patch 2026-08-13 14:18:45 -06:00
Dark-Alex-17 2658ca776e feat: Installed duckdb into the coyote image
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-12 19:29:26 -06:00
Dark-Alex-17 f8682102a0 docs: Added duckdb prerequisite
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-12 19:07:31 -06:00
Dark-Alex-17 3fa0f5c428 feat: Support managing MCP servers from the CLI directly
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-12 17:55:33 -06:00
Dark-Alex-17 68135b97d1 feat: append new built-in rag__query function to RAG contexts to allow further querying by LLMs
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-12 16:25:50 -06:00
Dark-Alex-17 b87a3460c4 fix: detect duplicate tool call IDs client-side before sending to Claude
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-12 13:33:24 -06:00
Dark-Alex-17 cb23da6490 feat: support static file bundling with sbx-mixins
CI / All (ubuntu-latest) (push) Failing after 29s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
2026-08-12 13:18:38 -06:00
Alex Clarke 2a40a5a81d Merge pull request #14 from Dark-Alex-17/feat/rag-driver-abstraction-v3
CI / All (ubuntu-latest) (push) Failing after 31s
CI / All (macos-latest) (push) Has been cancelled
CI / All (windows-latest) (push) Has been cancelled
feat(rag): pluggable RagProvider abstraction with DuckDB and Qdrant drivers
2026-08-12 12:45:12 -06:00
Dark-Alex-17 ebba976a27 fmt: applied formatting 2026-08-12 12:07:38 -06:00
Dark-Alex-17 c84f9522e9 feat(rag): offer the storage driver when an agent initializes its RAG
Agent startup and graph rag nodes both run an interactive wizard when their
knowledge base has not been built, but neither offered the driver choice that
interactive named-RAG creation has, so both silently produced a yaml store.

A plain agent was the worse of the two: AgentConfig carries only documents, so
there was no way to get a duckdb RAG for one, interactively or declaratively. A
graph node could at least declare driver: in the workflow.

Agent startup now passes prompt_for_driver, and a rag node whose wizard runs is
asked too. The prompt is skipped when the node already declares a driver, and
sits inside the not-fully-specified branch after the non-interactive bail, so
declarative workflows and headless runs are unchanged. Temp RAGs still pass
false: they are deleted on the next run, so a persistent store would only leave
a sidecar behind.

The prompt moves to select_rag_driver rather than being duplicated.
2026-08-12 12:02:15 -06:00
Dark-Alex-17 81ed769f8a fix(rag): warn when a duckdb store is empty but files are indexed
A duckdb RAG is two files. The .yaml deliberately carries no vectors, and
open() runs CREATE TABLE IF NOT EXISTS, so a .yaml copied or synced without its
.duckdb sidecar produces a fresh empty store, hydrates to nothing, and answers
every query with nothing while .info rag still lists every indexed file.

Neither existing guard catches it: the anti-wipe check in rebuild_indexes needs
existing > 0, and the mandatory ? on hydration needs a genuine error, while an
absent store is the same Ok(empty) as a RAG with nothing indexed yet.

Warn rather than bail, so a store deleted on purpose still loads and can be
rebuilt.
2026-08-12 11:26:16 -06:00
Dark-Alex-17 6f7defe25f style: removed redundant comment 2026-08-12 10:56:34 -06:00
Dark-Alex-17 b837f82d7e fix(rag): keep a local Qdrant off an ambient proxy
Reverts the global proxy rework in 54685be and narrows it to the provider.

That commit took over proxy detection for every client in order to exempt
loopback and private ranges. Too broad: reqwest's detection also reads macOS
System Settings and the Windows registry behind its system-proxy feature, which
sits in its default set. Coyote disables default features today, so hand-rolling
the environment lookup happened to match — but re-enabling defaults later would
silently restore that support for main and not for the hand-rolled version. It
also made an explicitly configured proxy skip local hosts, which nobody asked
for: a proxy named for a LAN endpoint should be used.

build_client and utils are byte-identical to main again. The bypass now lives in
QdrantProvider::make_client, which is the only place that knows the target host,
and applies solely when that host is loopback, link-local, private or .local. A
public or cloud-hosted store keeps whatever the environment configures.

Also drops apply_proxy: with build_client reverted there was one caller left, and
set_proxy already covers it.

Both #[ignore]d live tests still pass against a Qdrant on loopback while an
ambient proxy that rejects it is in force.
2026-08-12 10:52:31 -06:00
Dark-Alex-17 54685be9a2 fix: keep loopback and LAN traffic off an ambient proxy
Pre-existing on main, not introduced by the driver work, but it makes a local
RAG backend unusable so it belongs with this change.

build_client only called set_proxy when a client had configured one of its own.
With nothing configured, reqwest's own detection applied, which sends every
request through a *_PROXY variable including ones bound for 127.0.0.1 or a LAN
address. A proxy cannot usefully forward those, and anything that intercepts
proxied traffic answers on behalf of a service that is running perfectly well,
so the error names the proxy rather than the store and reads as a Coyote fault.

Concretely, an installed Socket Firewall exports HTTP_PROXY to the processes it
wraps and rejects hosts outside its allow list. That turned a healthy Ollama on
the LAN into 'error decoding response body: expected value at line 2 column 1' —
its HTML refusal page parsed as JSON — and a loopback Qdrant into an HTTP 405.

Proxy handling is now always applied and always exempts loopback and private
ranges, with NO_PROXY merged in since replacing reqwest's detection also
replaces its handling of that variable. HTTP_PROXY and HTTPS_PROXY are kept
separate because they are allowed to differ. An explicitly configured proxy
still wins, and '-' still means none.

This also supersedes the unconditional no_proxy() added to the Qdrant client in
af9622d: that made it the only client to ignore a proxy outright, on a
justification I got wrong. It now shares this path, so a remote store behind a
real proxy keeps working.
2026-08-12 10:33:03 -06:00
Dark-Alex-17 4dd6e794b2 docs: removed redundant comment 2026-08-11 22:19:14 -06:00
Dark-Alex-17 af9622d31c fix(rag): stop routing Qdrant requests through an ambient proxy
make_client used a bare reqwest builder, which honours whatever proxy the
environment advertises. That made it the only HTTP client in Coyote to do so:
utils::set_proxy discards ambient settings and applies only Coyote's configured
proxy, and every other client goes through it.

The symptom is that a perfectly healthy Qdrant is unreachable and the error
belongs to the interposing proxy, not the store, so it reads as a Coyote or
Qdrant fault. Locally an installed Socket Firewall answered `.rag attach`
against 127.0.0.1:6333 with an HTML 'Connection Required' page and HTTP 405.

Both #[ignore]d live tests now pass against a real Qdrant; they failed with that
same 405 before this change, which is the first time either has run green.

A remote store that genuinely needs Coyote's configured proxy is a follow-up:
that means threading the proxy config into the provider.
2026-08-11 22:15:04 -06:00
Dark-Alex-17 6f586bd535 style: cleanup 2026-08-11 22:07:35 -06:00
Dark-Alex-17 78740db170 chore: ignore the .coyote workspace directory
It holds generated workspace state and should never be committed.
2026-08-11 22:03:45 -06:00
Dark-Alex-17 64d594f4ee refactor(rag): discover driver_config secrets by grammar, not field name
Sandbox provisioning only ever looked at driver_config["api_key"], so a driver
whose credential is called anything else would have been silently unprovisioned
inside a sandbox. It now scans every driver_config value and treats any that is
a secret placeholder as a credential, which is the same rule resolve_driver_config
already used at point of use.

The first one binds to the RAG's own service id, which is what the generated
mixin declares; any others register under their own names, as MCP secrets do.
The mixin still carries a single credential entry, so a driver needing two bound
secrets remains a follow-up.

Also drops the placeholder parser added in 74bc613. crate::vault::SECRET_RE is
already the canonical definition and was already imported here, so that was a
third implementation of the same grammar. Requiring the whole value to match is
what keeps a literal key from being read as a secret name and printed.

The api_key check is gone from RagData::validate: a generic config validator
should not know a provider's field names.
2026-08-11 22:03:45 -06:00
Dark-Alex-17 1322d73c7b style: further cleanup 2026-08-11 21:50:46 -06:00
Dark-Alex-17 de91ffa517 fix(rag): treat a zero min_score as no floor on Qdrant searches
parse_search_hits filtered on score > min_score, and the only caller passes
0.0. Qdrant Euclid collections score by negative distance, so every hit was
dropped and an attached Euclid collection returned nothing at all, silently.

This is the same trap the surrounding code already documents: score_threshold
is deliberately not sent because it is metric-aware and a 0.0 floor filters
everything out on Euclid. The local filter then reproduced it exactly. Only a
positive floor is now treated as a floor.
2026-08-11 21:07:20 -06:00
Dark-Alex-17 74bc613d94 fix(rag): address Copilot review findings on the driver abstraction
Five review comments, all real:

- hybrid_search ran its vector and keyword legs sequentially after the
  provider refactor; main ran them under tokio::join!. Restores the
  concurrency while keeping the degrade-on-error keyword behaviour, so a
  remote provider no longer pays two serial round trips per query.

- inject_rag_secrets derived a vault secret name by trimming braces, which
  leaves a literal key untouched. A RAG holding a plaintext api_key therefore
  looked the secret up by its own value and printed it to stderr on failure.
  Parsing is now strict and a non-placeholder is skipped with a warning that
  names no credential.

- validate() now refuses a driver_config.api_key that is not a {{NAME}}
  placeholder, so a plaintext key cannot reach the RAG YAML at all.

- Rag::create's catch-all arm treated any unrecognised driver as yaml. A typo
  built a yaml store, paid to embed the corpus, persisted the bad driver and
  only failed on the next run. Unknown drivers now fail immediately.

- The qdrant arm's error was written for a developer; it now tells the user
  that only attached collections are readable and points at .rag attach.
2026-08-11 21:04:21 -06:00
Dark-Alex-17 6d0a5550fe style: Removed some redundant comments 2026-08-11 20:56:51 -06:00
Dark-Alex-17 7b1c0342b4 fix(rag): delete the DuckDB write-ahead log alongside the store
Deleting a RAG removed its .duckdb file but left the sibling .duckdb.wal
behind. DuckDB only removes that log on a clean close, so any RAG whose
process was killed left one on disk, and creating a new RAG under the same
name let it inherit a write-ahead log describing someone else's data.

The test helper already cleaned the log up after itself, which is why no
test noticed the production path did not.
2026-08-11 16:55:19 -06:00
48 changed files with 6202 additions and 502 deletions
+4
View File
@@ -5,3 +5,7 @@
.idea/ .idea/
/coyote.iml /coyote.iml
/.idea/ /.idea/
.coyote/**
.sisyphus/**
.coyote-project.json
.coyote/memory/
Generated
+43 -8
View File
@@ -1962,6 +1962,16 @@ dependencies = [
"darling_macro 0.23.0", "darling_macro 0.23.0",
] ]
[[package]]
name = "darling"
version = "0.24.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "88490bf1b990d87eaaa7ac8aa887f629a08e7359765b4911faf63c3763347d23"
dependencies = [
"darling_core 0.24.0",
"darling_macro 0.24.0",
]
[[package]] [[package]]
name = "darling_core" name = "darling_core"
version = "0.20.11" version = "0.20.11"
@@ -1989,6 +1999,19 @@ dependencies = [
"syn 2.0.119", "syn 2.0.119",
] ]
[[package]]
name = "darling_core"
version = "0.24.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "084e274f91c482280130e1e34e0b8d6e66776a060d7b6de7b84289ca778868c4"
dependencies = [
"ident_case",
"proc-macro2",
"quote",
"strsim",
"syn 3.0.3",
]
[[package]] [[package]]
name = "darling_macro" name = "darling_macro"
version = "0.20.11" version = "0.20.11"
@@ -2011,6 +2034,17 @@ dependencies = [
"syn 2.0.119", "syn 2.0.119",
] ]
[[package]]
name = "darling_macro"
version = "0.24.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68f5792fa0d41cd2325ce0ffa64f0a340eaebd4971a3a0c5e1ffd2cc488a355e"
dependencies = [
"darling_core 0.24.0",
"quote",
"syn 3.0.3",
]
[[package]] [[package]]
name = "defmt" name = "defmt"
version = "1.1.1" version = "1.1.1"
@@ -5303,12 +5337,12 @@ dependencies = [
[[package]] [[package]]
name = "rmcp" name = "rmcp"
version = "1.8.0" version = "3.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d1f571c72940a19d9532fe52dbea8bc9912bf1d766c2970bb824056b86f3f59" checksum = "c8dddc5b1924b9a59fba420166160ca2c4663a4e01803e52eda33070f56d63c8"
dependencies = [ dependencies = [
"async-trait", "base64 0.23.1",
"base64 0.22.1", "bytes",
"chrono", "chrono",
"futures", "futures",
"http 1.5.0", "http 1.5.0",
@@ -5326,19 +5360,20 @@ dependencies = [
"tokio-stream", "tokio-stream",
"tokio-util", "tokio-util",
"tracing", "tracing",
"uuid",
] ]
[[package]] [[package]]
name = "rmcp-macros" name = "rmcp-macros"
version = "1.8.0" version = "3.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1aad0035b69380782d78ea95b508327e6deaa2235909053e596eea8f27b5e1d5" checksum = "6898e24cd16342b59bfa8a53c2c04b9cf62fc8a2cfea57b9c038b09984bfc521"
dependencies = [ dependencies = [
"darling 0.23.0", "darling 0.24.0",
"proc-macro2", "proc-macro2",
"quote", "quote",
"serde_json", "serde_json",
"syn 2.0.119", "syn 3.0.3",
] ]
[[package]] [[package]]
+2 -2
View File
@@ -84,7 +84,7 @@ duct = "1.0.0"
argc = "1.23.0" argc = "1.23.0"
strum_macros = "0.27.2" strum_macros = "0.27.2"
indoc = "2.0.6" indoc = "2.0.6"
rmcp = { version = "1.5.0", features = [ rmcp = { version = "3.1.2", features = [
"client", "client",
"transport-child-process", "transport-child-process",
"transport-streamable-http-client-reqwest", "transport-streamable-http-client-reqwest",
@@ -150,4 +150,4 @@ path = "src/main.rs"
[profile.release] [profile.release]
lto = true lto = true
strip = true strip = true
opt-level = "z" opt-level = "z"
+15
View File
@@ -36,6 +36,21 @@ RUN set -euo pipefail; \
install -m 0755 "$TMPDIR/usql_static" /usr/local/bin/usql; \ install -m 0755 "$TMPDIR/usql_static" /usr/local/bin/usql; \
rm -rf "$TMPDIR" rm -rf "$TMPDIR"
RUN set -euo pipefail; \
DUCKDB_VERSION=1.5.5; \
case "${TARGETARCH}" in \
amd64) DUCKDB_ARCH=amd64 ;; \
arm64) DUCKDB_ARCH=arm64 ;; \
*) echo "Unsupported TARGETARCH: ${TARGETARCH}" >&2; exit 1 ;; \
esac; \
TMPDIR=$(mktemp -d); \
curl -fsSL --retry 3 \
"https://github.com/duckdb/duckdb/releases/download/v${DUCKDB_VERSION}/duckdb_cli-linux-${DUCKDB_ARCH}.gz" \
-o "$TMPDIR/duckdb.gz"; \
gunzip "$TMPDIR/duckdb.gz"; \
install -m 0755 "$TMPDIR/duckdb" /usr/local/bin/duckdb; \
rm -rf "$TMPDIR"
USER 1000 USER 1000
RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \ RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+2
View File
@@ -70,6 +70,8 @@ Coyote requires the following tools to be installed on your system:
* **Cargo:** `cargo install ast-grep --locked` * **Cargo:** `cargo install ast-grep --locked`
* **npm:** `npm i -g @ast-grep/cli` * **npm:** `npm i -g @ast-grep/cli`
* Optional: if `ast-grep` is not installed, the `ast_grep` tool reports it and agents fall back to `fs_grep` * Optional: if `ast-grep` is not installed, the `ast_grep` tool reports it and agents fall back to `fs_grep`
* [duckdb](https://duckdb.org/) (for fast, local RAGs)
* `curl https://install.duckdb.org | sh`
These tools are used to provide various functionalities within Coyote, such as document processing, JSON manipulation, These tools are used to provide various functionalities within Coyote, such as document processing, JSON manipulation,
etc., and they are used within agents and tools. etc., and they are used within agents and tools.
+5 -2
View File
@@ -31,7 +31,7 @@ settings:
max_loop_iterations: 20 max_loop_iterations: 20
log_state_snapshots: true log_state_snapshots: true
validate_before_run: true validate_before_run: true
timeout: 1800 timeout: 14400
initial_state: initial_state:
project_dir: '' project_dir: ''
@@ -90,6 +90,7 @@ nodes:
Project directory: {{project_dir}} Project directory: {{project_dir}}
prompt: '{{initial_prompt}}' prompt: '{{initial_prompt}}'
tools: [] tools: []
timeout: 300
output_schema: output_schema:
type: object type: object
properties: properties:
@@ -254,6 +255,7 @@ nodes:
- fs_patch - fs_patch
- execute_command - execute_command
max_iterations: 100 max_iterations: 100
timeout: 1800
state_updates: state_updates:
last_node_output: '{{output}}' last_node_output: '{{output}}'
fallback: end_failure fallback: end_failure
@@ -327,6 +329,7 @@ nodes:
- fs_ls - fs_ls
- execute_command - execute_command
max_iterations: 15 max_iterations: 15
timeout: 600
output_schema: output_schema:
type: object type: object
properties: properties:
@@ -384,4 +387,4 @@ nodes:
{{build_output}} {{build_output}}
Last tests output: Last tests output:
{{tests_output}} {{tests_output}}
+3
View File
@@ -14,6 +14,9 @@ variables:
- name: project_dir - name: project_dir
description: Absolute path to the project the plan targets - the ground truth for pointer verification description: Absolute path to the project the plan targets - the ground truth for pointer verification
default: '.' default: '.'
- name: auto_confirm
description: Auto-confirm command execution
default: '1'
global_tools: global_tools:
- ast_grep.sh - ast_grep.sh
+1 -2
View File
@@ -21,8 +21,7 @@
}, },
"iwe": { "iwe": {
"type": "stdio", "type": "stdio",
"command": "iwec", "command": "iwec"
"args": ["--project", "."]
} }
} }
} }
+23 -2
View File
@@ -14,6 +14,11 @@ set -e
# is the most common cause of "unable to apply patch" failures, especially in files with sed/jq/regex pipelines or # is the most common cause of "unable to apply patch" failures, especially in files with sed/jq/regex pipelines or
# embedded Python with quoted strings. # embedded Python with quoted strings.
# - Hunks are applied in order; the first hunk that fails aborts the whole patch — later hunks are NOT attempted. # - Hunks are applied in order; the first hunk that fails aborts the whole patch — later hunks are NOT attempted.
# - Hunks anchor at the FIRST exact match of their context in the file, so include enough context to make each hunk
# unique. Hunks must appear in file order.
# - Every hunk needs at least one context or removed line to anchor it; a hunk containing only additions is an error.
# - Unrecognized lines inside a hunk are hard errors: every hunk line must start with ' ' (context), '-' (removal),
# or '+' (addition).
# - If you've edited this file in earlier tool calls, fs_cat it again before composing the patch. A stale view of the file # - If you've edited this file in earlier tool calls, fs_cat it again before composing the patch. A stale view of the file
# produces context lines that no longer match. # produces context lines that no longer match.
# - On failure the error message names the failing hunk and shows the expected-vs-actual line. Fix that specific line and # - On failure the error message names the failing hunk and shows the expected-vs-actual line. Fix that specific line and
@@ -33,7 +38,11 @@ source "$LLM_PROMPT_UTILS_FILE"
# shellcheck disable=SC2154 # shellcheck disable=SC2154
main() { main() {
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")" # Command substitution strips *all* trailing newlines and `jq -r` appends one
# of its own, so read with `-j` and pin the real end of the content with a
# sentinel that is removed afterwards.
argc_contents="$(jq -j '.content' <<< "$LLM_TOOL_RAW_JSON"; printf x)"
argc_contents="${argc_contents%x}"
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")" argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
if [[ ! -f "$argc_path" ]]; then if [[ ! -f "$argc_path" ]]; then
@@ -41,7 +50,19 @@ main() {
exit 1 exit 1
fi fi
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"))" # Same sentinel guard on the patched result, otherwise the trailing newline
# is stripped again on the way back out. `rc` preserves patch_file's exit
# status so a failure still aborts under `set -e`.
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"); rc=$?; printf x; exit "$rc")"
new_contents="${new_contents%x}"
# awk newline-terminates every printed line, so patching a file that lacks
# a final newline would silently add one. Preserve the original file's
# final-newline state instead.
if [[ -n "$(tail -c 1 "$argc_path")" ]]; then
new_contents="${new_contents%$'\n'}"
fi
printf "%s" "$new_contents" | git diff --no-index "$argc_path" - || true printf "%s" "$new_contents" | git diff --no-index "$argc_path" - || true
guard_operation "Apply changes?" guard_operation "Apply changes?"
+6 -1
View File
@@ -15,7 +15,12 @@ source "$LLM_PROMPT_UTILS_FILE"
# shellcheck disable=SC2154 # shellcheck disable=SC2154
main() { main() {
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")" # Command substitution strips *all* trailing newlines and `jq -r` appends one
# of its own, so read with `-j` and pin the real end of the content with a
# sentinel that is removed afterwards. Without this every written file loses
# its final newline, which breaks formatters such as `cargo fmt --check`.
argc_contents="$(jq -j '.content' <<< "$LLM_TOOL_RAW_JSON"; printf x)"
argc_contents="${argc_contents%x}"
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")" argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
if [[ -f "$argc_path" ]]; then if [[ -f "$argc_path" ]]; then
+92 -18
View File
@@ -186,7 +186,9 @@ input() {
} }
confirm() { confirm() {
trap "stty echo; exit" EXIT # stty targets stdin, which is /dev/null when the host spawns tool scripts;
# point it at the real terminal and stay quiet when there isn't one.
trap "stty echo </dev/tty 2>/dev/null; exit" EXIT
_prompt_text "$1 (y/N)" _prompt_text "$1 (y/N)"
echo -en "\033[36m\c " >&2 echo -en "\033[36m\c " >&2
@@ -229,7 +231,7 @@ list() {
declare first_row declare first_row
first_row=$((last_row - opts_count + 1)) first_row=$((last_row - opts_count + 1))
trap "_cursor_blink_on; stty echo; exit" 2 trap "_cursor_blink_on; stty echo </dev/tty 2>/dev/null; exit" 2
_cursor_blink_off _cursor_blink_off
@@ -275,7 +277,7 @@ checkbox() {
declare first_row declare first_row
first_row=$((last_row - opts_count + 1)) first_row=$((last_row - opts_count + 1))
trap "_cursor_blink_on; stty echo; exit" 2 trap "_cursor_blink_on; stty echo </dev/tty 2>/dev/null; exit" 2
_cursor_blink_off _cursor_blink_off
@@ -403,7 +405,7 @@ range() {
declare current_row declare current_row
current_row=$((first_row - 1)) current_row=$((first_row - 1))
trap "_cursor_blink_on; stty echo; exit" 2 trap "_cursor_blink_on; stty echo </dev/tty 2>/dev/null; exit" 2
_cursor_blink_off _cursor_blink_off
@@ -528,7 +530,11 @@ guard_operation() {
# + print(f"Hello {name}") # + print(f"Hello {name}")
patch_file() { patch_file() {
awk ' awk '
FNR == NR { function isHeaderPair(i) {
return (patchLines[i] ~ /^--- / && patchLines[i+1] ~ /^\+\+\+ / && patchLines[i+2] ~ /^@@/)
}
FILENAME == ARGV[1] {
lines[FNR] = $0 lines[FNR] = $0
next; next;
} }
@@ -547,11 +553,6 @@ patch_file() {
while (patchLineIndex <= totalPatchLines) { while (patchLineIndex <= totalPatchLines) {
line = patchLines[patchLineIndex] line = patchLines[patchLineIndex]
if (line ~ /^--- / || line ~ /^\+\+\+ /) {
patchLineIndex++
continue
}
if (line ~ /^@@/) { if (line ~ /^@@/) {
mode = "hunk" mode = "hunk"
hunkIndex++ hunkIndex++
@@ -560,7 +561,22 @@ patch_file() {
} }
if (mode == "hunk") { if (mode == "hunk") {
while (patchLineIndex <= totalPatchLines && line ~ /^[-+ ]|^\s*$/ && line !~ /^--- /) { while (patchLineIndex <= totalPatchLines) {
line = patchLines[patchLineIndex]
if (line ~ /^\\ No newline/) {
patchLineIndex++
continue
}
if (isHeaderPair(patchLineIndex)) {
break
}
if (line !~ /^[-+ ]/ && line !~ /^[ \t]*$/) {
break
}
sanitizedLine = substr(line, 2) sanitizedLine = substr(line, 2)
if (line !~ /^\+/) { if (line !~ /^\+/) {
@@ -574,13 +590,37 @@ patch_file() {
} }
patchLineIndex++ patchLineIndex++
line = patchLines[patchLineIndex]
} }
mode = "none" mode = "none"
} else { continue
patchLineIndex++
} }
if (isHeaderPair(patchLineIndex)) {
patchLineIndex += 2
continue
}
if (line ~ /^\\ No newline/) {
patchLineIndex++
continue
}
if (hunkIndex == 0) {
# Preamble before the first hunk: tolerate prose, code fences, and lone headers.
patchLineIndex++
continue
}
if (line ~ /^[ \t]*$/ || line ~ /^```/) {
patchLineIndex++
continue
}
print "error: unrecognized line in patch (line " patchLineIndex "): " line > "/dev/stderr"
print "" > "/dev/stderr"
print "Every line inside a hunk must start with \" \" (context), \"-\" (removal), or \"+\" (addition)." > "/dev/stderr"
exit 1
} }
if (hunkIndex == 0) { if (hunkIndex == 0) {
@@ -593,6 +633,36 @@ patch_file() {
} }
totalHunks = hunkIndex totalHunks = hunkIndex
if (totalLines == 0) {
for (h = 1; h <= totalHunks; h++) {
if (hunkTotalOriginalLines[h] > 0) {
print "error: unable to apply patch" > "/dev/stderr"
print "" > "/dev/stderr"
print "Hunk " h " expects existing content but the file is empty." > "/dev/stderr"
exit 1
}
}
for (h = 1; h <= totalHunks; h++) {
for (i = 1; i <= hunkTotalUpdatedLines[h]; i++) {
print hunkUpdatedLines[h,i]
}
}
exit 0
}
for (h = 1; h <= totalHunks; h++) {
if (hunkTotalOriginalLines[h] == 0) {
print "error: unable to apply patch" > "/dev/stderr"
print "" > "/dev/stderr"
print "Hunk " h " contains no context or removed lines; include at least one" > "/dev/stderr"
print "context line so the hunk can be anchored." > "/dev/stderr"
exit 1
}
}
hunkIndex = 1 hunkIndex = 1
for (lineIndex = 1; lineIndex <= totalLines; lineIndex++) { for (lineIndex = 1; lineIndex <= totalLines; lineIndex++) {
@@ -603,7 +673,7 @@ patch_file() {
nextLineIndex = lineIndex + 1 nextLineIndex = lineIndex + 1
for (i = 2; i <= hunkTotalOriginalLines[hunkIndex]; i++) { for (i = 2; i <= hunkTotalOriginalLines[hunkIndex]; i++) {
if (lines[nextLineIndex] != hunkOriginalLines[hunkIndex,i]) { if (nextLineIndex > totalLines || lines[nextLineIndex] != hunkOriginalLines[hunkIndex,i]) {
if (i - 1 > bestPartialLen[hunkIndex]) { if (i - 1 > bestPartialLen[hunkIndex]) {
bestPartialLen[hunkIndex] = i - 1 bestPartialLen[hunkIndex] = i - 1
bestPartialAnchorLine[hunkIndex] = lineIndex bestPartialAnchorLine[hunkIndex] = lineIndex
@@ -646,9 +716,13 @@ patch_file() {
print "" > "/dev/stderr" print "" > "/dev/stderr"
print "Closest match: anchored at file line " bestPartialAnchorLine[failingHunk] ", matched " bestPartialLen[failingHunk] " of " hunkTotalOriginalLines[failingHunk] " original lines before diverging." > "/dev/stderr" print "Closest match: anchored at file line " bestPartialAnchorLine[failingHunk] ", matched " bestPartialLen[failingHunk] " of " hunkTotalOriginalLines[failingHunk] " original lines before diverging." > "/dev/stderr"
print "" > "/dev/stderr" print "" > "/dev/stderr"
print "At file line " bestPartialDivergeLine[failingHunk] " (hunk original line " bestPartialHunkPos[failingHunk] "):" > "/dev/stderr" if (bestPartialDivergeLine[failingHunk] > totalLines) {
print " expected: " bestPartialExpected[failingHunk] > "/dev/stderr" print "The hunk expects additional lines beyond the end of the file (file has " totalLines " lines)." > "/dev/stderr"
print " actual: " bestPartialActual[failingHunk] > "/dev/stderr" } else {
print "At file line " bestPartialDivergeLine[failingHunk] " (hunk original line " bestPartialHunkPos[failingHunk] "):" > "/dev/stderr"
print " expected: " bestPartialExpected[failingHunk] > "/dev/stderr"
print " actual: " bestPartialActual[failingHunk] > "/dev/stderr"
}
} }
print "" > "/dev/stderr" print "" > "/dev/stderr"
+1 -1
View File
@@ -2,7 +2,7 @@
description: Navigate and curate markdown knowledge bases (plan repos, spec repos, companion docs) with IWE graph tools. Load when the workspace is or contains a markdown knowledge base and the task involves finding, reading, or reorganizing plans, specs, designs, or notes. Activates the iwe MCP server rooted at the current directory. description: Navigate and curate markdown knowledge bases (plan repos, spec repos, companion docs) with IWE graph tools. Load when the workspace is or contains a markdown knowledge base and the task involves finding, reading, or reorganizing plans, specs, designs, or notes. Activates the iwe MCP server rooted at the current directory.
enabled_mcp_servers: iwe enabled_mcp_servers: iwe
--- ---
You are working with a markdown knowledge base through IWE, a graph-based knowledge tool. The `iwe` MCP server is rooted at the current working directory (`--project .`), so the knowledge base is the directory Coyote was launched in. IWE derives structure from links: a link on its own line is an *inclusion link* (parent-child hierarchy); a link inside text is an *inline reference* (cross-reference, produces backlinks). The server watches the filesystem, so external edits are picked up automatically — never ask for a restart. You are working with a markdown knowledge base through IWE, a graph-based knowledge tool. The `iwe` MCP server is rooted at the current working directory, so the knowledge base is the directory Coyote was launched in. IWE derives structure from links: a link on its own line is an *inclusion link* (parent-child hierarchy); a link inside text is an *inline reference* (cross-reference, produces backlinks). The server watches the filesystem, so external edits are picked up automatically — never ask for a restart.
## When to use this (and when not) ## When to use this (and when not)
+2 -1
View File
@@ -292,6 +292,7 @@ clients:
# extra: # extra:
# proxy: socks5://127.0.0.1:1080 # Set proxy # proxy: socks5://127.0.0.1:1080 # Set proxy
# connect_timeout: 10 # Set timeout in seconds for connect to api # connect_timeout: 10 # Set timeout in seconds for connect to api
# read_timeout: 300 # Set timeout in seconds for a read stall (no bytes received); 0 disables (default: 300)
# See https://platform.openai.com/docs/quickstart # See https://platform.openai.com/docs/quickstart
- type: openai - type: openai
@@ -525,4 +526,4 @@ clients:
- type: openai-compatible - type: openai-compatible
name: voyageai name: voyageai
api_base: https://api.voyageai.com/v1 api_base: https://api.voyageai.com/v1
api_key: '{{VOYAGEAI_API_KEY}}' # You can either hard-code or inject secrets from the Coyote vault api_key: '{{VOYAGEAI_API_KEY}}' # You can either hard-code or inject secrets from the Coyote vault
+5
View File
@@ -841,6 +841,11 @@
referrer: coyote referrer: coyote
echo_pkce_in_token_exchange: true echo_pkce_in_token_exchange: true
models: models:
- name: grok-4.6
input_price: 2
output_price: 6
max_input_tokens: 500000
supports_function_calling: true
- name: grok-4.5 - name: grok-4.5
input_price: 2 input_price: 2
output_price: 6 output_price: 6
+161 -1
View File
@@ -13,6 +13,20 @@ use is_terminal::IsTerminal;
use std::collections::HashSet; use std::collections::HashSet;
use std::io::{Read, stdin}; use std::io::{Read, stdin};
#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
pub enum McpTransportArg {
Stdio,
Http,
Sse,
}
#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum McpScopeArg {
#[default]
User,
Workspace,
}
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)] #[command(author, version, about, long_about = None)]
#[command( #[command(
@@ -41,10 +55,15 @@ use std::io::{Read, stdin};
"list_skills", "skill", "tail_logs", "completions", "update", "list_skills", "skill", "tail_logs", "completions", "update",
]) ])
), ),
group(
ArgGroup::new("mcp-action")
.args(["mcp_add", "mcp_remove", "mcp_list", "mcp_get"])
.multiple(false)
),
)] )]
pub struct Cli { pub struct Cli {
/// Input text /// Input text
#[arg(trailing_var_arg = true)] #[arg(allow_hyphen_values = true)]
text: Vec<String>, text: Vec<String>,
/// Select a LLM model /// Select a LLM model
@@ -224,6 +243,57 @@ pub struct Cli {
#[arg(long, exclusive = true, value_name = "SERVER_NAME", help_heading = "Authentication", add = ArgValueCompleter::new(mcp_server_completer))] #[arg(long, exclusive = true, value_name = "SERVER_NAME", help_heading = "Authentication", add = ArgValueCompleter::new(mcp_server_completer))]
pub auth_mcp: Option<String>, pub auth_mcp: Option<String>,
/// Add an MCP server. Use `-- <cmd> [args...]` for stdio, or `--url <URL>` for http/sse.
#[arg(long, value_name = "NAME", help_heading = "MCP Servers")]
pub mcp_add: Option<String>,
/// Remove an MCP server by name
#[arg(long, value_name = "NAME", help_heading = "MCP Servers", add = ArgValueCompleter::new(mcp_server_completer))]
pub mcp_remove: Option<String>,
/// List all configured MCP servers (user + workspace scopes)
#[arg(long, help_heading = "MCP Servers")]
pub mcp_list: bool,
/// Show the JSON config for one MCP server
#[arg(long, value_name = "NAME", help_heading = "MCP Servers", add = ArgValueCompleter::new(mcp_server_completer))]
pub mcp_get: Option<String>,
/// Transport for --mcp-add: stdio (default when `--` present), http, or sse
#[arg(
long,
value_enum,
value_name = "TRANSPORT",
help_heading = "MCP Servers"
)]
pub transport: Option<McpTransportArg>,
/// URL for http/sse MCP server (used with --mcp-add)
#[arg(long, value_name = "URL", help_heading = "MCP Servers")]
pub url: Option<String>,
/// Scope for MCP config: user (~/.config/coyote/functions/mcp.json) or workspace (./.coyote/mcp.json). Default: user
#[arg(long, value_enum, value_name = "SCOPE", help_heading = "MCP Servers")]
pub scope: Option<McpScopeArg>,
/// Environment variable for stdio MCP server (repeatable): --env KEY=VALUE
#[arg(long, value_name = "KEY=VALUE", help_heading = "MCP Servers")]
pub env: Vec<String>,
/// HTTP header for http/sse MCP server (repeatable): --header "Name: Value"
#[arg(long, value_name = "HEADER", help_heading = "MCP Servers")]
pub header: Vec<String>,
/// Working directory for stdio MCP server
#[arg(long, value_name = "PATH", value_hint = ValueHint::AnyPath, help_heading = "MCP Servers")]
pub cwd: Option<String>,
/// OAuth client ID for http/sse MCP server
#[arg(long, value_name = "ID", help_heading = "MCP Servers")]
pub client_id: Option<String>,
/// OAuth client secret for http/sse MCP server (use {{NAME}} to reference a vault secret)
#[arg(long, value_name = "SECRET", help_heading = "MCP Servers")]
pub client_secret: Option<String>,
/// OAuth callback port for http/sse MCP server
#[arg(long, value_name = "PORT", help_heading = "MCP Servers")]
pub callback_port: Option<u16>,
/// OAuth redirect host for http/sse MCP server
#[arg(long, value_name = "HOST", help_heading = "MCP Servers")]
pub redirect_host: Option<String>,
/// Overwrite an existing MCP server (with --mcp-add) or skip confirmation (with --mcp-remove)
#[arg(long, help_heading = "MCP Servers")]
pub mcp_force: bool,
/// Launch Coyote inside a Docker sandbox (via `sbx`); name defaults to current directory basename /// Launch Coyote inside a Docker sandbox (via `sbx`); name defaults to current directory basename
#[arg(long, value_name = "NAME", help_heading = "Sandbox")] #[arg(long, value_name = "NAME", help_heading = "Sandbox")]
pub sandbox: Option<Option<String>>, pub sandbox: Option<Option<String>>,
@@ -254,6 +324,15 @@ pub struct Cli {
/// Generate static shell completion scripts /// Generate static shell completion scripts
#[arg(long, value_name = "SHELL", value_enum, help_heading = "Shell")] #[arg(long, value_name = "SHELL", value_enum, help_heading = "Shell")]
pub completions: Option<ShellCompletion>, pub completions: Option<ShellCompletion>,
/// Stdio command for --mcp-add: everything after `--` is passed to the server verbatim
#[arg(
last = true,
allow_hyphen_values = true,
value_name = "CMD",
help_heading = "MCP Servers"
)]
pub mcp_command: Vec<String>,
} }
impl Cli { impl Cli {
@@ -633,4 +712,85 @@ mod tests {
fn parse_sandbox_is_exclusive() { fn parse_sandbox_is_exclusive() {
assert!(Cli::try_parse_from(["coyote", "--sandbox", "--agent", "foo"]).is_err()); assert!(Cli::try_parse_from(["coyote", "--sandbox", "--agent", "foo"]).is_err());
} }
#[test]
fn parse_mcp_add_stdio_with_trailing_command() {
let cli = parse(&[
"--mcp-add",
"myserver",
"--",
"npx",
"some-server",
"--flag",
"arg1",
]);
assert_eq!(cli.mcp_add, Some("myserver".to_string()));
assert_eq!(
cli.mcp_command,
vec!["npx", "some-server", "--flag", "arg1"]
);
assert!(cli.text.is_empty());
}
#[test]
fn parse_mcp_add_stdio_with_env_and_command() {
let cli = parse(&[
"--mcp-add",
"s",
"--env",
"API_KEY={{API_KEY}}",
"--env",
"MODE=dev",
"--",
"npx",
"srv",
]);
assert_eq!(cli.mcp_add, Some("s".to_string()));
assert_eq!(cli.env, vec!["API_KEY={{API_KEY}}", "MODE=dev"]);
assert_eq!(cli.mcp_command, vec!["npx", "srv"]);
}
#[test]
fn parse_mcp_add_http_with_header() {
let cli = parse(&[
"--mcp-add",
"notion",
"--transport",
"http",
"--url",
"https://mcp.notion.com/mcp",
"--header",
"Authorization: Bearer {{NOTION_TOKEN}}",
]);
assert_eq!(cli.mcp_add, Some("notion".to_string()));
assert!(matches!(cli.transport, Some(McpTransportArg::Http)));
assert_eq!(cli.url, Some("https://mcp.notion.com/mcp".to_string()));
assert_eq!(cli.header, vec!["Authorization: Bearer {{NOTION_TOKEN}}"]);
assert!(cli.mcp_command.is_empty());
}
#[test]
fn parse_mcp_list_flag() {
let cli = parse(&["--mcp-list"]);
assert!(cli.mcp_list);
}
#[test]
fn parse_mcp_scope_workspace() {
let cli = parse(&["--mcp-list", "--scope", "workspace"]);
assert!(cli.mcp_list);
assert!(matches!(cli.scope, Some(McpScopeArg::Workspace)));
}
#[test]
fn parse_mcp_action_group_is_exclusive() {
assert!(Cli::try_parse_from(["coyote", "--mcp-list", "--mcp-get", "foo"]).is_err());
}
#[test]
fn parse_trailing_text_unchanged_without_dash_dash() {
let cli = parse(&["hello", "world"]);
assert_eq!(cli.text, vec!["hello", "world"]);
assert!(cli.mcp_command.is_empty());
}
} }
+116 -1
View File
@@ -2,6 +2,7 @@ use anyhow::{Result, anyhow};
use chrono::Utc; use chrono::Utc;
use indexmap::IndexMap; use indexmap::IndexMap;
use parking_lot::RwLock; use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::LazyLock; use std::sync::LazyLock;
type AccessTokenEntry = (String, i64, Option<String>); type AccessTokenEntry = (String, i64, Option<String>);
@@ -9,6 +10,12 @@ type AccessTokenEntry = (String, i64, Option<String>);
static ACCESS_TOKENS: LazyLock<RwLock<IndexMap<String, AccessTokenEntry>>> = static ACCESS_TOKENS: LazyLock<RwLock<IndexMap<String, AccessTokenEntry>>> =
LazyLock::new(|| RwLock::new(IndexMap::new())); LazyLock::new(|| RwLock::new(IndexMap::new()));
/// Tokens a provider rejected (401) despite being locally unexpired.
/// Maps client name → the exact rejected token so a concurrently-refreshed
/// different token is never distrusted by mistake.
static REJECTED_TOKENS: LazyLock<RwLock<HashMap<String, String>>> =
LazyLock::new(|| RwLock::new(HashMap::new()));
pub fn get_access_token(client_name: &str) -> Result<String> { pub fn get_access_token(client_name: &str) -> Result<String> {
ACCESS_TOKENS ACCESS_TOKENS
.read() .read()
@@ -30,7 +37,7 @@ pub fn is_valid_access_token(client_name: &str) -> bool {
Some(v) => v, Some(v) => v,
None => return false, None => return false,
}; };
!token.is_empty() && Utc::now().timestamp() < *expires_at !token.is_empty() && Utc::now().timestamp() < *expires_at && !is_rejected(client_name, token)
} }
pub fn set_access_token( pub fn set_access_token(
@@ -45,3 +52,111 @@ pub fn set_access_token(
entry.1 = expires_at; entry.1 = expires_at;
entry.2 = account_id; entry.2 = account_id;
} }
/// Compare-and-invalidate a provider-rejected token.
///
/// Only if the currently-cached token EQUALS `rejected` is the cache entry
/// removed and the rejection marker recorded; a concurrently-refreshed
/// different token is left untouched and no marker is set.
///
/// Returns true if a cache entry existed for this client at all (whether or
/// not it matched `rejected`) — i.e. the client is token-authed and a retry
/// after refresh is worthwhile. Returns false when there is no entry
/// (API-key clients).
pub fn distrust_access_token(client_name: &str, rejected: &str) -> bool {
let mut access_tokens = ACCESS_TOKENS.write();
let (token, _, _) = match access_tokens.get(client_name) {
Some(v) => v,
None => return false,
};
if token == rejected {
access_tokens.shift_remove(client_name);
REJECTED_TOKENS
.write()
.insert(client_name.to_string(), rejected.to_string());
}
true
}
pub fn is_rejected(client_name: &str, token: &str) -> bool {
REJECTED_TOKENS
.read()
.get(client_name)
.is_some_and(|rejected| rejected == token)
}
pub fn clear_rejected(client_name: &str) {
REJECTED_TOKENS.write().remove(client_name);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn distrust_removes_matching_token_and_sets_marker() {
let client = "distrust-match-test";
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
assert!(distrust_access_token(client, "at-1"));
assert!(get_access_token(client).is_err(), "cache entry not removed");
assert!(is_rejected(client, "at-1"), "marker not set");
}
#[test]
fn distrust_keeps_differing_token_and_skips_marker() {
let client = "distrust-differ-test";
set_access_token(client, "at-new".into(), Utc::now().timestamp() + 3600, None);
assert!(distrust_access_token(client, "at-old"));
assert_eq!(get_access_token(client).unwrap(), "at-new");
assert!(!is_rejected(client, "at-old"), "marker set for stale token");
}
#[test]
fn distrust_returns_false_without_cache_entry() {
let client = "distrust-missing-test";
assert!(!distrust_access_token(client, "at-1"));
assert!(!is_rejected(client, "at-1"));
}
#[test]
fn is_valid_access_token_false_for_rejected_token() {
let client = "rejected-valid-test";
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
assert!(is_valid_access_token(client));
distrust_access_token(client, "at-1");
// A concurrent in-flight prepare re-caches the rejected file token
// between mark and refresh; it must still be treated as invalid.
set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None);
assert!(!is_valid_access_token(client));
}
#[test]
fn clear_rejected_clears_marker_and_clients_are_isolated() {
let client_a = "rejected-isolation-a";
let client_b = "rejected-isolation-b";
set_access_token(client_a, "at-1".into(), Utc::now().timestamp() + 3600, None);
distrust_access_token(client_a, "at-1");
assert!(is_rejected(client_a, "at-1"));
assert!(
!is_rejected(client_b, "at-1"),
"marker leaked across clients"
);
clear_rejected(client_b);
assert!(
is_rejected(client_a, "at-1"),
"wrong client's marker cleared"
);
clear_rejected(client_a);
assert!(!is_rejected(client_a, "at-1"));
}
}
+115 -5
View File
@@ -1,3 +1,4 @@
use std::collections::HashSet;
use std::mem; use std::mem;
use super::access_token::get_access_token; use super::access_token::get_access_token;
@@ -368,17 +369,32 @@ pub fn claude_build_chat_completions_body(
] ]
} else { } else {
// One pair per round: Claude can reuse tool_use IDs across API calls. // One pair per round: Claude can reuse tool_use IDs across API calls.
// A round boundary is detected by the presence of round text, but
// rounds where the model emitted only tool calls (no narration)
// carry no text marker. As a backstop, also split whenever a
// tool_use ID would repeat within the current assistant message —
// the API rejects duplicate tool_use IDs in a single message.
let mut messages = vec![]; let mut messages = vec![];
let mut assistant_parts: Vec<serde_json::Value> = vec![]; let mut assistant_parts: Vec<serde_json::Value> = vec![];
let mut user_parts: Vec<serde_json::Value> = vec![]; let mut user_parts: Vec<serde_json::Value> = vec![];
let mut chunk_ids: HashSet<&str> = HashSet::new();
for (index, tool_result) in tool_results.iter().enumerate() { for (index, tool_result) in tool_results.iter().enumerate() {
if index > 0 && tool_result.text.is_some() { let id_collision = tool_result
.call
.id
.as_deref()
.is_some_and(|id| chunk_ids.contains(id));
if index > 0 && (tool_result.text.is_some() || id_collision) {
messages.push( messages.push(
json!({ "role": "assistant", "content": assistant_parts }), json!({ "role": "assistant", "content": assistant_parts }),
); );
messages.push(json!({ "role": "user", "content": user_parts })); messages.push(json!({ "role": "user", "content": user_parts }));
assistant_parts = vec![]; assistant_parts = vec![];
user_parts = vec![]; user_parts = vec![];
chunk_ids.clear();
}
if let Some(id) = tool_result.call.id.as_deref() {
chunk_ids.insert(id);
} }
for block in &tool_result.thinking { for block in &tool_result.thinking {
assistant_parts.push(json!(block)); assistant_parts.push(json!(block));
@@ -485,10 +501,7 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
if let Some(v) = item["thinking"].as_str() { if let Some(v) = item["thinking"].as_str() {
thinking.push(ThinkingBlock::Thinking { thinking.push(ThinkingBlock::Thinking {
thinking: v.to_string(), thinking: v.to_string(),
signature: item["signature"] signature: item["signature"].as_str().unwrap_or_default().to_string(),
.as_str()
.unwrap_or_default()
.to_string(),
}); });
} }
} }
@@ -535,3 +548,100 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
}; };
Ok(output) Ok(output)
} }
#[cfg(test)]
mod tests {
use super::*;
use crate::function::{ToolCall, ToolResult};
fn tool_result(id: &str, text: Option<&str>) -> ToolResult {
ToolResult {
call: ToolCall::new("fs_read".into(), json!({"path": "x"}), Some(id.into())),
output: json!("ok"),
text: text.map(|t| t.to_string()),
thinking: vec![],
}
}
fn build_body(tool_results: Vec<ToolResult>) -> Value {
let data = ChatCompletionsData {
messages: vec![
Message::new(MessageRole::User, MessageContent::Text("hello".to_string())),
Message::new(
MessageRole::Assistant,
MessageContent::ToolCalls(MessageContentToolCalls {
tool_results,
text: String::new(),
sequence: true,
}),
),
],
temperature: None,
top_p: None,
reasoning_effort: None,
functions: None,
stream: false,
};
claude_build_chat_completions_body(data, &Model::new("claude", "claude-test")).unwrap()
}
fn assert_unique_tool_use_ids_per_message(body: &Value) {
for message in body["messages"].as_array().unwrap() {
let Some(content) = message["content"].as_array() else {
continue;
};
let mut seen = HashSet::new();
for block in content {
if block["type"] == "tool_use" {
let id = block["id"].as_str().unwrap();
assert!(
seen.insert(id.to_string()),
"duplicate tool_use id `{id}` within a single assistant message: {message}"
);
}
}
}
}
#[test]
fn sequence_splits_on_round_text() {
let body = build_body(vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_C", Some("running another tool")),
]);
let messages = body["messages"].as_array().unwrap();
assert_eq!(messages.len(), 5, "body: {body}");
assert_unique_tool_use_ids_per_message(&body);
}
#[test]
fn sequence_splits_on_reused_id_in_textless_round() {
let body = build_body(vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_A", None),
]);
let messages = body["messages"].as_array().unwrap();
assert_eq!(messages.len(), 5, "body: {body}");
assert_unique_tool_use_ids_per_message(&body);
}
#[test]
fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() {
let body = build_body(vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_C", None),
]);
let messages = body["messages"].as_array().unwrap();
assert_eq!(messages.len(), 3, "body: {body}");
assert_unique_tool_use_ids_per_message(&body);
}
}
+289 -15
View File
@@ -1,5 +1,6 @@
use super::*; use super::*;
use super::access_token::{distrust_access_token, get_access_token};
use crate::config::{RenderMode, paths}; use crate::config::{RenderMode, paths};
use crate::{ use crate::{
config::{AppConfig, Input, RequestContext}, config::{AppConfig, Input, RequestContext},
@@ -56,12 +57,16 @@ pub trait Client: Sync + Send {
let mut builder = ReqwestClient::builder(); let mut builder = ReqwestClient::builder();
let extra = self.extra_config(); let extra = self.extra_config();
let timeout = extra.and_then(|v| v.connect_timeout).unwrap_or(10); let timeout = extra.and_then(|v| v.connect_timeout).unwrap_or(10);
let read_timeout = extra.and_then(|v| v.read_timeout).unwrap_or(300);
if let Some(proxy) = extra.and_then(|v| v.proxy.as_deref()) { if let Some(proxy) = extra.and_then(|v| v.proxy.as_deref()) {
builder = set_proxy(builder, proxy)?; builder = set_proxy(builder, proxy)?;
} }
if let Some(user_agent) = self.app_config().user_agent.as_ref() { if let Some(user_agent) = self.app_config().user_agent.as_ref() {
builder = builder.user_agent(user_agent); builder = builder.user_agent(user_agent);
} }
if read_timeout > 0 {
builder = builder.read_timeout(Duration::from_secs(read_timeout));
}
let client = builder let client = builder
.connect_timeout(Duration::from_secs(timeout)) .connect_timeout(Duration::from_secs(timeout))
.build() .build()
@@ -69,6 +74,11 @@ pub trait Client: Sync + Send {
Ok(client) Ok(client)
} }
/// On a 401 the cached access token is distrusted and the call retried
/// exactly once; the retry re-runs the per-client prepare step, which
/// sees the rejection marker, force-refreshes the token, and rebuilds
/// the whole request. A second 401 propagates the original error; any
/// other retry failure propagates as-is.
async fn chat_completions(&self, input: Input) -> Result<ChatCompletionsOutput> { async fn chat_completions(&self, input: Input) -> Result<ChatCompletionsOutput> {
if self.app_config().dry_run { if self.app_config().dry_run {
let content = input.echo_messages(); let content = input.echo_messages();
@@ -76,11 +86,30 @@ pub trait Client: Sync + Send {
} }
let client = self.build_client()?; let client = self.build_client()?;
let data = input.prepare_completion_data(self.model(), false)?; let data = input.prepare_completion_data(self.model(), false)?;
self.chat_completions_inner(&client, data) let err = match self.chat_completions_inner(&client, data).await {
.await Ok(output) => return Ok(output),
.with_context(|| "Failed to call chat-completions api") Err(err) => err,
};
let ret = if should_retry_auth(&err, self.name()) {
debug!(
"provider '{}' rejected access token (401); refreshing and retrying once",
self.name()
);
let data = input.prepare_completion_data(self.model(), false)?;
match self.chat_completions_inner(&client, data).await {
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
ret => ret,
}
} else {
Err(err)
};
ret.with_context(|| "Failed to call chat-completions api")
} }
/// Same retry-once-on-401 semantics as [`Self::chat_completions`], but
/// only while the handler has received nothing yet: retrying after
/// partial output has streamed would render it to the user twice. The
/// retry lives inside the same `select!` arm so abort stays responsive.
async fn chat_completions_streaming( async fn chat_completions_streaming(
&self, &self,
input: &Input, input: &Input,
@@ -97,7 +126,22 @@ pub trait Client: Sync + Send {
} }
let client = self.build_client()?; let client = self.build_client()?;
let data = input.prepare_completion_data(self.model(), true)?; let data = input.prepare_completion_data(self.model(), true)?;
self.chat_completions_streaming_inner(&client, handler, data).await let err = match self.chat_completions_streaming_inner(&client, handler, data).await {
Ok(()) => return Ok(()),
Err(err) => err,
};
if handler.has_received_content() || !should_retry_auth(&err, self.name()) {
return Err(err);
}
debug!(
"provider '{}' rejected access token (401); refreshing and retrying once",
self.name()
);
let data = input.prepare_completion_data(self.model(), true)?;
match self.chat_completions_streaming_inner(&client, handler, data).await {
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
ret => ret,
}
} => { } => {
handler.done(); handler.done();
ret.with_context(|| "Failed to call chat-completions api") ret.with_context(|| "Failed to call chat-completions api")
@@ -109,11 +153,27 @@ pub trait Client: Sync + Send {
} }
} }
/// Same retry-once-on-401 semantics as [`Self::chat_completions`]
/// (gemini OAuth embeddings route here).
async fn embeddings(&self, data: &EmbeddingsData) -> Result<Vec<Vec<f32>>> { async fn embeddings(&self, data: &EmbeddingsData) -> Result<Vec<Vec<f32>>> {
let client = self.build_client()?; let client = self.build_client()?;
self.embeddings_inner(&client, data) let err = match self.embeddings_inner(&client, data).await {
.await Ok(output) => return Ok(output),
.context("Failed to call embeddings api") Err(err) => err,
};
let ret = if should_retry_auth(&err, self.name()) {
debug!(
"provider '{}' rejected access token (401); refreshing and retrying once",
self.name()
);
match self.embeddings_inner(&client, data).await {
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
ret => ret,
}
} else {
Err(err)
};
ret.context("Failed to call embeddings api")
} }
async fn rerank(&self, data: &RerankData) -> Result<RerankOutput> { async fn rerank(&self, data: &RerankData) -> Result<RerankOutput> {
@@ -205,6 +265,7 @@ impl Default for ClientConfig {
pub struct ExtraConfig { pub struct ExtraConfig {
pub proxy: Option<String>, pub proxy: Option<String>,
pub connect_timeout: Option<u64>, pub connect_timeout: Option<u64>,
pub read_timeout: Option<u64>,
} }
#[derive(Debug, Clone, Deserialize, Default)] #[derive(Debug, Clone, Deserialize, Default)]
@@ -557,46 +618,90 @@ pub async fn noop_rerank(_builder: RequestBuilder, _model: &Model) -> Result<Rer
bail!("The client doesn't support rerank api") bail!("The client doesn't support rerank api")
} }
#[derive(Debug)]
pub struct ApiStatusError {
pub status: u16,
pub message: String,
}
impl std::fmt::Display for ApiStatusError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.message)
}
}
impl std::error::Error for ApiStatusError {}
/// True when the error chain bottoms out in an [`ApiStatusError`] with
/// status 401 EXACTLY. 403 (entitlement) and 429 (rate limit) are never
/// auth failures, and message text is never inspected.
fn is_auth_error(err: &anyhow::Error) -> bool {
err.downcast_ref::<ApiStatusError>()
.is_some_and(|api_err| api_err.status == 401)
}
/// Decides whether a 401 from `client_name` warrants a single retry after a
/// forced token refresh: the error must be a 401 [`ApiStatusError`], and the
/// client must have a cached access token to distrust (API-key clients have
/// none and never retry). Distrusting marks the exact rejected token so the
/// retry's prepare step force-refreshes it. There is deliberately no backoff:
/// the blast radius is bounded at one extra request per user-visible call.
///
/// Note: vertexai shares the ACCESS_TOKENS cache, so a 401 there also
/// triggers distrust+retry — deliberate.
fn should_retry_auth(err: &anyhow::Error, client_name: &str) -> bool {
if !is_auth_error(err) {
return false;
}
let Ok(token) = get_access_token(client_name) else {
return false;
};
distrust_access_token(client_name, &token)
}
pub fn catch_error(data: &Value, status: u16) -> Result<()> { pub fn catch_error(data: &Value, status: u16) -> Result<()> {
if (200..300).contains(&status) { if (200..300).contains(&status) {
return Ok(()); return Ok(());
} }
debug!("Invalid response, status: {status}, data: {data}"); debug!("Invalid response, status: {status}, data: {data}");
let api_error = |message: String| anyhow::Error::new(ApiStatusError { status, message });
if let Some(error) = data["error"].as_object() { if let Some(error) = data["error"].as_object() {
if let (Some(typ), Some(message)) = ( if let (Some(typ), Some(message)) = (
json_str_from_map(error, "type"), json_str_from_map(error, "type"),
json_str_from_map(error, "message"), json_str_from_map(error, "message"),
) { ) {
bail!("{message} (type: {typ})"); return Err(api_error(format!("{message} (type: {typ})")));
} else if let (Some(typ), Some(message)) = ( } else if let (Some(typ), Some(message)) = (
json_str_from_map(error, "code"), json_str_from_map(error, "code"),
json_str_from_map(error, "message"), json_str_from_map(error, "message"),
) { ) {
bail!("{message} (code: {typ})"); return Err(api_error(format!("{message} (code: {typ})")));
} }
} else if let Some(error) = data["errors"][0].as_object() { } else if let Some(error) = data["errors"][0].as_object() {
if let (Some(code), Some(message)) = ( if let (Some(code), Some(message)) = (
error.get("code").and_then(|v| v.as_u64()), error.get("code").and_then(|v| v.as_u64()),
json_str_from_map(error, "message"), json_str_from_map(error, "message"),
) { ) {
bail!("{message} (status: {code})") return Err(api_error(format!("{message} (status: {code})")));
} }
} else if let Some(error) = data[0]["error"].as_object() { } else if let Some(error) = data[0]["error"].as_object() {
if let (Some(status), Some(message)) = ( if let (Some(status), Some(message)) = (
json_str_from_map(error, "status"), json_str_from_map(error, "status"),
json_str_from_map(error, "message"), json_str_from_map(error, "message"),
) { ) {
bail!("{message} (status: {status})") return Err(api_error(format!("{message} (status: {status})")));
} }
} else if let (Some(detail), Some(status)) = (data["detail"].as_str(), data["status"].as_i64()) } else if let (Some(detail), Some(status)) = (data["detail"].as_str(), data["status"].as_i64())
{ {
bail!("{detail} (status: {status})"); return Err(api_error(format!("{detail} (status: {status})")));
} else if let Some(error) = data["error"].as_str() { } else if let Some(error) = data["error"].as_str() {
bail!("{error}"); return Err(api_error(error.to_string()));
} else if let Some(message) = data["message"].as_str() { } else if let Some(message) = data["message"].as_str() {
bail!("{message}"); return Err(api_error(message.to_string()));
} }
bail!("Invalid response data: {data} (status: {status})"); Err(api_error(format!(
"Invalid response data: {data} (status: {status})"
)))
} }
pub fn json_str_from_map<'a>( pub fn json_str_from_map<'a>(
@@ -737,3 +842,172 @@ fn prompt_input_string(desc: &str, required: bool, help_message: Option<&str>) -
let text = text.prompt()?; let text = text.prompt()?;
Ok(text) Ok(text)
} }
#[cfg(test)]
mod tests {
use super::*;
use super::super::access_token::{is_rejected, set_access_token};
fn catch_error_message(data: &Value, status: u16) -> String {
catch_error(data, status).unwrap_err().to_string()
}
#[test]
fn test_catch_error_display_json_with_type() {
let data = json!({"error": {"type": "invalid_request_error", "message": "Bad request"}});
assert_eq!(
catch_error_message(&data, 400),
"Bad request (type: invalid_request_error)"
);
}
#[test]
fn test_catch_error_display_json_with_code() {
let data = json!({"error": {"code": "rate_limited", "message": "Too many requests"}});
assert_eq!(
catch_error_message(&data, 429),
"Too many requests (code: rate_limited)"
);
}
#[test]
fn test_catch_error_display_errors_array() {
let data = json!({"errors": [{"code": 7000, "message": "No route"}]});
assert_eq!(catch_error_message(&data, 404), "No route (status: 7000)");
}
#[test]
fn test_catch_error_display_array_error_status() {
let data = json!([{"error": {"status": "PERMISSION_DENIED", "message": "Denied"}}]);
assert_eq!(
catch_error_message(&data, 403),
"Denied (status: PERMISSION_DENIED)"
);
}
#[test]
fn test_catch_error_display_detail_status() {
let data = json!({"detail": "Not found", "status": 404});
assert_eq!(catch_error_message(&data, 404), "Not found (status: 404)");
}
#[test]
fn test_catch_error_display_error_string() {
let data = json!({"error": "Something went wrong"});
assert_eq!(catch_error_message(&data, 500), "Something went wrong");
}
#[test]
fn test_catch_error_display_message_string() {
let data = json!({"message": "Unauthorized"});
assert_eq!(catch_error_message(&data, 401), "Unauthorized");
}
#[test]
fn test_catch_error_display_fallback() {
let data = json!({"unexpected": true});
assert_eq!(
catch_error_message(&data, 500),
format!("Invalid response data: {data} (status: 500)")
);
}
#[test]
fn test_catch_error_ok_on_success_status() {
let data = json!({"error": {"type": "x", "message": "y"}});
assert!(catch_error(&data, 200).is_ok());
assert!(catch_error(&data, 299).is_ok());
}
#[test]
fn test_catch_error_downcast_through_context_chain() {
let data = json!({"error": {"type": "authentication_error", "message": "Invalid key"}});
let err = catch_error(&data, 401)
.context("Failed to call chat-completions api")
.unwrap_err();
let api_err = err
.downcast_ref::<ApiStatusError>()
.expect("should downcast through context chain");
assert_eq!(api_err.status, 401);
assert_eq!(api_err.message, "Invalid key (type: authentication_error)");
}
#[test]
fn test_catch_error_preserves_status() {
let data = json!({"message": "Unauthorized"});
let err = catch_error(&data, 401).unwrap_err();
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 401);
let data = json!({"detail": "Rate limited", "status": 429});
let err = catch_error(&data, 429).unwrap_err();
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 429);
// The struct carries the outer HTTP status even when the body embeds another code
let data = json!({"errors": [{"code": 7000, "message": "No route"}]});
let err = catch_error(&data, 429).unwrap_err();
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 429);
}
/// Wrapped in `.context(...)` so every test below proves the downcast
/// works through an anyhow context chain, as in the trait methods.
fn api_status_error(status: u16) -> anyhow::Error {
anyhow::Error::new(ApiStatusError {
status,
message: format!("error (status: {status})"),
})
.context("Failed to call chat-completions api")
}
fn cache_token(client: &str, token: &str) {
set_access_token(
client,
token.into(),
chrono::Utc::now().timestamp() + 3600,
None,
);
}
#[test]
fn test_should_retry_auth_401_with_cached_token() {
let client = "should-retry-auth-401";
cache_token(client, "at-1");
assert!(should_retry_auth(&api_status_error(401), client));
assert!(is_rejected(client, "at-1"), "rejected marker not set");
}
#[test]
fn test_should_retry_auth_non_401_statuses() {
let client = "should-retry-auth-non-401";
cache_token(client, "at-1");
for status in [403, 429, 500] {
assert!(
!should_retry_auth(&api_status_error(status), client),
"retried on {status}"
);
}
assert_eq!(get_access_token(client).unwrap(), "at-1");
assert!(!is_rejected(client, "at-1"), "marker set without a 401");
}
#[test]
fn test_should_retry_auth_non_api_status_error() {
let client = "should-retry-auth-non-api";
cache_token(client, "at-1");
let err = anyhow::anyhow!("connection reset").context("Failed to call embeddings api");
assert!(!should_retry_auth(&err, client));
assert_eq!(get_access_token(client).unwrap(), "at-1");
assert!(!is_rejected(client, "at-1"));
}
#[test]
fn test_should_retry_auth_401_without_cached_token() {
let client = "should-retry-auth-no-token";
assert!(!should_retry_auth(&api_status_error(401), client));
assert!(!is_rejected(client, "at-1"));
}
}
+410 -43
View File
@@ -1,14 +1,14 @@
use super::access_token::{is_valid_access_token, set_access_token}; use super::access_token::{clear_rejected, is_rejected, is_valid_access_token, set_access_token};
use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider; use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider;
use super::{ClientConfig, ProviderModels}; use super::{ClientConfig, ProviderModels};
use crate::config::paths; use crate::config::paths;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Error, Result, anyhow, bail};
use base64::Engine; use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD; use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use chrono::Utc; use chrono::Utc;
use indexmap::IndexMap; use indexmap::IndexMap;
use inquire::Text; use inquire::Text;
use reqwest::{Client as ReqwestClient, RequestBuilder}; use reqwest::{Client as ReqwestClient, RequestBuilder, StatusCode};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
@@ -16,6 +16,10 @@ use std::collections::HashMap;
use std::fs; use std::fs;
use std::io::{BufRead, BufReader, Write}; use std::io::{BufRead, BufReader, Write};
use std::net::TcpListener; use std::net::TcpListener;
use std::path::PathBuf;
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tokio::sync;
use url::Url; use url::Url;
use uuid::Uuid; use uuid::Uuid;
@@ -197,10 +201,20 @@ pub struct OAuthTokens {
pub account_id: Option<String>, pub account_id: Option<String>,
} }
const TOKEN_ENDPOINT_TIMEOUT: Duration = Duration::from_secs(30);
pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> { pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> {
match provider.flow() { match provider.flow() {
OAuthFlow::Pkce => run_pkce_flow(provider, client_name).await, OAuthFlow::Pkce => run_pkce_flow(provider, client_name).await,
OAuthFlow::ClientCredentials => run_client_credentials_flow(provider, client_name).await, OAuthFlow::ClientCredentials => {
run_client_credentials_flow(provider, client_name).await?;
println!(
"Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.",
client_name,
provider.provider_name()
);
Ok(())
}
OAuthFlow::DeviceCode => run_device_code_flow(provider, client_name).await, OAuthFlow::DeviceCode => run_device_code_flow(provider, client_name).await,
} }
} }
@@ -301,12 +315,20 @@ async fn run_pkce_flow(provider: &dyn OAuthProvider, client_name: &str) -> Resul
let access_token = response["access_token"] let access_token = response["access_token"]
.as_str() .as_str()
.ok_or_else(|| anyhow!("Missing access_token in response: {response}"))? .ok_or_else(|| {
anyhow!(
"Missing access_token in response (keys: {})",
token_response_keys(&response)
)
})?
.to_string(); .to_string();
let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string()); let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string());
let expires_in = response["expires_in"] let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
.as_i64() anyhow!(
.ok_or_else(|| anyhow!("Missing expires_in in response: {response}"))?; "Missing expires_in in response (keys: {})",
token_response_keys(&response)
)
})?;
let expires_at = Utc::now().timestamp() + expires_in; let expires_at = Utc::now().timestamp() + expires_in;
@@ -334,7 +356,9 @@ async fn run_client_credentials_flow(
provider: &dyn OAuthProvider, provider: &dyn OAuthProvider,
client_name: &str, client_name: &str,
) -> Result<()> { ) -> Result<()> {
let client = ReqwestClient::new(); let client = ReqwestClient::builder()
.timeout(TOKEN_ENDPOINT_TIMEOUT)
.build()?;
let scopes = provider.scopes(); let scopes = provider.scopes();
let mut params: Vec<(&str, &str)> = vec![ let mut params: Vec<(&str, &str)> = vec![
("grant_type", "client_credentials"), ("grant_type", "client_credentials"),
@@ -349,11 +373,19 @@ async fn run_client_credentials_flow(
let access_token = response["access_token"] let access_token = response["access_token"]
.as_str() .as_str()
.ok_or_else(|| anyhow!("Missing access_token in client_credentials response: {response}"))? .ok_or_else(|| {
anyhow!(
"Missing access_token in client_credentials response (keys: {})",
token_response_keys(&response)
)
})?
.to_string(); .to_string();
let expires_in = response["expires_in"] let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
.as_i64() anyhow!(
.ok_or_else(|| anyhow!("Missing expires_in in client_credentials response: {response}"))?; "Missing expires_in in client_credentials response (keys: {})",
token_response_keys(&response)
)
})?;
let expires_at = Utc::now().timestamp() + expires_in; let expires_at = Utc::now().timestamp() + expires_in;
let tokens = OAuthTokens { let tokens = OAuthTokens {
@@ -363,11 +395,6 @@ async fn run_client_credentials_flow(
account_id: provider.extract_account_id(&response), account_id: provider.extract_account_id(&response),
}; };
save_oauth_tokens(client_name, &tokens)?; save_oauth_tokens(client_name, &tokens)?;
println!(
"Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.",
client_name,
provider.provider_name()
);
Ok(()) Ok(())
} }
@@ -417,19 +444,28 @@ async fn run_device_code_flow(provider: &dyn OAuthProvider, client_name: &str) -
let device_code = device_response["device_code"] let device_code = device_response["device_code"]
.as_str() .as_str()
.ok_or_else(|| { .ok_or_else(|| {
anyhow!("Missing device_code in device authorization response: {device_response}") anyhow!(
"Missing device_code in device authorization response (keys: {})",
token_response_keys(&device_response)
)
})? })?
.to_string(); .to_string();
let user_code = device_response["user_code"] let user_code = device_response["user_code"]
.as_str() .as_str()
.ok_or_else(|| { .ok_or_else(|| {
anyhow!("Missing user_code in device authorization response: {device_response}") anyhow!(
"Missing user_code in device authorization response (keys: {})",
token_response_keys(&device_response)
)
})? })?
.to_string(); .to_string();
let verification_uri = device_response["verification_uri"] let verification_uri = device_response["verification_uri"]
.as_str() .as_str()
.ok_or_else(|| { .ok_or_else(|| {
anyhow!("Missing verification_uri in device authorization response: {device_response}") anyhow!(
"Missing verification_uri in device authorization response (keys: {})",
token_response_keys(&device_response)
)
})? })?
.to_string(); .to_string();
let verification_uri_complete = device_response["verification_uri_complete"] let verification_uri_complete = device_response["verification_uri_complete"]
@@ -551,10 +587,76 @@ fn save_oauth_tokens(client_name: &str, tokens: &OAuthTokens) -> Result<()> {
fs::create_dir_all(parent)?; fs::create_dir_all(parent)?;
} }
let json = serde_json::to_string_pretty(tokens)?; let json = serde_json::to_string_pretty(tokens)?;
fs::write(path, json)?; // Write-then-rename so a crash mid-write never truncates the live token file.
let mut tmp = path.clone().into_os_string();
tmp.push(".tmp");
let tmp = PathBuf::from(tmp);
// Tokens are live credentials: create the file owner-only, not umask-default.
let mut options = fs::OpenOptions::new();
options.write(true).create(true).truncate(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options.mode(0o600);
}
options.open(&tmp)?.write_all(json.as_bytes())?;
fs::rename(&tmp, &path)?;
Ok(()) Ok(())
} }
pub(crate) fn token_response_keys(response: &Value) -> String {
match response.as_object() {
Some(map) => {
let keys: Vec<&str> = map.keys().map(String::as_str).collect();
format!("[{}]", keys.join(", "))
}
None => "<non-object response>".to_string(),
}
}
fn parse_refresh_response(
status: StatusCode,
response: &Value,
previous_refresh_token: Option<&str>,
) -> Result<(String, Option<String>, i64)> {
if let Some(error) = response["error"].as_str() {
let description = response["error_description"]
.as_str()
.unwrap_or("no description");
if matches!(error, "invalid_grant" | "invalid_token") {
bail!(
"OAuth refresh token was rejected ({error}: {description}). Please re-authenticate."
);
}
bail!("Token refresh failed ({error}: {description})");
}
if !status.is_success() {
bail!("Token refresh failed with HTTP status {status}");
}
let access_token = response["access_token"]
.as_str()
.ok_or_else(|| {
anyhow!(
"Missing access_token in refresh response (keys: {})",
token_response_keys(response)
)
})?
.to_string();
let refresh_token = response["refresh_token"]
.as_str()
.or(previous_refresh_token)
.map(str::to_string);
let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
anyhow!(
"Missing expires_in in refresh response (keys: {})",
token_response_keys(response)
)
})?;
Ok((access_token, refresh_token, expires_in))
}
pub async fn refresh_oauth_token( pub async fn refresh_oauth_token(
client: &ReqwestClient, client: &ReqwestClient,
provider: &dyn OAuthProvider, provider: &dyn OAuthProvider,
@@ -577,19 +679,23 @@ pub async fn refresh_oauth_token(
], ],
); );
let response: Value = request.send().await?.json().await?; let (status, response) = tokio::time::timeout(TOKEN_ENDPOINT_TIMEOUT, async {
let response = request.send().await?;
let status = response.status();
let body: Value = response.json().await?;
Ok::<_, Error>((status, body))
})
.await
.map_err(|_| {
anyhow!(
"Token refresh for '{}' timed out after {}s",
client_name,
TOKEN_ENDPOINT_TIMEOUT.as_secs()
)
})??;
let access_token = response["access_token"] let (access_token, refresh_token, expires_in) =
.as_str() parse_refresh_response(status, &response, tokens.refresh_token.as_deref())?;
.ok_or_else(|| anyhow!("Missing access_token in refresh response: {response}"))?
.to_string();
let refresh_token = response["refresh_token"]
.as_str()
.map(|s| s.to_string())
.or_else(|| tokens.refresh_token.clone());
let expires_in = response["expires_in"]
.as_i64()
.ok_or_else(|| anyhow!("Missing expires_in in refresh response: {response}"))?;
let expires_at = Utc::now().timestamp() + expires_in; let expires_at = Utc::now().timestamp() + expires_in;
@@ -609,6 +715,20 @@ pub async fn refresh_oauth_token(
Ok(new_tokens) Ok(new_tokens)
} }
/// Per-client lock so concurrent requests perform a single refresh.
/// Returns a clone of the Arc so the parking_lot guard is dropped before the
/// caller awaits on the tokio mutex.
fn refresh_guard(client_name: &str) -> Arc<sync::Mutex<()>> {
static GUARDS: OnceLock<parking_lot::Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
OnceLock::new();
GUARDS
.get_or_init(Default::default)
.lock()
.entry(client_name.to_string())
.or_default()
.clone()
}
pub async fn prepare_oauth_access_token( pub async fn prepare_oauth_access_token(
client: &ReqwestClient, client: &ReqwestClient,
provider: &dyn OAuthProvider, provider: &dyn OAuthProvider,
@@ -623,16 +743,40 @@ pub async fn prepare_oauth_access_token(
None => return Ok(false), None => return Ok(false),
}; };
let tokens = if Utc::now().timestamp() >= tokens.expires_at { let tokens = if Utc::now().timestamp() >= tokens.expires_at
match provider.flow() { || is_rejected(client_name, &tokens.access_token)
OAuthFlow::Pkce | OAuthFlow::DeviceCode => { {
refresh_oauth_token(client, provider, client_name, &tokens).await? let guard = refresh_guard(client_name);
} let _guard = guard.lock().await;
OAuthFlow::ClientCredentials => {
run_client_credentials_flow(provider, client_name).await?; // A concurrent caller may have refreshed while we waited for the
load_oauth_tokens(client_name) // lock; a valid in-memory token means the winner already populated
.ok_or_else(|| anyhow!("Token file missing after client_credentials refresh"))? // the cache.
if is_valid_access_token(client_name) {
return Ok(true);
}
let tokens = match load_oauth_tokens(client_name) {
Some(t) => t,
None => return Ok(false),
};
if Utc::now().timestamp() >= tokens.expires_at
|| is_rejected(client_name, &tokens.access_token)
{
match provider.flow() {
OAuthFlow::Pkce | OAuthFlow::DeviceCode => {
refresh_oauth_token(client, provider, client_name, &tokens).await?
}
OAuthFlow::ClientCredentials => {
run_client_credentials_flow(provider, client_name).await?;
load_oauth_tokens(client_name).ok_or_else(|| {
anyhow!("Token file missing after client_credentials refresh")
})?
}
} }
} else {
tokens
} }
} else { } else {
tokens tokens
@@ -644,6 +788,9 @@ pub async fn prepare_oauth_access_token(
tokens.expires_at, tokens.expires_at,
tokens.account_id, tokens.account_id,
); );
// Clear even when the refresh returned the same token (some IdPs reuse
// JWTs within validity); otherwise every request re-hits the token endpoint.
clear_rejected(client_name);
Ok(true) Ok(true)
} }
@@ -886,11 +1033,55 @@ pub(crate) fn client_config_info(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::ffi::OsString;
use std::path::PathBuf;
use std::str; use std::str;
use std::time::UNIX_EPOCH;
use super::*; use super::*;
use crate::client::access_token::{distrust_access_token, get_access_token};
use crate::client::openai_compatible::OpenAICompatibleConfig; use crate::client::openai_compatible::OpenAICompatibleConfig;
use crate::client::{ModelData, ProviderModels}; use crate::client::{ModelData, ProviderModels};
use crate::utils::get_env_name;
use serial_test::serial;
use std::{env, time::SystemTime};
fn with_temp_cache<F: FnOnce()>(f: F) {
struct Restore {
key: String,
prev: Option<OsString>,
root: PathBuf,
}
impl Drop for Restore {
fn drop(&mut self) {
unsafe {
match self.prev.take() {
Some(v) => env::set_var(&self.key, v),
None => env::remove_var(&self.key),
}
}
let _ = fs::remove_dir_all(&self.root);
}
}
let unique = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
let root = env::temp_dir().join(format!("coyote-client-oauth-test-{unique}"));
fs::create_dir_all(&root).unwrap();
let env_key = get_env_name("cache_dir");
let prev = env::var_os(&env_key);
unsafe {
env::set_var(&env_key, &root);
}
let _restore = Restore {
key: env_key,
prev,
root,
};
f();
}
fn base_config() -> OAuthConfig { fn base_config() -> OAuthConfig {
OAuthConfig { OAuthConfig {
@@ -1468,4 +1659,180 @@ scopes:
"body missing grant_type param: {body}" "body missing grant_type param: {body}"
); );
} }
#[test]
#[serial]
fn save_oauth_tokens_roundtrips_and_leaves_no_tmp_file() {
with_temp_cache(|| {
let tokens = OAuthTokens {
access_token: "at-123".into(),
refresh_token: Some("rt-456".into()),
expires_at: 1234567890,
account_id: Some("acct-789".into()),
};
save_oauth_tokens("atomic-test", &tokens).unwrap();
let loaded = load_oauth_tokens("atomic-test").unwrap();
assert_eq!(loaded.access_token, "at-123");
assert_eq!(loaded.refresh_token.as_deref(), Some("rt-456"));
assert_eq!(loaded.expires_at, 1234567890);
assert_eq!(loaded.account_id.as_deref(), Some("acct-789"));
let dir = paths::oauth_tokens_dir();
let leftover_tmp = fs::read_dir(&dir)
.unwrap()
.any(|e| e.unwrap().file_name().to_string_lossy().ends_with(".tmp"));
assert!(!leftover_tmp, "temp file left behind in {dir:?}");
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = fs::metadata(paths::token_file("atomic-test"))
.unwrap()
.permissions()
.mode();
assert_eq!(mode & 0o777, 0o600, "token file mode was {mode:o}");
}
});
}
#[test]
#[serial]
fn prepare_rejected_valid_file_token_attempts_refresh_branch() {
with_temp_cache(|| {
let client_name = "prepare-rejected-branch-test";
let expires_at = Utc::now().timestamp() + 3600;
save_oauth_tokens(
client_name,
&OAuthTokens {
access_token: "rejected-at".into(),
refresh_token: None,
expires_at,
account_id: None,
},
)
.unwrap();
set_access_token(client_name, "rejected-at".into(), expires_at, None);
assert!(distrust_access_token(client_name, "rejected-at"));
let err = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(prepare_oauth_access_token(
&ReqwestClient::new(),
&ResourceStubProvider,
client_name,
))
.unwrap_err()
.to_string();
// The timestamp-valid but rejected file token must not be trusted;
// the refresh branch is taken and bails on the missing refresh token.
assert!(err.contains("No refresh token"), "unexpected error: {err}");
});
}
#[test]
#[serial]
fn prepare_trusts_differing_unmarked_valid_file_token() {
with_temp_cache(|| {
let client_name = "prepare-differing-token-test";
let expires_at = Utc::now().timestamp() + 3600;
set_access_token(client_name, "rejected-at".into(), expires_at, None);
assert!(distrust_access_token(client_name, "rejected-at"));
save_oauth_tokens(
client_name,
&OAuthTokens {
access_token: "fresh-at".into(),
refresh_token: None,
expires_at,
account_id: None,
},
)
.unwrap();
let ready = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(prepare_oauth_access_token(
&ReqwestClient::new(),
&ResourceStubProvider,
client_name,
))
.unwrap();
assert!(ready);
assert_eq!(get_access_token(client_name).unwrap(), "fresh-at");
assert!(
!is_rejected(client_name, "rejected-at"),
"marker not cleared after successful prepare"
);
});
}
#[test]
fn parse_refresh_response_invalid_grant_redacts_and_prompts_reauth() {
let response = serde_json::json!({
"error": "invalid_grant",
"error_description": "refresh token revoked",
"refresh_token": "planted-secret-token",
});
let err = parse_refresh_response(StatusCode::BAD_REQUEST, &response, Some("old-rt"))
.unwrap_err()
.to_string();
assert!(err.contains("re-authenticate"), "unexpected error: {err}");
assert!(
!err.contains("planted-secret-token"),
"error leaked token material: {err}"
);
}
#[test]
fn token_response_keys_lists_keys_without_values() {
let response = serde_json::json!({
"access_token": "secret-at",
"token_type": "SecretBearer",
});
let keys = token_response_keys(&response);
assert!(keys.contains("access_token"), "missing key name: {keys}");
assert!(keys.contains("token_type"), "missing key name: {keys}");
assert!(!keys.contains("secret-at"), "leaked value: {keys}");
assert!(!keys.contains("SecretBearer"), "leaked value: {keys}");
}
#[test]
fn parse_refresh_response_rotates_refresh_token_when_present() {
let response = serde_json::json!({
"access_token": "new-at",
"refresh_token": "new-rt",
"expires_in": 3600,
});
let (access_token, refresh_token, expires_in) =
parse_refresh_response(StatusCode::OK, &response, Some("old-rt")).unwrap();
assert_eq!(access_token, "new-at");
assert_eq!(refresh_token.as_deref(), Some("new-rt"));
assert_eq!(expires_in, 3600);
}
#[test]
fn parse_refresh_response_keeps_old_refresh_token_when_absent() {
let response = serde_json::json!({
"access_token": "new-at",
"expires_in": 3600,
});
let (_, refresh_token, _) =
parse_refresh_response(StatusCode::OK, &response, Some("old-rt")).unwrap();
assert_eq!(refresh_token.as_deref(), Some("old-rt"));
}
} }
+54 -5
View File
@@ -1,4 +1,4 @@
use super::{ThinkingBlock, ToolCall, catch_error}; use super::{ApiStatusError, ThinkingBlock, ToolCall, catch_error};
use crate::utils::AbortSignal; use crate::utils::AbortSignal;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Result, anyhow, bail};
@@ -176,6 +176,14 @@ impl SseHandler {
self.thinking.push(block); self.thinking.push(block);
} }
/// Whether any output (text, tool calls, or thinking blocks) has been
/// accumulated. `Client::chat_completions_streaming` gates its 401 retry
/// on this: content already streamed to the user would be rendered a
/// second time by a retry, so partial responses are never retried.
pub fn has_received_content(&self) -> bool {
!self.buffer.is_empty() || !self.tool_calls.is_empty() || !self.thinking.is_empty()
}
pub fn abort(&self) -> AbortSignal { pub fn abort(&self) -> AbortSignal {
self.abort_signal.clone() self.abort_signal.clone()
} }
@@ -224,10 +232,14 @@ where
let data: Value = match text.parse() { let data: Value = match text.parse() {
Ok(data) => data, Ok(data) => data,
Err(_) => { Err(_) => {
bail!( return Err(ApiStatusError {
"Invalid response data: {text} (status: {})", status: status.as_u16(),
status.as_u16() message: format!(
); "Invalid response data: {text} (status: {})",
status.as_u16()
),
}
.into());
} }
}; };
catch_error(&data, status.as_u16())?; catch_error(&data, status.as_u16())?;
@@ -418,6 +430,43 @@ mod tests {
assert!(error_message.contains("test_function_loop")); assert!(error_message.contains("test_function_loop"));
} }
fn new_handler() -> (SseHandler, tokio::sync::mpsc::UnboundedReceiver<SseEvent>) {
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
let abort_signal = crate::utils::create_abort_signal();
(SseHandler::new(sender, abort_signal), receiver)
}
#[test]
fn test_has_received_content_text() {
let (mut handler, _rx) = new_handler();
assert!(!handler.has_received_content());
handler.text("hello").unwrap();
assert!(handler.has_received_content());
}
#[test]
fn test_has_received_content_tool_call() {
let (mut handler, _rx) = new_handler();
assert!(!handler.has_received_content());
let call = ToolCall::new("test_function".to_string(), json!({"param": 1}), None);
handler.tool_call(call).unwrap();
assert!(handler.has_received_content());
}
#[test]
fn test_has_received_content_thinking() {
let (mut handler, _rx) = new_handler();
assert!(!handler.has_received_content());
handler.thinking_block(ThinkingBlock::Thinking {
thinking: "hmm".to_string(),
signature: "sig".to_string(),
});
assert!(handler.has_received_content());
}
fn split_chunks(text: &str) -> Vec<Vec<u8>> { fn split_chunks(text: &str) -> Vec<Vec<u8>> {
let len = text.len(); let len = text.len();
let cut1 = random_range(1..len - 1); let cut1 = random_range(1..len - 1);
+12 -3
View File
@@ -4,6 +4,7 @@ use crate::{
client::Model, client::Model,
config::memory, config::memory,
function::{Functions, run_llm_function}, function::{Functions, run_llm_function},
graph, rag,
}; };
use super::rag_cache::RagKey; use super::rag_cache::RagKey;
@@ -185,7 +186,7 @@ impl Agent {
&rag_path_clone, &rag_path_clone,
&document_paths, &document_paths,
abort, abort,
false, true,
) )
.await .await
}) })
@@ -247,6 +248,10 @@ impl Agent {
} }
} }
if rag.is_some() && app.function_calling_support && graph_for_rag.is_none() {
functions.append_rag_query_functions();
}
agent_config.replace_tools_placeholder(&functions); agent_config.replace_tools_placeholder(&functions);
Ok(Self { Ok(Self {
@@ -1021,11 +1026,11 @@ async fn init_graph_rags(
// Graph validation catches this too, but it is skipped when // Graph validation catches this too, but it is skipped when
// `validate_before_run` is off, so this guard is the load-bearing one. // `validate_before_run` is off, so this guard is the load-bearing one.
if let Some(driver) = &rag_node.driver if let Some(driver) = &rag_node.driver
&& let Some(message) = crate::graph::validator::rag_driver_error(driver) && let Some(message) = graph::validator::rag_driver_error(driver)
{ {
bail!("rag node '{node_id}': {message}"); bail!("rag node '{node_id}': {message}");
} }
let config = rag_init_config(rag_node); let mut config = rag_init_config(rag_node);
let fully_specified = config.embedding_model.is_some() let fully_specified = config.embedding_model.is_some()
&& config.chunk_size.is_some() && config.chunk_size.is_some()
&& config.chunk_overlap.is_some(); && config.chunk_overlap.is_some();
@@ -1051,6 +1056,10 @@ async fn init_graph_rags(
initialized. RAG initialization is required for this agent." initialized. RAG initialization is required for this agent."
); );
} }
if config.driver.is_none() {
config.driver = Some(rag::select_rag_driver()?);
}
} }
let document_paths = let document_paths =
+13 -9
View File
@@ -1,6 +1,6 @@
use crate::client::{ClientConfig, Model, ModelType, list_models}; use crate::client::{ClientConfig, Model, ModelType, list_models};
use crate::render::{MarkdownRender, RenderOptions}; use crate::render::{MarkdownRender, RenderOptions};
use crate::utils::{IS_STDOUT_TERMINAL, NO_COLOR, decode_bin, get_env_name}; use crate::utils::{IS_STDOUT_TERMINAL, NO_COLOR, decode_bin, drain_stale_tty_input, get_env_name};
use super::paths; use super::paths;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Result, anyhow, bail};
@@ -616,14 +616,18 @@ impl AppConfig {
if self.highlight && self.theme.is_none() { if self.highlight && self.theme.is_none() {
if let Some(v) = super::read_env_value::<String>(&get_env_name("theme")) { if let Some(v) = super::read_env_value::<String>(&get_env_name("theme")) {
self.theme = v; self.theme = v;
} else if *IS_STDOUT_TERMINAL } else if *IS_STDOUT_TERMINAL {
&& let Ok(color_scheme) = color_scheme(QueryOptions::default()) if let Ok(color_scheme) = color_scheme(QueryOptions::default()) {
{ let theme = match color_scheme {
let theme = match color_scheme { ColorScheme::Dark => "dark",
ColorScheme::Dark => "dark", ColorScheme::Light => "light",
ColorScheme::Light => "light", };
}; self.theme = Some(theme.into());
self.theme = Some(theme.into()); }
// The OSC/DA1 reply can arrive after colorsaurus stops reading
// (observed under zellij-in-kitty). Drain any late reply bytes so
// they are neither echoed nor read as line-editor input.
drain_stale_tty_input();
} }
} }
if let Some(v) = super::read_env_value::<String>(&get_env_name("left_prompt")) { if let Some(v) = super::read_env_value::<String>(&get_env_name("left_prompt")) {
+13 -19
View File
@@ -1,6 +1,6 @@
use crate::mcp::{ use crate::mcp::{
ConnectedServer, JsonField, McpServer, McpTransportType, is_auth_required_error, oauth, ConnectedServer, JsonField, McpAuthRequired, McpServer, McpTransportType,
spawn_mcp_server, is_auth_required_error, resolve_http_auth, spawn_mcp_server,
}; };
use anyhow::Result; use anyhow::Result;
@@ -102,23 +102,17 @@ impl McpFactory {
return Ok(existing); return Ok(existing);
} }
let bearer_token = if spec.is_remote() { let (auth, auth_reason) = resolve_http_auth(name, spec).await;
oauth::load_valid_mcp_token(name) let handle = spawn_mcp_server(spec, log_path, auth).await.map_err(|e| {
} else { if is_auth_required_error(&e) {
None e.context(McpAuthRequired {
}; server: name.to_string(),
let handle = spawn_mcp_server(spec, log_path, bearer_token) reason: auth_reason,
.await })
.map_err(|e| { } else {
if is_auth_required_error(&e) { e
e.context(format!( }
"MCP server '{name}' requires OAuth authentication. \ })?;
Run `coyote --auth-mcp {name}` or `.mcp auth {name}` in the REPL to authenticate."
))
} else {
e
}
})?;
self.insert_active(key, &handle); self.insert_active(key, &handle);
Ok(handle) Ok(handle)
} }
+18 -1
View File
@@ -148,6 +148,16 @@ pub fn sbx_kit_hash_file() -> PathBuf {
sbx_kit_dir().join(SBX_KIT_HASH_FILE) sbx_kit_dir().join(SBX_KIT_HASH_FILE)
} }
pub fn sandbox_mixin_hashes_dir() -> PathBuf {
cache_dir().join("sandbox-mixin-hashes")
}
pub fn sandbox_mixin_hash_file(sandbox_name: &str) -> PathBuf {
// Sandbox names are sanitized by the caller, but never trust a path
// component: a stray separator must not escape the hash directory.
sandbox_mixin_hashes_dir().join(format!("{}.hash", sandbox_name.replace('/', "_")))
}
pub fn sbx_mixin_kits_dir() -> PathBuf { pub fn sbx_mixin_kits_dir() -> PathBuf {
cache_dir().join(SBX_MIXIN_KITS_DIR_NAME) cache_dir().join(SBX_MIXIN_KITS_DIR_NAME)
} }
@@ -437,6 +447,10 @@ pub(crate) fn remove_rag_sidecars(dir: &Path, name: &str) -> Result<()> {
if duckdb_path.exists() { if duckdb_path.exists() {
let _ = remove_file(&duckdb_path); let _ = remove_file(&duckdb_path);
} }
let wal_path = dir.join(format!("{name}.duckdb.wal"));
if wal_path.exists() {
let _ = remove_file(&wal_path);
}
let mixin_path = dir.join(format!("{name}.sbx-mixin.yaml")); let mixin_path = dir.join(format!("{name}.sbx-mixin.yaml"));
if mixin_path.exists() { if mixin_path.exists() {
remove_file(&mixin_path).with_context(|| { remove_file(&mixin_path).with_context(|| {
@@ -894,16 +908,19 @@ mod tests {
} }
#[test] #[test]
fn remove_rag_sidecars_removes_both() { fn remove_rag_sidecars_removes_duckdb_wal_and_mixin() {
let root = sidecar_temp_dir("rag-sidecars-both"); let root = sidecar_temp_dir("rag-sidecars-both");
let duckdb = root.join("docs.duckdb"); let duckdb = root.join("docs.duckdb");
let wal = root.join("docs.duckdb.wal");
let mixin = root.join("docs.sbx-mixin.yaml"); let mixin = root.join("docs.sbx-mixin.yaml");
fs::write(&duckdb, "db").unwrap(); fs::write(&duckdb, "db").unwrap();
fs::write(&wal, "wal").unwrap();
fs::write(&mixin, "mixin").unwrap(); fs::write(&mixin, "mixin").unwrap();
remove_rag_sidecars(&root, "docs").unwrap(); remove_rag_sidecars(&root, "docs").unwrap();
assert!(!duckdb.exists(), "the .duckdb sidecar must be removed"); assert!(!duckdb.exists(), "the .duckdb sidecar must be removed");
assert!(!wal.exists(), "the .duckdb.wal sidecar must be removed");
assert!( assert!(
!mixin.exists(), !mixin.exists(),
"the .sbx-mixin.yaml sidecar must be removed" "the .sbx-mixin.yaml sidecar must be removed"
+26 -12
View File
@@ -16,12 +16,13 @@ use super::{MessageContentToolCalls, prompts};
use crate::client::{Model, ModelType, list_models}; use crate::client::{Model, ModelType, list_models};
use crate::function::{ use crate::function::{
FunctionDeclaration, Functions, ToolCallTracker, ToolResult, memory::MEMORY_FUNCTION_PREFIX, FunctionDeclaration, Functions, ToolCallTracker, ToolResult, memory::MEMORY_FUNCTION_PREFIX,
skill::SKILL_FUNCTION_PREFIX, supervisor::SUPERVISOR_FUNCTION_PREFIX, rag_query::RAG_FUNCTION_PREFIX, skill::SKILL_FUNCTION_PREFIX,
todo::TODO_FUNCTION_PREFIX, user_interaction::USER_FUNCTION_PREFIX, supervisor::SUPERVISOR_FUNCTION_PREFIX, todo::TODO_FUNCTION_PREFIX,
user_interaction::USER_FUNCTION_PREFIX,
}; };
use crate::mcp::{ use crate::mcp::{
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
MCP_SEARCH_META_FUNCTION_NAME_PREFIX, is_auth_required_error, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error,
}; };
use crate::rag::Rag; use crate::rag::Rag;
use crate::supervisor::Supervisor; use crate::supervisor::Supervisor;
@@ -715,6 +716,7 @@ impl RequestContext {
pub fn exit_rag(&mut self) -> Result<()> { pub fn exit_rag(&mut self) -> Result<()> {
self.rag.take(); self.rag.take();
self.tool_scope.functions.remove_rag_query_functions();
Ok(()) Ok(())
} }
@@ -1137,6 +1139,7 @@ impl RequestContext {
&& !v.name.starts_with("agent__") && !v.name.starts_with("agent__")
&& !v.name.starts_with("memory__") && !v.name.starts_with("memory__")
&& !v.name.starts_with("skill__") && !v.name.starts_with("skill__")
&& !v.name.starts_with("rag__")
}) })
.map(|v| v.name.clone()) .map(|v| v.name.clone())
.collect() .collect()
@@ -1957,7 +1960,8 @@ impl RequestContext {
|| (!matches!(role.skills_enabled(), Some(false)) || (!matches!(role.skills_enabled(), Some(false))
&& v.name.starts_with(SKILL_FUNCTION_PREFIX)) && v.name.starts_with(SKILL_FUNCTION_PREFIX))
|| (self.auto_continue_config().enabled || (self.auto_continue_config().enabled
&& v.name.starts_with(TODO_FUNCTION_PREFIX))) && v.name.starts_with(TODO_FUNCTION_PREFIX))
|| v.name.starts_with(RAG_FUNCTION_PREFIX))
&& !existing.contains(&v.name) && !existing.contains(&v.name)
}) })
.cloned() .cloned()
@@ -1987,6 +1991,7 @@ impl RequestContext {
|| v.name.starts_with(TODO_FUNCTION_PREFIX) || v.name.starts_with(TODO_FUNCTION_PREFIX)
|| v.name.starts_with(SUPERVISOR_FUNCTION_PREFIX) || v.name.starts_with(SUPERVISOR_FUNCTION_PREFIX)
|| v.name.starts_with(MEMORY_FUNCTION_PREFIX) || v.name.starts_with(MEMORY_FUNCTION_PREFIX)
|| v.name.starts_with(RAG_FUNCTION_PREFIX)
}); });
} }
@@ -3422,7 +3427,11 @@ impl RequestContext {
{ {
Ok(handle) => handles.push((id.clone(), handle)), Ok(handle) => handles.push((id.clone(), handle)),
Err(e) if is_auth_required_error(&e) => { Err(e) if is_auth_required_error(&e) => {
auth_required.push(id.clone()) let reason = e
.downcast_ref::<McpAuthRequired>()
.map(|a| a.reason)
.unwrap_or(McpAuthReason::NotAuthenticated);
auth_required.push((id.clone(), reason));
} }
Err(e) => return Err(e), Err(e) => return Err(e),
} }
@@ -3439,11 +3448,8 @@ impl RequestContext {
for (id, handle) in handles { for (id, handle) in handles {
mcp_runtime.insert(id, handle); mcp_runtime.insert(id, handle);
} }
for id in auth_required { for (id, reason) in auth_required {
eprintln!( eprintln!("Warning: {}", McpAuthRequired { server: id, reason });
"Warning: MCP server '{id}' requires OAuth authentication and was not started. \
Run `.mcp auth {id}` (or `coyote --auth-mcp {id}`) to authenticate and attach it."
);
} }
} }
} }
@@ -3467,6 +3473,12 @@ impl RequestContext {
if self.should_register_memory_tools() { if self.should_register_memory_tools() {
functions.append_memory_functions(); functions.append_memory_functions();
} }
if self.rag.is_some()
&& app.function_calling_support
&& !self.agent.as_ref().is_some_and(|a| a.is_graph())
{
functions.append_rag_query_functions();
}
let tool_tracker = self.tool_scope.tool_tracker.clone(); let tool_tracker = self.tool_scope.tool_tracker.clone();
self.tool_scope = ToolScope { self.tool_scope = ToolScope {
@@ -4136,7 +4148,7 @@ impl RequestContext {
super::TEMP_RAG_NAME, super::TEMP_RAG_NAME,
&rag_path, &rag_path,
&[], &[],
abort_signal, abort_signal.clone(),
false, false,
) )
.await?, .await?,
@@ -4172,10 +4184,11 @@ impl RequestContext {
}; };
self.rag = Some(rag); self.rag = Some(rag);
self.rag_key = rag_key; self.rag_key = rag_key;
self.refresh_tool_scope(abort_signal).await?;
Ok(()) Ok(())
} }
pub async fn attach_rag(&mut self, name: &str) -> Result<()> { pub async fn attach_rag(&mut self, name: &str, abort_signal: AbortSignal) -> Result<()> {
let rag_path = self.rag_file(name); let rag_path = self.rag_file(name);
if rag_path.exists() { if rag_path.exists() {
bail!( bail!(
@@ -4192,6 +4205,7 @@ impl RequestContext {
self.rag_cache().insert(key.clone(), &rag); self.rag_cache().insert(key.clone(), &rag);
self.rag = Some(rag); self.rag = Some(rag);
self.rag_key = Some(key); self.rag_key = Some(key);
self.refresh_tool_scope(abort_signal).await?;
Ok(()) Ok(())
} }
+466 -97
View File
@@ -1,4 +1,5 @@
pub(crate) mod memory; pub(crate) mod memory;
pub(crate) mod rag_query;
pub(crate) mod skill; pub(crate) mod skill;
pub(crate) mod supervisor; pub(crate) mod supervisor;
pub(crate) mod todo; pub(crate) mod todo;
@@ -23,20 +24,22 @@ use futures_util::future;
use indexmap::IndexMap; use indexmap::IndexMap;
use indoc::formatdoc; use indoc::formatdoc;
use memory::MEMORY_FUNCTION_PREFIX; use memory::MEMORY_FUNCTION_PREFIX;
use rag_query::RAG_FUNCTION_PREFIX;
use rust_embed::Embed; use rust_embed::Embed;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::{Value, json}; use serde_json::{Value, json};
use skill::SKILL_FUNCTION_PREFIX; use skill::SKILL_FUNCTION_PREFIX;
use std::collections::VecDeque;
use std::ffi::OsStr; use std::ffi::OsStr;
use std::fs::File; use std::fs::File;
use std::io::{Read, Write}; use std::io::{Read, Write};
use std::sync::atomic::Ordering; use std::sync::atomic::{AtomicU64, Ordering};
use std::{collections::VecDeque, thread};
use std::{ use std::{
collections::{HashMap, HashSet}, collections::{HashMap, HashSet},
env, fs, io, env, fs, io,
path::{Path, PathBuf}, path::{Path, PathBuf},
process::{Command, Stdio}, process::{Command, Stdio},
time::{Duration, Instant},
}; };
use strum_macros::AsRefStr; use strum_macros::AsRefStr;
use supervisor::SUPERVISOR_FUNCTION_PREFIX; use supervisor::SUPERVISOR_FUNCTION_PREFIX;
@@ -135,6 +138,114 @@ fn extract_shebang_runtime(path: &Path) -> Option<String> {
} }
} }
pub(crate) fn write_file_atomic(
path: &Path,
content: &str,
#[cfg_attr(not(unix), expect(unused))] mode: Option<u32>,
) -> Result<()> {
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;
if fs::read_to_string(path).is_ok_and(|existing| existing == content) {
#[cfg(unix)]
if let Some(mode) = mode {
fs::set_permissions(path, fs::Permissions::from_mode(mode))?;
}
return Ok(());
}
let file_name = path
.file_name()
.and_then(OsStr::to_str)
.ok_or_else(|| anyhow!("Unable to extract file name from path: {}", path.display()))?;
static TMP_COUNTER: AtomicU64 = AtomicU64::new(0);
let tmp = path.with_file_name(format!(
".{file_name}.tmp.{}.{}",
std::process::id(),
TMP_COUNTER.fetch_add(1, Ordering::Relaxed)
));
fs::write(&tmp, content)?;
#[cfg(unix)]
if let Some(mode) = mode {
fs::set_permissions(&tmp, fs::Permissions::from_mode(mode))?;
}
if let Err(err) = fs::rename(&tmp, path) {
let _ = fs::remove_file(&tmp);
return Err(err.into());
}
Ok(())
}
fn tool_source_stems() -> Result<HashSet<String>> {
let mut stems = HashSet::new();
let tools_dir = paths::global_tools_dir();
if !tools_dir.exists() {
return Ok(stems);
}
for entry in fs::read_dir(&tools_dir)? {
let path = entry?.path();
if path.is_file()
&& let Some(stem) = path.file_stem().and_then(OsStr::to_str)
{
stems.insert(stem.to_string());
}
}
Ok(stems)
}
fn bin_entry_stem(file_name: &str) -> &str {
let name = file_name.strip_prefix("run-").unwrap_or(file_name);
Path::new(name)
.file_stem()
.and_then(OsStr::to_str)
.unwrap_or(name)
}
fn prune_stale_bin_entries(
bin_dir: &Path,
valid_stems: &HashSet<String>,
extra_valid_stem: Option<&str>,
) -> Result<()> {
if !bin_dir.exists() {
fs::create_dir_all(bin_dir)?;
return Ok(());
}
for entry in fs::read_dir(bin_dir)? {
let entry = entry?;
let path = entry.path();
if path.is_dir() {
debug!(
"Removing unexpected directory in bin dir: {}",
path.display()
);
fs::remove_dir_all(&path)?;
continue;
}
let file_name = entry.file_name();
let Some(file_name) = file_name.to_str() else {
continue;
};
let stem = bin_entry_stem(file_name);
if valid_stems.contains(stem) || extra_valid_stem == Some(stem) {
continue;
}
debug!("Removing stale bin entry: {}", path.display());
fs::remove_file(&path)?;
}
Ok(())
}
pub async fn eval_tool_calls( pub async fn eval_tool_calls(
ctx: &mut RequestContext, ctx: &mut RequestContext,
mut calls: Vec<ToolCall>, mut calls: Vec<ToolCall>,
@@ -183,42 +294,29 @@ pub async fn eval_tool_calls(
}) })
.collect(); .collect();
for (idx, call, result) in future::join_all(futs).await { for (idx, call, result) in future::join_all(futs).await {
indexed_results.push((idx, ToolResult::new(call, normalize_tool_result(result?)))); let value = match result {
Ok(v) => normalize_tool_result(v),
Err(e) => json!({"tool_call_error": format!("{e}")}),
};
indexed_results.push((idx, ToolResult::new(call, value)));
} }
} }
for (idx, call) in sequential_calls { for (idx, call) in sequential_calls {
let result = call.eval(ctx).await?; let value = match call.eval(ctx).await {
indexed_results.push((idx, ToolResult::new(call, normalize_tool_result(result)))); Ok(v) => normalize_tool_result(v),
Err(e) => json!({
"tool_call_error": format!(
"{e}. This tool is not available or the call failed; use only tools listed in your catalog."
)
}),
};
indexed_results.push((idx, ToolResult::new(call, value)));
} }
indexed_results.sort_unstable_by_key(|(idx, _)| *idx); indexed_results.sort_unstable_by_key(|(idx, _)| *idx);
output = indexed_results.into_iter().map(|(_, r)| r).collect(); output = indexed_results.into_iter().map(|(_, r)| r).collect();
if !output.is_empty() {
let (has_escalations, summary) = if ctx.current_depth == 0
&& let Some(queue) = ctx.root_escalation_queue()
&& queue.has_pending()
{
(true, queue.pending_summary())
} else {
(false, vec![])
};
if has_escalations {
let notification = json!({
"pending_escalations": summary,
"instruction": "Child agents are BLOCKED waiting for your reply. Call agent__reply_escalation for each pending escalation to unblock them."
});
let synthetic_call = ToolCall::new(
"__escalation_notification".to_string(),
json!({}),
Some("escalation_check".to_string()),
);
output.push(ToolResult::new(synthetic_call, notification));
}
}
{ {
let max_chars = ctx let max_chars = ctx
.agent .agent
@@ -233,6 +331,14 @@ pub async fn eval_tool_calls(
} }
} }
if ctx.current_depth == 0
&& let Some(queue) = ctx.root_escalation_queue()
&& queue.has_pending()
&& let Some(last) = output.last_mut()
{
inject_escalation_notification(last, queue.pending_summary());
}
Ok(output) Ok(output)
} }
@@ -249,6 +355,24 @@ fn normalize_tool_result(result: Value) -> Value {
} }
} }
fn inject_escalation_notification(last: &mut ToolResult, summary: Vec<Value>) {
let instruction = "Child agents are BLOCKED waiting for your reply. \
Call agent__reply_escalation for each pending escalation to unblock them.";
match &mut last.output {
Value::Object(map) => {
map.insert("pending_escalations".into(), json!(summary));
map.insert("escalation_instruction".into(), json!(instruction));
}
other => {
*other = json!({
"output": other.take(),
"pending_escalations": summary,
"escalation_instruction": instruction,
});
}
}
}
#[derive(Debug, Clone, Deserialize, Serialize)] #[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ToolResult { pub struct ToolResult {
pub call: ToolCall, pub call: ToolCall,
@@ -305,7 +429,6 @@ impl Functions {
})?; })?;
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) }; let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
let file_path = paths::functions_dir().join(file.as_ref()); let file_path = paths::functions_dir().join(file.as_ref());
#[cfg_attr(not(unix), expect(unused))]
let is_script = file_path let is_script = file_path
.extension() .extension()
.and_then(OsStr::to_str) .and_then(OsStr::to_str)
@@ -322,14 +445,7 @@ impl Functions {
ensure_parent_exists(&file_path)?; ensure_parent_exists(&file_path)?;
info!("Creating function file: {}", file_path.display()); info!("Creating function file: {}", file_path.display());
let mut function_file = File::create(&file_path)?; write_file_atomic(&file_path, content, is_script.then_some(0o755))?;
function_file.write_all(content.as_bytes())?;
#[cfg(unix)]
if is_script {
use std::os::unix::fs::PermissionsExt;
fs::set_permissions(&file_path, fs::Permissions::from_mode(0o755))?;
}
} }
Ok(()) Ok(())
@@ -380,7 +496,7 @@ impl Functions {
} }
pub fn init(visible_tools: &[String]) -> Result<Self> { pub fn init(visible_tools: &[String]) -> Result<Self> {
Self::clear_global_functions_bin_dir()?; Self::remove_stale_global_function_binaries()?;
let declarations = Self { let declarations = Self {
declarations: Self::build_global_tool_declarations(visible_tools)?, declarations: Self::build_global_tool_declarations(visible_tools)?,
@@ -396,7 +512,7 @@ impl Functions {
} }
pub fn init_agent(name: &str, global_tools: &[String]) -> Result<Self> { pub fn init_agent(name: &str, global_tools: &[String]) -> Result<Self> {
Self::clear_agent_bin_dir(name)?; Self::remove_stale_agent_bin_entries(name)?;
let global_tools_declarations = if !global_tools.is_empty() { let global_tools_declarations = if !global_tools.is_empty() {
info!("Loading global tools for agent: {name}: {global_tools:?}"); info!("Loading global tools for agent: {name}: {global_tools:?}");
@@ -495,6 +611,16 @@ impl Functions {
.extend(user_interaction::user_interaction_function_declarations()); .extend(user_interaction::user_interaction_function_declarations());
} }
pub fn append_rag_query_functions(&mut self) {
self.declarations
.extend(rag_query::rag_query_function_declarations());
}
pub fn remove_rag_query_functions(&mut self) {
self.declarations
.retain(|f| !f.name.starts_with(RAG_FUNCTION_PREFIX));
}
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) { pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
let mut invoke_function_properties = IndexMap::new(); let mut invoke_function_properties = IndexMap::new();
invoke_function_properties.insert( invoke_function_properties.insert(
@@ -700,38 +826,23 @@ impl Functions {
Ok(()) Ok(())
} }
fn clear_agent_bin_dir(name: &str) -> Result<()> { fn remove_stale_agent_bin_entries(name: &str) -> Result<()> {
let agent_bin_directory = paths::agent_bin_dir(name); let agent_bin_directory = paths::agent_bin_dir(name);
if !agent_bin_directory.exists() {
debug!(
"Creating agent bin directory: {}",
agent_bin_directory.display()
);
fs::create_dir_all(&agent_bin_directory)?;
} else {
debug!(
"Clearing existing agent bin directory: {}",
agent_bin_directory.display()
);
clear_dir(&agent_bin_directory)?;
}
Ok(()) debug!(
"Pruning stale entries in agent bin directory: {}",
agent_bin_directory.display()
);
prune_stale_bin_entries(&agent_bin_directory, &tool_source_stems()?, Some(name))
} }
fn clear_global_functions_bin_dir() -> Result<()> { fn remove_stale_global_function_binaries() -> Result<()> {
let bin_dir = paths::functions_bin_dir(); let bin_dir = paths::functions_bin_dir();
if !bin_dir.exists() {
fs::create_dir_all(&bin_dir)?;
}
info!( info!("Pruning stale function binaries in {}", bin_dir.display());
"Clearing existing function binaries in {}",
bin_dir.display()
);
clear_dir(&bin_dir)?;
Ok(()) prune_stale_bin_entries(&bin_dir, &tool_source_stems()?, None)
} }
fn build_agent_tool_binaries(name: &str) -> Result<()> { fn build_agent_tool_binaries(name: &str) -> Result<()> {
@@ -832,11 +943,7 @@ impl Functions {
"{prompt_utils_file}", "{prompt_utils_file}",
&to_script_path(&paths::bash_prompt_utils_file().to_string_lossy()), &to_script_path(&paths::bash_prompt_utils_file().to_string_lossy()),
); );
if binary_script_file.exists() { write_file_atomic(&binary_script_file, &content, None)?;
fs::remove_file(&binary_script_file)?;
}
let mut script_file = File::create(&binary_script_file)?;
script_file.write_all(content.as_bytes())?;
info!( info!(
"Building binary for function: {} ({})", "Building binary for function: {} ({})",
@@ -898,8 +1005,7 @@ impl Functions {
{run} "{wrapper_binary}" %*"#, {run} "{wrapper_binary}" %*"#,
); );
let mut file = File::create(&binary_file)?; write_file_atomic(&binary_file, &content, None)?;
file.write_all(content.as_bytes())?;
Ok(()) Ok(())
} }
@@ -911,8 +1017,6 @@ impl Functions {
binary_type: BinaryType, binary_type: BinaryType,
custom_runtime: Option<&str>, custom_runtime: Option<&str>,
) -> Result<()> { ) -> Result<()> {
use std::os::unix::prelude::PermissionsExt;
let binary_file = match binary_type { let binary_file = match binary_type {
BinaryType::Tool(None) => paths::functions_bin_dir().join(binary_name), BinaryType::Tool(None) => paths::functions_bin_dir().join(binary_name),
BinaryType::Tool(Some(agent_name)) => { BinaryType::Tool(Some(agent_name)) => {
@@ -981,31 +1085,16 @@ impl Functions {
.parent() .parent()
.expect("Failed to get parent directory of binary file"); .expect("Failed to get parent directory of binary file");
let script_file = bin_dir.join(format!("run-{binary_name}.ts")); let script_file = bin_dir.join(format!("run-{binary_name}.ts"));
if script_file.exists() { write_file_atomic(&script_file, &content, Some(0o755))?;
fs::remove_file(&script_file)?;
}
let mut sf = File::create(&script_file)?;
sf.write_all(content.as_bytes())?;
fs::set_permissions(&script_file, fs::Permissions::from_mode(0o755))?;
let ts_runtime = custom_runtime.unwrap_or("tsx"); let ts_runtime = custom_runtime.unwrap_or("tsx");
let wrapper = format!( let wrapper = format!(
"#!/bin/sh\nexec {ts_runtime} \"{}\" \"$@\"\n", "#!/bin/sh\nexec {ts_runtime} \"{}\" \"$@\"\n",
script_file.display() script_file.display()
); );
if binary_file.exists() { write_file_atomic(&binary_file, &wrapper, Some(0o755))?;
fs::remove_file(&binary_file)?;
}
let mut wf = File::create(&binary_file)?;
wf.write_all(wrapper.as_bytes())?;
fs::set_permissions(&binary_file, fs::Permissions::from_mode(0o755))?;
} else { } else {
if binary_file.exists() { write_file_atomic(&binary_file, &content, Some(0o755))?;
fs::remove_file(&binary_file)?;
}
let mut file = File::create(&binary_file)?;
file.write_all(content.as_bytes())?;
fs::set_permissions(&binary_file, fs::Permissions::from_mode(0o755))?;
} }
Ok(()) Ok(())
@@ -1252,6 +1341,15 @@ impl ToolCall {
json!({"tool_call_error": error_msg}) json!({"tool_call_error": error_msg})
}) })
} }
_ if cmd_name.starts_with(RAG_FUNCTION_PREFIX) => {
rag_query::handle_rag_tool(ctx, &cmd_name, &json_data)
.await
.unwrap_or_else(|e| {
let error_msg = format!("RAG query failed: {e}");
eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️")));
json!({"tool_call_error": error_msg})
})
}
_ => match run_llm_function(cmd_name, cmd_args, envs, agent_name) { _ => match run_llm_function(cmd_name, cmd_args, envs, agent_name) {
Ok(Some(contents)) => serde_json::from_str(&contents) Ok(Some(contents)) => serde_json::from_str(&contents)
.ok() .ok()
@@ -1435,6 +1533,7 @@ pub fn run_llm_function(
let mut child = Command::new(&cmd_name) let mut child = Command::new(&cmd_name)
.args(&cmd_args) .args(&cmd_args)
.envs(envs) .envs(envs)
.stdin(Stdio::null())
.stdout(Stdio::piped()) .stdout(Stdio::piped())
.stderr(Stdio::piped()) .stderr(Stdio::piped())
.spawn() .spawn()
@@ -1443,7 +1542,7 @@ pub fn run_llm_function(
let stdout = child.stdout.take().expect("Failed to capture stdout"); let stdout = child.stdout.take().expect("Failed to capture stdout");
let stderr = child.stderr.take().expect("Failed to capture stderr"); let stderr = child.stderr.take().expect("Failed to capture stderr");
let stdout_thread = std::thread::spawn(move || { let stdout_thread = thread::spawn(move || {
let mut buffer = [0; 1024]; let mut buffer = [0; 1024];
let mut reader = stdout; let mut reader = stdout;
let mut out = io::stdout(); let mut out = io::stdout();
@@ -1470,7 +1569,7 @@ pub fn run_llm_function(
buf buf
}); });
let stderr_thread = std::thread::spawn(move || { let stderr_thread = thread::spawn(move || {
let mut buffer = [0; 1024]; let mut buffer = [0; 1024];
let mut reader = stderr; let mut reader = stderr;
let mut err = io::stderr(); let mut err = io::stderr();
@@ -1497,9 +1596,39 @@ pub fn run_llm_function(
buf buf
}); });
let status = child let timeout_secs = env::var("COYOTE_TOOL_TIMEOUT")
.wait() .ok()
.map_err(|err| anyhow!("Unable to run {command_name}, {err}"))?; .and_then(|v| v.parse::<u64>().ok())
.unwrap_or(1800);
let deadline = (timeout_secs > 0).then(|| Instant::now() + Duration::from_secs(timeout_secs));
let status = loop {
match child.try_wait() {
Ok(Some(status)) => break status,
Ok(None) => {}
Err(err) => bail!("Unable to run {command_name}, {err}"),
}
if let Some(deadline) = deadline
&& Instant::now() >= deadline
{
let _ = child.kill();
let _ = child.wait();
drop(stdout_thread);
drop(stderr_thread);
let tool_error_message = format!(
"Tool call '{command_name}' timed out after {timeout_secs}s and was killed (set COYOTE_TOOL_TIMEOUT to adjust; 0 = unlimited)"
);
eprintln!(
"{}",
muted_warning_text(&format!("⚠️ {tool_error_message} ⚠️"))
);
let error_json = json!({"tool_call_error": tool_error_message});
debug!("Tool call error: {error_json:?}");
return Ok(Some(error_json.to_string()));
}
thread::sleep(Duration::from_millis(100));
};
let stdout_bytes = stdout_thread.join().unwrap_or_default(); let stdout_bytes = stdout_thread.join().unwrap_or_default();
let stderr_bytes = stderr_thread.join().unwrap_or_default(); let stderr_bytes = stderr_thread.join().unwrap_or_default();
@@ -1519,6 +1648,11 @@ pub fn run_llm_function(
if !stdout.is_empty() { if !stdout.is_empty() {
error_json["stdout"] = json!(stdout); error_json["stdout"] = json!(stdout);
} }
if let Ok(contents) = fs::read_to_string(&tmp_file)
&& !contents.trim().is_empty()
{
error_json["output"] = json!(contents);
}
debug!("Tool call error: {error_json:?}"); debug!("Tool call error: {error_json:?}");
return Ok(Some(error_json.to_string())); return Ok(Some(error_json.to_string()));
} }
@@ -1688,7 +1822,10 @@ fn format_json_colored_keys(value: &serde_json::Value) -> String {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::config::{AppState, WorkingMode};
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
use serde_json::json; use serde_json::json;
use std::sync::Arc;
fn call(name: &str, id: Option<&str>) -> ToolCall { fn call(name: &str, id: Option<&str>) -> ToolCall {
ToolCall::new(name.to_string(), json!({}), id.map(|s| s.to_string())) ToolCall::new(name.to_string(), json!({}), id.map(|s| s.to_string()))
@@ -1698,11 +1835,98 @@ mod tests {
ToolCall::new(name.to_string(), args, Some("id1".to_string())) ToolCall::new(name.to_string(), args, Some("id1".to_string()))
} }
fn run_async<F: Future>(f: F) -> F::Output {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(f)
}
fn submit_escalation(queue: &EscalationQueue, id: &str) {
let (tx, _rx) = tokio::sync::oneshot::channel();
queue.submit(EscalationRequest {
id: id.to_string(),
from_agent_id: "a1".into(),
from_agent_name: "explore".into(),
question: "What do?".into(),
options: None,
reply_tx: tx,
});
}
#[test] #[test]
fn normalize_tool_result_substitutes_done_for_null() { fn normalize_tool_result_substitutes_done_for_null() {
assert_eq!(normalize_tool_result(Value::Null), json!("DONE")); assert_eq!(normalize_tool_result(Value::Null), json!("DONE"));
} }
#[test]
fn inject_escalation_notification_extends_object_output() {
let mut result = ToolResult::new(call("t", Some("id-1")), json!({"status": "ok"}));
inject_escalation_notification(&mut result, vec![json!({"escalation_id": "esc_1"})]);
assert_eq!(result.output["status"], "ok");
assert_eq!(
result.output["pending_escalations"],
json!([{"escalation_id": "esc_1"}])
);
assert!(
result.output["escalation_instruction"]
.as_str()
.unwrap()
.contains("agent__reply_escalation")
);
assert!(result.text.is_none());
}
#[test]
fn inject_escalation_notification_wraps_non_object_output() {
let mut result = ToolResult::new(call("t", Some("id-1")), json!("DONE"));
inject_escalation_notification(&mut result, vec![json!({"escalation_id": "esc_2"})]);
assert_eq!(result.output["output"], json!("DONE"));
assert_eq!(
result.output["pending_escalations"],
json!([{"escalation_id": "esc_2"}])
);
assert!(result.output["escalation_instruction"].is_string());
}
#[test]
fn eval_tool_calls_soft_fails_unknown_tool() {
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
let calls = vec![call("__escalation_notification", Some("id-1"))];
let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap();
assert_eq!(results.len(), 1);
let err = results[0].output["tool_call_error"].as_str().unwrap();
assert!(err.contains("Unexpected call"));
assert!(err.contains("use only tools listed in your catalog"));
}
#[test]
fn eval_tool_calls_injects_escalations_into_last_result() {
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
let queue = ctx.ensure_root_escalation_queue();
submit_escalation(&queue, "esc_1");
let calls = vec![call("unknown_tool", Some("id-1"))];
let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap();
assert_eq!(results.len(), 1);
assert!(
results
.iter()
.all(|r| r.call.name != "__escalation_notification")
);
let out = &results[0].output;
assert!(out["tool_call_error"].is_string());
assert_eq!(out["pending_escalations"][0]["escalation_id"], "esc_1");
assert!(
out["escalation_instruction"]
.as_str()
.unwrap()
.contains("agent__reply_escalation")
);
}
#[test] #[test]
fn normalize_tool_result_preserves_non_null_values() { fn normalize_tool_result_preserves_non_null_values() {
assert_eq!( assert_eq!(
@@ -2130,4 +2354,149 @@ mod tests {
let tc = call_with_args("t", json!(42)); let tc = call_with_args("t", json!(42));
assert!(tc.parse_arguments().is_err()); assert!(tc.parse_arguments().is_err());
} }
#[test]
fn write_file_atomic_writes_and_skips_unchanged() {
let dir = temp_file("-atomic-", "");
fs::create_dir_all(&dir).unwrap();
let path = dir.join("shim");
write_file_atomic(&path, "one", Some(0o755)).unwrap();
assert_eq!(fs::read_to_string(&path).unwrap(), "one");
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
assert_eq!(
fs::metadata(&path).unwrap().permissions().mode() & 0o777,
0o755
);
}
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
let ino = fs::metadata(&path).unwrap().ino();
write_file_atomic(&path, "one", Some(0o755)).unwrap();
assert_eq!(
fs::metadata(&path).unwrap().ino(),
ino,
"unchanged content must not be rewritten"
);
}
write_file_atomic(&path, "two", None).unwrap();
assert_eq!(fs::read_to_string(&path).unwrap(), "two");
assert_eq!(
fs::read_dir(&dir).unwrap().count(),
1,
"no tmp files left behind"
);
fs::remove_dir_all(&dir).unwrap();
}
#[test]
fn write_file_atomic_concurrent_writers_to_same_target() {
let dir = temp_file("-atomic-concurrent-", "");
fs::create_dir_all(&dir).unwrap();
let path = dir.join("shim");
let contents: Vec<String> = (0..8)
.map(|i| format!("#!/bin/sh\necho writer-{i}\n"))
.collect();
thread::scope(|scope| {
for content in &contents {
scope.spawn(|| {
for _ in 0..50 {
write_file_atomic(&path, content, Some(0o755)).unwrap();
}
});
}
});
let final_content = fs::read_to_string(&path).unwrap();
assert!(
contents.contains(&final_content),
"final content must be one writer's complete content, got: {final_content:?}"
);
assert_eq!(
fs::read_dir(&dir).unwrap().count(),
1,
"no tmp files left behind"
);
fs::remove_dir_all(&dir).unwrap();
}
#[cfg(unix)]
#[test]
fn run_llm_function_includes_llm_output_on_nonzero_exit() {
let result = run_llm_function(
"bash".into(),
vec![
"-c".into(),
"echo partial-output >> \"$LLM_OUTPUT\"; echo err-text >&2; exit 3".into(),
],
HashMap::new(),
None,
)
.unwrap()
.expect("nonzero exit must return an error payload");
let json: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(
json["tool_call_error"]
.as_str()
.unwrap()
.contains("exited with code 3")
);
assert_eq!(json["stderr"], "err-text");
assert_eq!(json["output"], "partial-output\n");
}
#[test]
fn bin_entry_stem_strips_run_prefix_and_extension() {
assert_eq!(bin_entry_stem("fs_grep"), "fs_grep");
assert_eq!(bin_entry_stem("fs_grep.cmd"), "fs_grep");
assert_eq!(bin_entry_stem("run-web_search.ts"), "web_search");
assert_eq!(bin_entry_stem("run-fs_grep.sh"), "fs_grep");
}
#[test]
fn prune_stale_bin_entries_removes_only_stale_files() {
let dir = temp_file("-prune-", "");
fs::create_dir_all(dir.join("nested")).unwrap();
for name in [
"fs_grep",
"run-web_search.ts",
"old_tool",
"run-old_tool.ts",
"myagent",
] {
fs::write(dir.join(name), "x").unwrap();
}
let valid_stems: HashSet<String> = ["fs_grep", "web_search"]
.iter()
.map(|s| s.to_string())
.collect();
prune_stale_bin_entries(&dir, &valid_stems, Some("myagent")).unwrap();
assert!(dir.join("fs_grep").exists());
assert!(dir.join("run-web_search.ts").exists());
assert!(dir.join("myagent").exists(), "agent binary must survive");
assert!(!dir.join("old_tool").exists());
assert!(!dir.join("run-old_tool.ts").exists());
assert!(!dir.join("nested").exists(), "directories must be removed");
fs::remove_dir_all(&dir).unwrap();
}
#[test]
fn prune_stale_bin_entries_creates_missing_dir() {
let dir = temp_file("-prune-missing-", "");
prune_stale_bin_entries(&dir, &HashSet::new(), None).unwrap();
assert!(dir.is_dir());
fs::remove_dir_all(&dir).unwrap();
}
} }
+101
View File
@@ -0,0 +1,101 @@
use super::{FunctionDeclaration, JsonSchema};
use crate::config::RequestContext;
use anyhow::{Result, anyhow};
use indexmap::IndexMap;
use serde_json::{Value, json};
pub const RAG_FUNCTION_PREFIX: &str = "rag__";
pub fn rag_query_function_declarations() -> Vec<FunctionDeclaration> {
vec![FunctionDeclaration {
name: format!("{RAG_FUNCTION_PREFIX}query"),
description: "Search the RAG knowledge base attached to this session and return \
the most relevant text chunks with their source paths. The relevant \
context has already been injected into the prompt up-front; use this \
tool to pull additional context on-demand when the initial retrieval \
does not fully answer the question. Prefer specific, keyword-rich queries."
.to_string(),
parameters: JsonSchema {
type_value: Some("object".to_string()),
properties: Some(IndexMap::from([
(
"query".to_string(),
JsonSchema {
type_value: Some("string".to_string()),
description: Some(
"Natural language search query used to retrieve relevant chunks."
.into(),
),
..Default::default()
},
),
(
"top_k".to_string(),
JsonSchema {
type_value: Some("integer".to_string()),
description: Some(
"Maximum number of chunks to return. Defaults to the RAG's \
configured top_k when omitted."
.into(),
),
..Default::default()
},
),
])),
required: Some(vec!["query".to_string()]),
..Default::default()
},
agent: false,
}]
}
pub async fn handle_rag_tool(
ctx: &mut RequestContext,
cmd_name: &str,
args: &Value,
) -> Result<Value> {
let action = cmd_name
.strip_prefix(RAG_FUNCTION_PREFIX)
.unwrap_or(cmd_name);
match action {
"query" => handle_query(ctx, args).await,
_ => Err(anyhow!("Unknown RAG action: {action}")),
}
}
async fn handle_query(ctx: &RequestContext, args: &Value) -> Result<Value> {
let rag = ctx
.rag
.clone()
.ok_or_else(|| anyhow!("No RAG is attached to this session"))?;
let query = args
.get("query")
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'query' is required"))?;
let top_k = args
.get("top_k")
.and_then(Value::as_u64)
.map(|v| v as usize)
.unwrap_or_else(|| rag.configured_top_k());
let rerank_model = rag.configured_reranker().map(|s| s.to_string());
let chunks = rag
.search_chunks(query, top_k, rerank_model.as_deref())
.await?;
let chunks_json: Vec<Value> = chunks
.into_iter()
.map(|(text, source)| json!({ "text": text, "source": source }))
.collect();
Ok(json!({
"rag_name": rag.name(),
"count": chunks_json.len(),
"chunks": chunks_json,
}))
}
+109 -1
View File
@@ -85,7 +85,17 @@ pub fn check_pending_agents_guardrail(ctx: &mut RequestContext) -> GuardrailActi
} }
ctx.pending_agents_guardrail_count += 1; ctx.pending_agents_guardrail_count += 1;
GuardrailAction::Inject(build_pending_agents_guardrail_prompt(&pending)) let mut prompt = build_pending_agents_guardrail_prompt(&pending);
if let Some(queue) = ctx.root_escalation_queue()
&& queue.has_pending()
{
let summary = serde_json::to_string(&queue.pending_summary()).unwrap_or_default();
prompt.push_str(&format!(
"\n\nAdditionally, child agents have pending escalations blocking them. Reply to each \
via `agent__reply_escalation` first:\n{summary}"
));
}
GuardrailAction::Inject(prompt)
} }
pub fn escalation_function_declarations() -> Vec<FunctionDeclaration> { pub fn escalation_function_declarations() -> Vec<FunctionDeclaration> {
@@ -1996,4 +2006,102 @@ mod tests {
let q2 = ctx.ensure_root_escalation_queue(); let q2 = ctx.ensure_root_escalation_queue();
assert!(Arc::ptr_eq(&q1, &q2)); assert!(Arc::ptr_eq(&q1, &q2));
} }
#[test]
fn guardrail_prompt_mentions_pending_escalations() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let mut ctx = ctx_with_supervisor(4, 3);
let join_handle = tokio::spawn(async {
time::sleep(Duration::from_secs(60)).await;
Ok(AgentResult {
id: "slow".into(),
agent_name: "test".into(),
output: String::new(),
exit_status: AgentExitStatus::Completed,
})
});
let handle = AgentHandle {
id: "slow".into(),
agent_name: "test".into(),
depth: 1,
inbox: Arc::new(Inbox::new()),
abort_signal: create_abort_signal(),
join_handle,
child_supervisor: None,
};
ctx.supervisor
.as_ref()
.unwrap()
.write()
.register(handle)
.unwrap();
let queue = ctx.ensure_root_escalation_queue();
let (tx, _rx) = tokio::sync::oneshot::channel();
queue.submit(EscalationRequest {
id: "esc_9".into(),
from_agent_id: "a1".into(),
from_agent_name: "explore".into(),
question: "Which option?".into(),
options: None,
reply_tx: tx,
});
match check_pending_agents_guardrail(&mut ctx) {
GuardrailAction::Inject(prompt) => {
assert!(prompt.contains("agent__reply_escalation"));
assert!(prompt.contains("esc_9"));
}
_ => panic!("expected Inject action"),
}
});
}
#[test]
fn guardrail_prompt_omits_escalations_when_none_pending() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let mut ctx = ctx_with_supervisor(4, 3);
let join_handle = tokio::spawn(async {
time::sleep(Duration::from_secs(60)).await;
Ok(AgentResult {
id: "slow".into(),
agent_name: "test".into(),
output: String::new(),
exit_status: AgentExitStatus::Completed,
})
});
let handle = AgentHandle {
id: "slow".into(),
agent_name: "test".into(),
depth: 1,
inbox: Arc::new(Inbox::new()),
abort_signal: create_abort_signal(),
join_handle,
child_supervisor: None,
};
ctx.supervisor
.as_ref()
.unwrap()
.write()
.register(handle)
.unwrap();
match check_pending_agents_guardrail(&mut ctx) {
GuardrailAction::Inject(prompt) => {
assert!(!prompt.contains("agent__reply_escalation"));
}
_ => panic!("expected Inject action"),
}
});
}
} }
+64 -1
View File
@@ -250,7 +250,30 @@ impl GraphExecutor {
branch_tasks.push(task); branch_tasks.push(task);
} }
let joined = join_all(branch_tasks).await; let joined = match graph_timeout {
Some(t) => {
let remaining = t.saturating_sub(start.elapsed());
let abort_handles: Vec<_> = branch_tasks
.iter()
.map(|task| task.abort_handle())
.collect();
match tokio::time::timeout(remaining, join_all(branch_tasks)).await {
Ok(joined) => joined,
Err(_) => {
for handle in abort_handles {
handle.abort();
}
bail!(
"Graph '{}' timed out after {}s during super-step with frontier {:?}",
graph.name,
t.as_secs(),
sorted_frontier(&frontier)
);
}
}
}
None => join_all(branch_tasks).await,
};
let mut branch_writes: Vec<BranchWrites> = Vec::new(); let mut branch_writes: Vec<BranchWrites> = Vec::new();
let mut next_frontier: HashSet<String> = HashSet::new(); let mut next_frontier: HashSet<String> = HashSet::new();
@@ -793,4 +816,44 @@ nodes:
"error should list both End nodes: {err}" "error should list both End nodes: {err}"
); );
} }
#[tokio::test]
async fn graph_timeout_interrupts_in_flight_super_step() {
if !cmd_available("bash") {
eprintln!("skipping: bash not available");
return;
}
let ws = TestWorkspace::new();
ws.write_script("sleeper.sh", "#!/bin/bash\nsleep 10\necho '{}'\n");
let yaml = r#"
name: inflight_timeout_test
start: sleeper
settings:
timeout: 1
nodes:
sleeper:
type: script
script: sleeper.sh
state_updates: {}
next: done
done:
type: end
output: "done"
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
let mut ctx = make_ctx();
let abort = create_abort_signal();
let result = GraphExecutor::new(graph, &ws.dir)
.execute(&mut ctx, abort)
.await;
assert!(result.is_err(), "expected in-flight timeout to error");
let err = format!("{:#}", result.unwrap_err());
assert!(
err.contains("timed out after 1s during super-step"),
"error should report during-super-step timeout: {err}"
);
assert!(err.contains("sleeper"), "error should name frontier: {err}");
}
} }
+1
View File
@@ -54,6 +54,7 @@ impl ScriptExecutor {
let mut cmd = build_command(language, &script_path)?; let mut cmd = build_command(language, &script_path)?;
cmd.stdout(Stdio::piped()); cmd.stdout(Stdio::piped());
cmd.stderr(Stdio::piped()); cmd.stderr(Stdio::piped());
cmd.kill_on_drop(true);
cmd.envs(&self.extra_envs); cmd.envs(&self.extra_envs);
cmd.env("AUTO_CONFIRM", "true"); cmd.env("AUTO_CONFIRM", "true");
match &state_repr { match &state_repr {
+12
View File
@@ -196,6 +196,18 @@ async fn main() -> Result<()> {
return Ok(()); return Ok(());
} }
let mcp_action =
cli.mcp_list || cli.mcp_get.is_some() || cli.mcp_remove.is_some() || cli.mcp_add.is_some();
if mcp_action {
let cfg = Config::load_with_interpolation(true).await?;
let app_config = AppConfig::from_config(cfg)?;
let vault = Vault::init(&app_config)?;
mcp::manage::handle(&cli, &vault)?;
return Ok(());
}
if vault_flags { if vault_flags {
let cfg = Config::load_with_interpolation(true).await?; let cfg = Config::load_with_interpolation(true).await?;
let app_config = AppConfig::from_config(cfg)?; let app_config = AppConfig::from_config(cfg)?;
+651
View File
@@ -0,0 +1,651 @@
use crate::mcp::oauth::{force_refresh_mcp_token, load_or_refresh_mcp_token};
use http::{HeaderName, HeaderValue};
use log::debug;
use rmcp::model::ClientJsonRpcMessage;
use rmcp::transport::common::client_side_sse::BoxedSseResponse;
use rmcp::transport::streamable_http_client::{
AuthRequiredError, StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse,
};
use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;
/// [`StreamableHttpClient`] wrapper that injects the OAuth bearer token for an
/// MCP server on every request instead of pinning it at spawn time, so tokens
/// refreshed mid-session take effect without reconnecting.
///
/// A caller-supplied `auth_header` always passes through untouched; only a
/// `None` header is filled from the stored token. When the wrapper injected
/// the token and a POST comes back 401, it forces a token refresh and retries
/// exactly once (see [`Self::post_with_retry`]).
#[derive(Clone)]
pub struct McpOAuthClient<C = reqwest::Client> {
inner: C,
server: Arc<str>,
}
impl<C> McpOAuthClient<C> {
pub fn new(inner: C, server: &str) -> Self {
Self {
inner,
server: Arc::from(server),
}
}
}
impl<C: StreamableHttpClient + Sync> McpOAuthClient<C> {
/// Resolves the effective auth header. Caller-supplied values pass through
/// untouched; `None` is filled from the stored token for this server.
/// Returns the header plus whether the wrapper injected it. Errors with
/// [`StreamableHttpError::AuthRequired`] (without contacting the server)
/// when no usable token exists.
async fn resolve_auth(
&self,
auth_header: Option<String>,
) -> Result<(Option<String>, bool), StreamableHttpError<C::Error>> {
if auth_header.is_some() {
return Ok((auth_header, false));
}
match load_or_refresh_mcp_token(&self.server).await.into_token() {
Some(token) => Ok((Some(token), true)),
None => Err(self.auth_required()),
}
}
fn auth_required(&self) -> StreamableHttpError<C::Error> {
StreamableHttpError::AuthRequired(AuthRequiredError::new(format!(
"no valid OAuth token for MCP server '{server}'; \
run `.mcp auth {server}` to re-authenticate",
server = self.server
)))
}
async fn post_with_retry<F, Fut>(
&self,
auth_header: Option<String>,
mut post: F,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<C::Error>>
where
F: FnMut(Option<String>) -> Fut,
Fut: Future<Output = Result<StreamableHttpPostResponse, StreamableHttpError<C::Error>>>,
{
let (auth, injected) = self.resolve_auth(auth_header).await?;
let (original, rejected) = match (post(auth.clone()).await, auth) {
(Err(err @ StreamableHttpError::AuthRequired(_)), Some(rejected)) if injected => {
(err, rejected)
}
(result, _) => return result,
};
debug!(
"MCP server '{}' rejected the injected token; forcing a refresh and retrying once",
self.server
);
let Some(token) = force_refresh_mcp_token(&self.server, &rejected).await else {
return Err(original);
};
match post(Some(token)).await {
Err(StreamableHttpError::AuthRequired(_)) => {
debug!(
"Retry after forced token refresh was rejected again by MCP server '{}'",
self.server
);
Err(original)
}
result => result,
}
}
}
impl<C: StreamableHttpClient + Sync> StreamableHttpClient for McpOAuthClient<C> {
type Error = C::Error;
async fn post_message(
&self,
uri: Arc<str>,
message: ClientJsonRpcMessage,
session_id: Option<Arc<str>>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
self.post_with_retry(auth_header, |auth| {
self.inner.post_message(
uri.clone(),
message.clone(),
session_id.clone(),
auth,
custom_headers.clone(),
)
})
.await
}
/// Overridden rather than left to the trait default: the default impl
/// delegates to [`Self::post_message`], silently dropping the
/// transport-wide SSE event size limit. Delegating to the inner client's
/// size-enforcing variant keeps the limit applied at the raw byte layer.
async fn post_message_with_max_sse_event_size(
&self,
uri: Arc<str>,
message: ClientJsonRpcMessage,
session_id: Option<Arc<str>>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
max_sse_event_size: usize,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
self.post_with_retry(auth_header, |auth| {
self.inner.post_message_with_max_sse_event_size(
uri.clone(),
message.clone(),
session_id.clone(),
auth,
custom_headers.clone(),
max_sse_event_size,
)
})
.await
}
async fn delete_session(
&self,
uri: Arc<str>,
session_id: Arc<str>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<(), StreamableHttpError<Self::Error>> {
let (auth, _) = self.resolve_auth(auth_header).await?;
self.inner
.delete_session(uri, session_id, auth, custom_headers)
.await
}
async fn get_stream(
&self,
uri: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<BoxedSseResponse, StreamableHttpError<Self::Error>> {
let (auth, _) = self.resolve_auth(auth_header).await?;
self.inner
.get_stream(uri, session_id, last_event_id, auth, custom_headers)
.await
}
/// Overridden for the same reason as
/// [`Self::post_message_with_max_sse_event_size`]: the trait default
/// bypasses SSE event size enforcement.
async fn get_stream_with_max_sse_event_size(
&self,
uri: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
max_sse_event_size: usize,
) -> Result<BoxedSseResponse, StreamableHttpError<Self::Error>> {
let (auth, _) = self.resolve_auth(auth_header).await?;
self.inner
.get_stream_with_max_sse_event_size(
uri,
session_id,
last_event_id,
auth,
custom_headers,
max_sse_event_size,
)
.await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::paths;
use crate::mcp::oauth::test_support::with_temp_cache;
use futures_util::StreamExt;
use parking_lot::Mutex;
use serial_test::serial;
use std::convert::Infallible;
use std::fs;
use std::sync::atomic::{AtomicUsize, Ordering};
const FRESH: i64 = 9999999999;
type Calls = Arc<Mutex<Vec<(&'static str, Option<String>)>>>;
type AfterFirstCallAction = Option<Box<dyn FnOnce() + Send + 'static>>;
/// Inner client that records `(method, auth_header)` per call, rejects the
/// first `reject_times` POSTs/streams with `AuthRequired` (then the next
/// `transport_error_times` with a non-auth error), and runs an optional
/// side effect after the first call (to mutate token files between the
/// initial attempt and the retry).
#[derive(Clone, Default)]
struct FakeInner {
calls: Calls,
reject_times: Arc<AtomicUsize>,
transport_error_times: Arc<AtomicUsize>,
after_first_call: Arc<Mutex<AfterFirstCallAction>>,
}
impl FakeInner {
fn record(
&self,
method: &'static str,
auth: Option<String>,
) -> Result<(), StreamableHttpError<Infallible>> {
self.calls.lock().push((method, auth));
if let Some(f) = self.after_first_call.lock().take() {
f();
}
if self.reject_times.load(Ordering::SeqCst) > 0 {
self.reject_times.fetch_sub(1, Ordering::SeqCst);
return Err(Self::rejection());
}
if self.transport_error_times.load(Ordering::SeqCst) > 0 {
self.transport_error_times.fetch_sub(1, Ordering::SeqCst);
return Err(StreamableHttpError::UnexpectedServerResponse(
"connection reset".into(),
));
}
Ok(())
}
fn rejection() -> StreamableHttpError<Infallible> {
StreamableHttpError::AuthRequired(AuthRequiredError::new(
"Bearer error=\"invalid_token\"".to_string(),
))
}
}
impl StreamableHttpClient for FakeInner {
type Error = Infallible;
async fn post_message(
&self,
_uri: Arc<str>,
_message: ClientJsonRpcMessage,
_session_id: Option<Arc<str>>,
auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
self.record("post_message", auth_header)?;
Ok(StreamableHttpPostResponse::Accepted)
}
async fn post_message_with_max_sse_event_size(
&self,
_uri: Arc<str>,
_message: ClientJsonRpcMessage,
_session_id: Option<Arc<str>>,
auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
_max_sse_event_size: usize,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
self.record("post_message_with_max_sse_event_size", auth_header)?;
Ok(StreamableHttpPostResponse::Accepted)
}
async fn delete_session(
&self,
_uri: Arc<str>,
_session_id: Arc<str>,
auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<(), StreamableHttpError<Self::Error>> {
self.record("delete_session", auth_header)?;
Ok(())
}
async fn get_stream(
&self,
_uri: Arc<str>,
_session_id: Option<Arc<str>>,
_last_event_id: Option<String>,
auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<BoxedSseResponse, StreamableHttpError<Self::Error>> {
self.record("get_stream", auth_header)?;
Ok(futures_util::stream::empty().boxed())
}
async fn get_stream_with_max_sse_event_size(
&self,
_uri: Arc<str>,
_session_id: Option<Arc<str>>,
_last_event_id: Option<String>,
auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
_max_sse_event_size: usize,
) -> Result<BoxedSseResponse, StreamableHttpError<Self::Error>> {
self.record("get_stream_with_max_sse_event_size", auth_header)?;
Ok(futures_util::stream::empty().boxed())
}
}
fn write_token_file(server: &str, access_token: &str, expires_at: i64) {
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
fs::write(
paths::token_file(&format!("mcp_{server}")),
format!(
r#"{{"access_token":"{access_token}","refresh_token":"r","expires_at":{expires_at}}}"#
),
)
.unwrap();
}
fn ping() -> ClientJsonRpcMessage {
serde_json::from_value(serde_json::json!({
"jsonrpc": "2.0",
"id": 1,
"method": "ping"
}))
.unwrap()
}
fn rt() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
}
fn post(
client: &McpOAuthClient<FakeInner>,
auth_header: Option<String>,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Infallible>> {
rt().block_on(client.post_message(
Arc::from("http://mcp.test/mcp"),
ping(),
None,
auth_header,
HashMap::new(),
))
}
#[test]
#[serial]
fn injects_token_from_disk_when_auth_header_none() {
with_temp_cache(|| {
write_token_file("wrapper-inject", "tok-live", FRESH);
let inner = FakeInner::default();
let client = McpOAuthClient::new(inner.clone(), "wrapper-inject");
let result = post(&client, None);
assert!(result.is_ok());
assert_eq!(
*inner.calls.lock(),
vec![("post_message", Some("tok-live".to_string()))]
);
});
}
#[test]
#[serial]
fn caller_supplied_auth_header_passes_through() {
with_temp_cache(|| {
write_token_file("wrapper-passthrough", "tok-disk", FRESH);
let inner = FakeInner::default();
let client = McpOAuthClient::new(inner.clone(), "wrapper-passthrough");
let result = post(&client, Some("caller-tok".to_string()));
assert!(result.is_ok());
assert_eq!(
*inner.calls.lock(),
vec![("post_message", Some("caller-tok".to_string()))]
);
});
}
#[test]
#[serial]
fn caller_supplied_header_rejection_propagates_without_refresh() {
with_temp_cache(|| {
// A fresh, different token sits on disk: if the injected guard
// were dropped, the wrapper would refresh and retry with it.
write_token_file("wrapper-caller-401", "tok-disk", FRESH);
let inner = FakeInner::default();
inner.reject_times.store(1, Ordering::SeqCst);
let client = McpOAuthClient::new(inner.clone(), "wrapper-caller-401");
let result = post(&client, Some("caller-tok".to_string()));
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
assert_eq!(
*inner.calls.lock(),
vec![("post_message", Some("caller-tok".to_string()))]
);
});
}
#[test]
#[serial]
fn missing_token_returns_auth_required_without_calling_inner() {
with_temp_cache(|| {
let inner = FakeInner::default();
let client = McpOAuthClient::new(inner.clone(), "wrapper-no-token");
let result = post(&client, None);
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
assert!(inner.calls.lock().is_empty());
});
}
#[test]
#[serial]
fn rejected_token_forces_refresh_and_retries_once() {
with_temp_cache(|| {
write_token_file("wrapper-retry", "tok-a", FRESH);
let inner = FakeInner::default();
inner.reject_times.store(1, Ordering::SeqCst);
// Simulate a concurrent refresh landing between the rejection and
// the forced refresh: the retry must carry the new token.
*inner.after_first_call.lock() = Some(Box::new(|| {
write_token_file("wrapper-retry", "tok-b", FRESH);
}));
let client = McpOAuthClient::new(inner.clone(), "wrapper-retry");
let result = post(&client, None);
assert!(result.is_ok());
assert_eq!(
*inner.calls.lock(),
vec![
("post_message", Some("tok-a".to_string())),
("post_message", Some("tok-b".to_string())),
]
);
});
}
#[test]
#[serial]
fn failed_force_refresh_propagates_original_error_after_one_call() {
with_temp_cache(|| {
write_token_file("wrapper-refresh-fail", "tok-a", FRESH);
let inner = FakeInner::default();
inner.reject_times.store(1, Ordering::SeqCst);
// Token file gone by refresh time: force refresh yields nothing.
*inner.after_first_call.lock() = Some(Box::new(|| {
fs::remove_file(paths::token_file("mcp_wrapper-refresh-fail")).unwrap();
}));
let client = McpOAuthClient::new(inner.clone(), "wrapper-refresh-fail");
let result = post(&client, None);
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
assert_eq!(inner.calls.lock().len(), 1);
});
}
#[test]
#[serial]
fn second_rejection_after_retry_propagates_original_error() {
with_temp_cache(|| {
write_token_file("wrapper-double-401", "tok-a", FRESH);
let inner = FakeInner::default();
inner.reject_times.store(2, Ordering::SeqCst);
// A changed token appears before the forced refresh, so the retry
// actually runs (an unchanged token would trigger a real refresh
// attempt, which fails without a cached registration).
*inner.after_first_call.lock() = Some(Box::new(|| {
write_token_file("wrapper-double-401", "tok-b", FRESH);
}));
let client = McpOAuthClient::new(inner.clone(), "wrapper-double-401");
let result = post(&client, None);
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
assert_eq!(inner.calls.lock().len(), 2);
});
}
#[test]
#[serial]
fn non_auth_retry_error_propagates_as_is() {
with_temp_cache(|| {
write_token_file("wrapper-retry-transport", "tok-a", FRESH);
let inner = FakeInner::default();
inner.reject_times.store(1, Ordering::SeqCst);
inner.transport_error_times.store(1, Ordering::SeqCst);
*inner.after_first_call.lock() = Some(Box::new(|| {
write_token_file("wrapper-retry-transport", "tok-b", FRESH);
}));
let client = McpOAuthClient::new(inner.clone(), "wrapper-retry-transport");
let result = post(&client, None);
assert!(matches!(
result,
Err(StreamableHttpError::UnexpectedServerResponse(_))
));
assert_eq!(inner.calls.lock().len(), 2);
});
}
#[test]
#[serial]
fn sized_post_delegates_to_inner_sized_variant() {
with_temp_cache(|| {
write_token_file("wrapper-sized-post", "tok-live", FRESH);
let inner = FakeInner::default();
let client = McpOAuthClient::new(inner.clone(), "wrapper-sized-post");
let result = rt().block_on(client.post_message_with_max_sse_event_size(
Arc::from("http://mcp.test/mcp"),
ping(),
None,
None,
HashMap::new(),
4096,
));
assert!(result.is_ok());
assert_eq!(
*inner.calls.lock(),
vec![(
"post_message_with_max_sse_event_size",
Some("tok-live".to_string())
)]
);
});
}
#[test]
#[serial]
fn sized_get_stream_delegates_to_inner_sized_variant() {
with_temp_cache(|| {
write_token_file("wrapper-sized-get", "tok-live", FRESH);
let inner = FakeInner::default();
let client = McpOAuthClient::new(inner.clone(), "wrapper-sized-get");
let result = rt().block_on(client.get_stream_with_max_sse_event_size(
Arc::from("http://mcp.test/mcp"),
None,
None,
None,
HashMap::new(),
4096,
));
assert!(result.is_ok());
assert_eq!(
*inner.calls.lock(),
vec![(
"get_stream_with_max_sse_event_size",
Some("tok-live".to_string())
)]
);
});
}
#[test]
#[serial]
fn get_stream_does_not_retry_on_rejection() {
with_temp_cache(|| {
write_token_file("wrapper-get-401", "tok-live", FRESH);
let inner = FakeInner::default();
inner.reject_times.store(1, Ordering::SeqCst);
let client = McpOAuthClient::new(inner.clone(), "wrapper-get-401");
let result = rt().block_on(client.get_stream(
Arc::from("http://mcp.test/mcp"),
None,
None,
None,
HashMap::new(),
));
assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_))));
assert_eq!(inner.calls.lock().len(), 1);
});
}
#[test]
#[serial]
fn delete_session_injects_token() {
with_temp_cache(|| {
write_token_file("wrapper-delete", "tok-live", FRESH);
let inner = FakeInner::default();
let client = McpOAuthClient::new(inner.clone(), "wrapper-delete");
let result = rt().block_on(client.delete_session(
Arc::from("http://mcp.test/mcp"),
Arc::from("session-1"),
None,
HashMap::new(),
));
assert!(result.is_ok());
assert_eq!(
*inner.calls.lock(),
vec![("delete_session", Some("tok-live".to_string()))]
);
});
}
#[test]
#[serial]
fn auth_required_error_contains_no_token_material() {
with_temp_cache(|| {
write_token_file("wrapper-redact", "stale-secret-token", 0);
let inner = FakeInner::default();
let client = McpOAuthClient::new(inner.clone(), "wrapper-redact");
let err = post(&client, None).unwrap_err();
let display = format!("{err}");
let debug = format!("{err:?}");
assert!(!display.contains("stale-secret-token"));
assert!(!debug.contains("stale-secret-token"));
assert!(inner.calls.lock().is_empty());
});
}
}
+443
View File
@@ -0,0 +1,443 @@
use crate::cli::{Cli, McpScopeArg, McpTransportArg};
use crate::config::{ensure_parent_exists, paths};
use crate::mcp::{JsonField, McpOAuthConfig, McpServer, McpServersConfig, McpTransportType};
use crate::vault::{SECRET_RE, Vault};
use anyhow::{Context, Result, anyhow, bail};
use indexmap::{IndexMap, IndexSet};
use inquire::Confirm;
use std::collections::HashSet;
use std::fs;
use std::path::{Path, PathBuf};
impl From<McpTransportArg> for McpTransportType {
fn from(value: McpTransportArg) -> Self {
match value {
McpTransportArg::Stdio => McpTransportType::Stdio,
McpTransportArg::Http => McpTransportType::Http,
McpTransportArg::Sse => McpTransportType::Sse,
}
}
}
pub fn handle(cli: &Cli, vault: &Vault) -> Result<()> {
if cli.mcp_list {
return handle_list(cli.scope);
}
if let Some(name) = &cli.mcp_get {
return handle_get(name, cli.scope);
}
if let Some(name) = &cli.mcp_remove {
return handle_remove(name, cli.scope, cli.mcp_force);
}
if let Some(name) = &cli.mcp_add {
return handle_add(cli, name, vault);
}
Ok(())
}
fn handle_list(scope: Option<McpScopeArg>) -> Result<()> {
let show_user = scope != Some(McpScopeArg::Workspace);
let show_workspace = scope != Some(McpScopeArg::User);
if show_user {
let user_path = paths::mcp_config_file();
let user_cfg = load_config_raw(&user_path)?;
println!("User ({})", user_path.display());
print_server_list(&user_cfg);
}
if show_workspace {
match paths::workspace_mcp_config_file() {
Some(ws_path) => {
let ws_cfg = load_config_raw(&ws_path)?;
if show_user {
println!();
}
println!("Workspace ({})", ws_path.display());
print_server_list(&ws_cfg);
}
None if scope == Some(McpScopeArg::Workspace) => {
println!("Workspace: no mcp.json found in current directory");
}
None => {}
}
}
Ok(())
}
fn print_server_list(cfg: &McpServersConfig) {
if cfg.mcp_servers.is_empty() {
println!(" (none)");
return;
}
let name_width = cfg.mcp_servers.keys().map(String::len).max().unwrap_or(0);
for (name, spec) in &cfg.mcp_servers {
let transport = match spec.transport_type {
McpTransportType::Stdio => "stdio",
McpTransportType::Http => "http",
McpTransportType::Sse => "sse",
};
let target = spec.url.clone().unwrap_or_else(|| {
let cmd = spec.command.clone().unwrap_or_default();
let args = spec.args.as_ref().map(|a| a.join(" ")).unwrap_or_default();
if args.is_empty() {
cmd
} else {
format!("{cmd} {args}")
}
});
println!(
" {name:<name_width$} {transport:<5} {target}",
name_width = name_width
);
}
}
fn handle_get(name: &str, scope: Option<McpScopeArg>) -> Result<()> {
let (path, cfg) = load_for_scope_or_search(name, scope)?;
let spec = cfg
.mcp_servers
.get(name)
.ok_or_else(|| anyhow!("MCP server '{name}' not found"))?;
let pretty =
serde_json::to_string_pretty(spec).context("failed to serialize MCP server config")?;
println!("# {}", path.display());
println!("{pretty}");
Ok(())
}
fn handle_remove(name: &str, scope: Option<McpScopeArg>, force: bool) -> Result<()> {
let (path, mut cfg) = load_for_scope_or_search(name, scope)?;
if !force {
let ok = Confirm::new(&format!(
"Remove MCP server '{name}' from {}?",
path.display()
))
.with_default(false)
.prompt()?;
if !ok {
println!("Aborted.");
return Ok(());
}
}
cfg.mcp_servers.shift_remove(name);
save_config(&path, &cfg)?;
println!("✓ Removed MCP server '{name}' from {}", path.display());
Ok(())
}
fn handle_add(cli: &Cli, name: &str, vault: &Vault) -> Result<()> {
validate_name(name)?;
let server = build_server(cli)?;
server.validate(name)?;
let scope = cli.scope.unwrap_or_default();
let path = write_path_for_scope(scope);
let mut cfg = load_config_raw(&path)?;
if cfg.mcp_servers.contains_key(name) && !cli.mcp_force {
let ok = Confirm::new(&format!(
"MCP server '{name}' already exists in {}. Overwrite?",
path.display()
))
.with_default(false)
.prompt()?;
if !ok {
println!("Aborted. Use --mcp-force to overwrite without prompting.");
return Ok(());
}
}
provision_secrets(cli, vault)?;
cfg.mcp_servers.insert(name.to_string(), server);
save_config(&path, &cfg)?;
println!("✓ Added MCP server '{name}' to {}", path.display());
Ok(())
}
fn validate_name(name: &str) -> Result<()> {
if name.is_empty() {
bail!("MCP server name cannot be empty");
}
if !name
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
{
bail!("Invalid MCP server name '{name}': only letters, digits, '-', and '_' are allowed");
}
Ok(())
}
fn build_server(cli: &Cli) -> Result<McpServer> {
let has_command = !cli.mcp_command.is_empty();
let has_url = cli.url.is_some();
let transport = cli
.transport
.map(McpTransportType::from)
.unwrap_or_else(|| {
if has_command {
McpTransportType::Stdio
} else {
McpTransportType::Http
}
});
match transport {
McpTransportType::Stdio => build_stdio(cli, has_url),
McpTransportType::Http | McpTransportType::Sse => build_remote(cli, transport, has_command),
}
}
fn build_stdio(cli: &Cli, has_url: bool) -> Result<McpServer> {
if cli.mcp_command.is_empty() {
bail!(
"stdio MCP server requires a command. Pass it after `--`, e.g. \
`--mcp-add NAME -- npx some-server --flag`"
);
}
if has_url {
bail!("stdio MCP server does not accept --url");
}
if !cli.header.is_empty() {
bail!("stdio MCP server does not accept --header");
}
if cli.client_id.is_some()
|| cli.client_secret.is_some()
|| cli.callback_port.is_some()
|| cli.redirect_host.is_some()
{
bail!("stdio MCP server does not accept OAuth flags");
}
let (cmd, args) = cli.mcp_command.split_first().unwrap();
let mut env: IndexMap<String, JsonField> = IndexMap::new();
for kv in &cli.env {
let (k, v) = kv
.split_once('=')
.ok_or_else(|| anyhow!("invalid --env value '{kv}': expected KEY=VALUE"))?;
if k.is_empty() {
bail!("invalid --env value '{kv}': KEY cannot be empty");
}
env.insert(k.to_string(), JsonField::Str(v.to_string()));
}
Ok(McpServer {
transport_type: McpTransportType::Stdio,
command: Some(cmd.clone()),
args: (!args.is_empty()).then(|| args.to_vec()),
env: (!env.is_empty()).then_some(env),
cwd: cli.cwd.clone(),
url: None,
headers: None,
oauth: None,
})
}
fn build_remote(cli: &Cli, transport: McpTransportType, has_command: bool) -> Result<McpServer> {
if has_command {
bail!(
"http/sse MCP server does not accept a trailing `-- <cmd>`. Use `--url` \
to specify the endpoint."
);
}
let url = cli
.url
.clone()
.ok_or_else(|| anyhow!("http/sse MCP server requires --url <URL>"))?;
if !cli.env.is_empty() {
bail!("http/sse MCP server does not accept --env; use --header instead");
}
if cli.cwd.is_some() {
bail!("http/sse MCP server does not accept --cwd");
}
let mut headers: IndexMap<String, String> = IndexMap::new();
for h in &cli.header {
let (name, value) = h
.split_once(':')
.ok_or_else(|| anyhow!("invalid --header value '{h}': expected 'Name: Value'"))?;
let name = name.trim();
let value = value.trim_start_matches(' ');
if name.is_empty() {
bail!("invalid --header value '{h}': header name cannot be empty");
}
headers.insert(name.to_string(), value.to_string());
}
let oauth = if cli.client_id.is_some()
|| cli.client_secret.is_some()
|| cli.callback_port.is_some()
|| cli.redirect_host.is_some()
{
Some(McpOAuthConfig {
client_id: cli.client_id.clone(),
client_secret: cli.client_secret.clone(),
callback_port: cli.callback_port,
redirect_host: cli.redirect_host.clone(),
})
} else {
None
};
Ok(McpServer {
transport_type: transport,
command: None,
args: None,
env: None,
cwd: None,
url: Some(url),
headers: (!headers.is_empty()).then_some(headers),
oauth,
})
}
fn provision_secrets(cli: &Cli, vault: &Vault) -> Result<()> {
let mut sources: Vec<&str> = Vec::new();
if let Some(s) = cli.url.as_deref() {
sources.push(s);
}
if let Some(s) = cli.client_secret.as_deref() {
sources.push(s);
}
if let Some(s) = cli.client_id.as_deref() {
sources.push(s);
}
if let Some(s) = cli.redirect_host.as_deref() {
sources.push(s);
}
if let Some(s) = cli.cwd.as_deref() {
sources.push(s);
}
sources.extend(cli.env.iter().map(String::as_str));
sources.extend(cli.header.iter().map(String::as_str));
let mut needed: IndexSet<String> = IndexSet::new();
for value in sources {
for caps in SECRET_RE.captures_iter(value).filter_map(Result::ok) {
if let Some(m) = caps.get(1) {
needed.insert(m.as_str().trim().to_string());
}
}
}
if needed.is_empty() {
return Ok(());
}
let existing: HashSet<String> = vault.list_secrets(false)?.into_iter().collect();
for name in needed {
if existing.contains(&name) {
continue;
}
eprintln!("Value references vault secret {{{{ {name} }}}} which is not stored yet.");
let ok = Confirm::new(&format!("Add '{name}' to the vault now?"))
.with_default(true)
.prompt()?;
if !ok {
bail!(
"Vault secret '{name}' is required by the config; aborting. \
Add it later with `coyote --add-secret {name}`."
);
}
vault.add_secret(&name)?;
}
Ok(())
}
fn load_for_scope_or_search(
name: &str,
scope: Option<McpScopeArg>,
) -> Result<(PathBuf, McpServersConfig)> {
if let Some(s) = scope {
let path = match s {
McpScopeArg::User => paths::mcp_config_file(),
McpScopeArg::Workspace => paths::workspace_mcp_config_file()
.ok_or_else(|| anyhow!("no workspace mcp.json found in the current directory"))?,
};
let cfg = load_config_raw(&path)?;
if !cfg.mcp_servers.contains_key(name) {
bail!(
"MCP server '{name}' not found in {} scope ({})",
scope_label(s),
path.display()
);
}
return Ok((path, cfg));
}
let user_path = paths::mcp_config_file();
let user_cfg = load_config_raw(&user_path)?;
if user_cfg.mcp_servers.contains_key(name) {
return Ok((user_path, user_cfg));
}
if let Some(ws_path) = paths::workspace_mcp_config_file() {
let ws_cfg = load_config_raw(&ws_path)?;
if ws_cfg.mcp_servers.contains_key(name) {
return Ok((ws_path, ws_cfg));
}
}
bail!("MCP server '{name}' not found in any scope");
}
fn write_path_for_scope(scope: McpScopeArg) -> PathBuf {
match scope {
McpScopeArg::User => paths::mcp_config_file(),
McpScopeArg::Workspace => paths::workspace_mcp_config_file()
.unwrap_or_else(|| paths::workspace_config_dir().join("mcp.json")),
}
}
fn scope_label(scope: McpScopeArg) -> &'static str {
match scope {
McpScopeArg::User => "user",
McpScopeArg::Workspace => "workspace",
}
}
fn load_config_raw(path: &Path) -> Result<McpServersConfig> {
if !path.exists() {
return Ok(McpServersConfig {
mcp_servers: IndexMap::new(),
});
}
let raw = fs::read_to_string(path)
.with_context(|| format!("failed to read MCP config at {}", path.display()))?;
if raw.trim().is_empty() {
return Ok(McpServersConfig {
mcp_servers: IndexMap::new(),
});
}
serde_json::from_str(&raw)
.with_context(|| format!("failed to parse MCP config at {}", path.display()))
}
fn save_config(path: &Path, config: &McpServersConfig) -> Result<()> {
ensure_parent_exists(path)?;
let serialized =
serde_json::to_string_pretty(config).context("failed to serialize MCP config")?;
let tmp = path.with_extension("json.tmp");
fs::write(&tmp, &serialized)
.with_context(|| format!("failed to write temporary MCP config at {}", tmp.display()))?;
fs::rename(&tmp, path)
.with_context(|| format!("failed to finalize MCP config at {}", path.display()))?;
Ok(())
}
+317 -13
View File
@@ -1,3 +1,5 @@
mod auth_client;
pub(crate) mod manage;
pub(crate) mod oauth; pub(crate) mod oauth;
mod sse_transport; mod sse_transport;
@@ -8,6 +10,7 @@ use crate::vault::Vault;
use crate::vault::interpolate_secrets; use crate::vault::interpolate_secrets;
use anyhow::Error; use anyhow::Error;
use anyhow::{Context, Result, anyhow}; use anyhow::{Context, Result, anyhow};
use auth_client::McpOAuthClient;
use futures_util::{StreamExt, TryStreamExt, stream}; use futures_util::{StreamExt, TryStreamExt, stream};
use http::{HeaderName, HeaderValue}; use http::{HeaderName, HeaderValue};
use indexmap::IndexMap; use indexmap::IndexMap;
@@ -20,6 +23,8 @@ use rmcp::{RoleClient, ServiceExt};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use sse_transport::LegacySseTransport; use sse_transport::LegacySseTransport;
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
use std::fmt;
use std::fmt::Display;
use std::fs::OpenOptions; use std::fs::OpenOptions;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::process::Stdio; use std::process::Stdio;
@@ -62,6 +67,8 @@ pub(crate) struct McpServersConfig {
pub(crate) struct McpOAuthConfig { pub(crate) struct McpOAuthConfig {
#[serde(rename = "clientId", skip_serializing_if = "Option::is_none")] #[serde(rename = "clientId", skip_serializing_if = "Option::is_none")]
pub client_id: Option<String>, pub client_id: Option<String>,
#[serde(rename = "clientSecret", skip_serializing_if = "Option::is_none")]
pub client_secret: Option<String>,
#[serde(rename = "callbackPort", skip_serializing_if = "Option::is_none")] #[serde(rename = "callbackPort", skip_serializing_if = "Option::is_none")]
pub callback_port: Option<u16>, pub callback_port: Option<u16>,
#[serde(rename = "redirectHost", skip_serializing_if = "Option::is_none")] #[serde(rename = "redirectHost", skip_serializing_if = "Option::is_none")]
@@ -323,18 +330,17 @@ impl McpRegistry {
.and_then(|c| c.mcp_servers.get(&id)) .and_then(|c| c.mcp_servers.get(&id))
.with_context(|| format!("MCP server not found in config: {id}"))?; .with_context(|| format!("MCP server not found in config: {id}"))?;
let bearer_token = if spec.is_remote() { let (auth, auth_reason) = resolve_http_auth(&id, spec).await;
oauth::load_valid_mcp_token(&id)
} else {
None
};
let service = match spawn_mcp_server(spec, self.log_path.as_deref(), bearer_token).await { let service = match spawn_mcp_server(spec, self.log_path.as_deref(), auth).await {
Ok(s) => s, Ok(s) => s,
Err(e) if is_auth_required_error(&e) => { Err(e) if is_auth_required_error(&e) => {
warn!( warn!(
"MCP server '{id}' requires OAuth authentication. \ "{}",
Run `coyote --auth-mcp {id}` or `.mcp auth {id}` in the REPL to authenticate." McpAuthRequired {
server: id,
reason: auth_reason,
}
); );
return Ok(None); return Ok(None);
} }
@@ -409,19 +415,76 @@ impl McpRegistry {
} }
} }
/// How a remote MCP server authenticates outgoing requests.
pub(crate) enum HttpAuth {
/// Only the static headers from the server spec; no OAuth token.
StaticOnly,
/// OAuth-managed: HTTP transports inject a fresh bearer token per request
/// via [`McpOAuthClient`] (ignoring the carried token); SSE transports
/// send the carried token as a static header.
Managed { server: String, token: String },
}
impl fmt::Debug for HttpAuth {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::StaticOnly => f.write_str("StaticOnly"),
Self::Managed { server, token: _ } => f
.debug_struct("Managed")
.field("server", server)
.field("token", &"<redacted>")
.finish(),
}
}
}
impl HttpAuth {
pub(crate) fn from_token_status(status: &oauth::McpTokenStatus, server: &str) -> Self {
match status {
oauth::McpTokenStatus::Token(token) => Self::Managed {
server: server.to_string(),
token: token.clone(),
},
oauth::McpTokenStatus::NotAuthenticated | oauth::McpTokenStatus::RefreshFailed => {
Self::StaticOnly
}
}
}
}
pub(crate) async fn resolve_http_auth(name: &str, spec: &McpServer) -> (HttpAuth, McpAuthReason) {
let token_status = if spec.is_remote() {
oauth::load_or_refresh_mcp_token(name).await
} else {
oauth::McpTokenStatus::NotAuthenticated
};
(
HttpAuth::from_token_status(&token_status, name),
McpAuthReason::from_token_status(&token_status),
)
}
pub(crate) async fn spawn_mcp_server( pub(crate) async fn spawn_mcp_server(
spec: &McpServer, spec: &McpServer,
log_path: Option<&Path>, log_path: Option<&Path>,
bearer_token: Option<String>, auth: HttpAuth,
) -> Result<Arc<ConnectedServer>> { ) -> Result<Arc<ConnectedServer>> {
match spec.transport_type { match spec.transport_type {
McpTransportType::Http => { McpTransportType::Http => {
let url = spec.url.as_deref().expect("validated: http spec has url"); let url = spec.url.as_deref().expect("validated: http spec has url");
let headers = merge_bearer_token(spec.headers.as_ref(), bearer_token); match auth {
spawn_http_mcp_server(url, headers.as_ref()).await HttpAuth::Managed { server, token: _ } => {
spawn_oauth_http_mcp_server(url, &server, spec.headers.as_ref()).await
}
HttpAuth::StaticOnly => spawn_http_mcp_server(url, spec.headers.as_ref()).await,
}
} }
McpTransportType::Sse => { McpTransportType::Sse => {
let url = spec.url.as_deref().expect("validated: sse spec has url"); let url = spec.url.as_deref().expect("validated: sse spec has url");
let bearer_token = match auth {
HttpAuth::Managed { server: _, token } => Some(token),
HttpAuth::StaticOnly => None,
};
let headers = merge_bearer_token(spec.headers.as_ref(), bearer_token); let headers = merge_bearer_token(spec.headers.as_ref(), bearer_token);
spawn_sse_mcp_server(url, headers.as_ref()).await spawn_sse_mcp_server(url, headers.as_ref()).await
} }
@@ -449,15 +512,66 @@ fn merge_bearer_token(
} }
(Some(h), Some(token)) => { (Some(h), Some(token)) => {
let mut m = h.clone(); let mut m = h.clone();
m.retain(|k, _| !k.eq_ignore_ascii_case("authorization"));
m.insert("Authorization".to_string(), format!("Bearer {token}")); m.insert("Authorization".to_string(), format!("Bearer {token}"));
Some(m) Some(m)
} }
} }
} }
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum McpAuthReason {
NotAuthenticated,
RefreshFailed,
TokenRejected,
}
impl McpAuthReason {
pub(crate) fn from_token_status(status: &oauth::McpTokenStatus) -> Self {
match status {
oauth::McpTokenStatus::Token(_) => Self::TokenRejected,
oauth::McpTokenStatus::NotAuthenticated => Self::NotAuthenticated,
oauth::McpTokenStatus::RefreshFailed => Self::RefreshFailed,
}
}
}
#[derive(Debug)]
pub(crate) struct McpAuthRequired {
pub server: String,
pub reason: McpAuthReason,
}
impl Display for McpAuthRequired {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let server = &self.server;
match self.reason {
McpAuthReason::NotAuthenticated => write!(
f,
"MCP server '{server}' requires OAuth authentication and was not started \
(no stored credentials). Run `.mcp auth {server}` (or `coyote --auth-mcp \
{server}`) to authenticate and attach it."
),
McpAuthReason::RefreshFailed => write!(
f,
"MCP server '{server}' was not started: stored OAuth token has expired and \
automatic refresh failed. Run `.mcp auth {server}` (or `coyote --auth-mcp \
{server}`) to re-authenticate and attach it."
),
McpAuthReason::TokenRejected => write!(
f,
"MCP server '{server}' was not started: the server rejected the stored OAuth \
token. Run `.mcp auth {server}` (or `coyote --auth-mcp {server}`) to \
re-authenticate and attach it."
),
}
}
}
pub(crate) fn is_auth_required_error(e: &Error) -> bool { pub(crate) fn is_auth_required_error(e: &Error) -> bool {
e.chain() e.downcast_ref::<McpAuthRequired>().is_some()
.any(|cause| cause.to_string().contains("Auth required")) || e.chain()
.any(|cause| cause.to_string().contains("Auth required"))
} }
async fn spawn_http_mcp_server( async fn spawn_http_mcp_server(
@@ -490,6 +604,66 @@ async fn spawn_http_mcp_server(
Ok(service) Ok(service)
} }
/// Builds the custom-header map for an OAuth-managed HTTP transport, dropping
/// any static `Authorization` entry case-insensitively: [`McpOAuthClient`]
/// owns that header, and a stale configured value must not collide with the
/// per-request token.
fn oauth_custom_headers(
headers: Option<&IndexMap<String, String>>,
) -> Result<HashMap<HeaderName, HeaderValue>> {
let mut custom = HashMap::new();
let Some(hdrs) = headers else {
return Ok(custom);
};
for (k, v) in hdrs {
if k.eq_ignore_ascii_case("authorization") {
continue;
}
let name = k
.parse::<HeaderName>()
.with_context(|| format!("Invalid header name: {k}"))?;
let value = v
.parse::<HeaderValue>()
.with_context(|| format!("Invalid header value for {k}"))?;
custom.insert(name, value);
}
Ok(custom)
}
async fn spawn_oauth_http_mcp_server(
url: &str,
server: &str,
headers: Option<&IndexMap<String, String>>,
) -> Result<Arc<ConnectedServer>> {
// Mirror rmcp's default_http_client, which `with_client` bypasses:
// idle pooling off avoids a documented TCP delayed-ACK stall, and
// redirects off keeps custom headers from being replayed to a redirect
// target.
let inner = reqwest::Client::builder()
.pool_max_idle_per_host(0)
.redirect(reqwest::redirect::Policy::none())
.build()
.context("Failed to build HTTP client for OAuth-managed MCP transport")?;
let client = McpOAuthClient::new(inner, server);
// `auth_header` stays None so the wrapper injects a fresh token per
// request; `reinit_on_expired_session` defaults to true in rmcp 3.1.2
// but is pinned explicitly because transparent session re-init is
// load-bearing for long-lived sessions.
let config = StreamableHttpClientTransportConfig::with_uri(url)
.custom_headers(oauth_custom_headers(headers)?)
.reinit_on_expired_session(true);
let transport = StreamableHttpClientTransport::with_client(client, config);
let service = Arc::new(
().serve(transport)
.await
.with_context(|| format!("Failed to connect to HTTP MCP server: {url}"))?,
);
Ok(service)
}
async fn spawn_sse_mcp_server( async fn spawn_sse_mcp_server(
url: &str, url: &str,
headers: Option<&IndexMap<String, String>>, headers: Option<&IndexMap<String, String>>,
@@ -1048,6 +1222,92 @@ mod tests {
assert_eq!(result["X-Custom"], "keep"); assert_eq!(result["X-Custom"], "keep");
} }
#[test]
fn merge_bearer_token_replaces_authorization_case_insensitively() {
let mut h = IndexMap::new();
h.insert("authorization".to_string(), "Bearer stale-1".to_string());
h.insert("AUTHORIZATION".to_string(), "Bearer stale-2".to_string());
h.insert("X-Custom".to_string(), "keep".to_string());
let result = merge_bearer_token(Some(&h), Some("newtoken".to_string())).unwrap();
assert_eq!(result.len(), 2);
assert_eq!(result["Authorization"], "Bearer newtoken");
assert_eq!(result["X-Custom"], "keep");
assert!(!result.contains_key("authorization"));
assert!(!result.contains_key("AUTHORIZATION"));
}
#[test]
fn http_auth_from_token_status_maps_token_to_managed() {
assert!(matches!(
HttpAuth::from_token_status(&oauth::McpTokenStatus::Token("tok".into()), "srv"),
HttpAuth::Managed { server, token } if server == "srv" && token == "tok"
));
assert!(matches!(
HttpAuth::from_token_status(&oauth::McpTokenStatus::NotAuthenticated, "srv"),
HttpAuth::StaticOnly
));
assert!(matches!(
HttpAuth::from_token_status(&oauth::McpTokenStatus::RefreshFailed, "srv"),
HttpAuth::StaticOnly
));
}
#[test]
fn http_auth_debug_redacts_token() {
let auth = HttpAuth::Managed {
server: "srv".into(),
token: "live-secret".into(),
};
let debug = format!("{auth:?}");
assert!(debug.contains("srv"));
assert!(debug.contains("<redacted>"));
assert!(!debug.contains("live-secret"));
}
#[test]
fn oauth_custom_headers_strips_authorization_case_insensitively() {
let mut h = IndexMap::new();
h.insert("Authorization".to_string(), "Bearer stale-1".to_string());
h.insert("authorization".to_string(), "Bearer stale-2".to_string());
h.insert("AUTHORIZATION".to_string(), "Bearer stale-3".to_string());
h.insert("X-Custom".to_string(), "keep".to_string());
let custom = oauth_custom_headers(Some(&h)).unwrap();
assert_eq!(custom.len(), 1);
assert_eq!(custom[&HeaderName::from_static("x-custom")], "keep");
}
#[test]
fn oauth_custom_headers_none_is_empty() {
assert!(oauth_custom_headers(None).unwrap().is_empty());
}
#[test]
fn oauth_custom_headers_rejects_invalid_header_name() {
let mut h = IndexMap::new();
h.insert("bad header".to_string(), "v".to_string());
assert!(oauth_custom_headers(Some(&h)).is_err());
}
#[test]
fn oauth_custom_headers_keeps_non_authorization_headers() {
let mut h = IndexMap::new();
h.insert("X-Api-Key".to_string(), "k".to_string());
h.insert("X-Trace".to_string(), "t".to_string());
let custom = oauth_custom_headers(Some(&h)).unwrap();
assert_eq!(custom.len(), 2);
assert_eq!(custom[&HeaderName::from_static("x-api-key")], "k");
assert_eq!(custom[&HeaderName::from_static("x-trace")], "t");
}
#[test] #[test]
fn is_auth_required_error_matches_rmcp_message() { fn is_auth_required_error_matches_rmcp_message() {
let e = anyhow!("Auth required, when send initialize request"); let e = anyhow!("Auth required, when send initialize request");
@@ -1071,4 +1331,48 @@ mod tests {
assert!(is_auth_required_error(&e)); assert!(is_auth_required_error(&e));
} }
#[test]
fn auth_reason_maps_token_status() {
assert_eq!(
McpAuthReason::from_token_status(&oauth::McpTokenStatus::Token("tok".into())),
McpAuthReason::TokenRejected
);
assert_eq!(
McpAuthReason::from_token_status(&oauth::McpTokenStatus::NotAuthenticated),
McpAuthReason::NotAuthenticated
);
assert_eq!(
McpAuthReason::from_token_status(&oauth::McpTokenStatus::RefreshFailed),
McpAuthReason::RefreshFailed
);
}
#[test]
fn mcp_auth_required_context_downcasts_with_reason() {
let e = anyhow!("Auth required, when send initialize request").context(McpAuthRequired {
server: "github".into(),
reason: McpAuthReason::RefreshFailed,
});
assert!(is_auth_required_error(&e));
let ctx = e.downcast_ref::<McpAuthRequired>().unwrap();
assert_eq!(ctx.server, "github");
assert_eq!(ctx.reason, McpAuthReason::RefreshFailed);
}
#[test]
fn mcp_auth_required_display_is_reason_specific() {
let msg = |reason| {
McpAuthRequired {
server: "github".into(),
reason,
}
.to_string()
};
assert!(msg(McpAuthReason::NotAuthenticated).contains("no stored credentials"));
assert!(msg(McpAuthReason::RefreshFailed).contains("expired and automatic refresh failed"));
assert!(msg(McpAuthReason::TokenRejected).contains("rejected the stored OAuth token"));
}
} }
+532 -33
View File
@@ -1,15 +1,26 @@
use crate::client::oauth::{OAuthProvider, TokenRequestFormat, load_oauth_tokens, run_oauth_flow}; use crate::client::oauth::{
OAuthProvider, OAuthTokens, TokenRequestFormat, load_oauth_tokens, refresh_oauth_token,
run_oauth_flow, token_response_keys,
};
use crate::config::paths; use crate::config::paths;
use anyhow::{Context, Result, anyhow}; use anyhow::{Context, Result, anyhow};
use chrono::Utc; use chrono::Utc;
use inquire::Text; use inquire::Text;
use log::warn; use log::{debug, warn};
use reqwest::Client; use reqwest::Client;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt;
use std::fs; use std::fs;
use std::net::TcpListener; use std::net::TcpListener;
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};
use tokio::sync;
use url::Url; use url::Url;
const REFRESH_HTTP_TIMEOUT: Duration = Duration::from_secs(10);
const REFRESH_FAILURE_BACKOFF: Duration = Duration::from_secs(60);
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
struct ProtectedResourceMetadata { struct ProtectedResourceMetadata {
#[serde(default)] #[serde(default)]
@@ -34,6 +45,10 @@ struct McpRegistration {
client_id: String, client_id: String,
#[serde(default)] #[serde(default)]
redirect_uri: Option<String>, redirect_uri: Option<String>,
#[serde(default)]
token_url: Option<String>,
#[serde(default)]
resource: Option<String>,
} }
struct DiscoveredOAuth { struct DiscoveredOAuth {
@@ -124,8 +139,19 @@ pub async fn run_mcp_oauth_flow(
None None
}; };
let (client_id, redirect_uri) = if let Some(reused) = cached_reuse { let (client_id, redirect_uri) = if let Some((client_id, redirect_uri)) = cached_reuse {
reused // Re-save so registrations cached before token_url/resource were
// persisted gain them, enabling token refresh next time.
if let Err(e) = save_registration(
server_name,
&client_id,
&redirect_uri,
&metadata.token_endpoint,
&resource,
) {
debug!("Failed to update cached MCP registration for '{server_name}': {e}");
}
(client_id, redirect_uri)
} else { } else {
let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0)); let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0));
let listener = TcpListener::bind(&bind_addr)?; let listener = TcpListener::bind(&bind_addr)?;
@@ -137,10 +163,7 @@ pub async fn run_mcp_oauth_flow(
id.to_string() id.to_string()
} else if let Some(reg_endpoint) = &metadata.registration_endpoint { } else if let Some(reg_endpoint) = &metadata.registration_endpoint {
match register_client(reg_endpoint, &redirect_uri).await { match register_client(reg_endpoint, &redirect_uri).await {
Ok(id) => { Ok(id) => id,
let _ = save_registration(server_name, &id, &redirect_uri);
id
}
Err(e) => { Err(e) => {
warn!("Dynamic client registration failed: {e}. Falling back to manual entry."); warn!("Dynamic client registration failed: {e}. Falling back to manual entry.");
Text::new("Enter the OAuth client ID for this MCP server:") Text::new("Enter the OAuth client ID for this MCP server:")
@@ -153,6 +176,18 @@ pub async fn run_mcp_oauth_flow(
.prompt() .prompt()
.context("Failed to read client ID")? .context("Failed to read client ID")?
}; };
// Persist regardless of how the client_id was obtained (DCR, config,
// or manual entry) so refresh_mcp_token can run the refresh_token
// grant later without interactive re-auth.
if let Err(e) = save_registration(
server_name,
&client_id,
&redirect_uri,
&metadata.token_endpoint,
&resource,
) {
debug!("Failed to cache MCP registration for '{server_name}': {e}");
}
(client_id, redirect_uri) (client_id, redirect_uri)
}; };
@@ -168,12 +203,164 @@ pub async fn run_mcp_oauth_flow(
run_oauth_flow(&provider, &mcp_token_key(server_name)).await run_oauth_flow(&provider, &mcp_token_key(server_name)).await
} }
pub fn load_valid_mcp_token(server_name: &str) -> Option<String> { #[derive(PartialEq, Eq)]
let tokens = load_oauth_tokens(&mcp_token_key(server_name))?; pub enum McpTokenStatus {
if Utc::now().timestamp() < tokens.expires_at { Token(String),
Some(tokens.access_token) NotAuthenticated,
} else { RefreshFailed,
None }
impl fmt::Debug for McpTokenStatus {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Token(_) => f.write_str("Token(<redacted>)"),
Self::NotAuthenticated => f.write_str("NotAuthenticated"),
Self::RefreshFailed => f.write_str("RefreshFailed"),
}
}
}
impl McpTokenStatus {
pub fn into_token(self) -> Option<String> {
match self {
Self::Token(token) => Some(token),
Self::NotAuthenticated | Self::RefreshFailed => None,
}
}
}
pub async fn load_or_refresh_mcp_token(server_name: &str) -> McpTokenStatus {
load_or_refresh_inner(server_name, None).await
}
/// Re-acquires a token after the server rejected the current one mid-session
/// (HTTP 401). The rejection proves the stored token is bad regardless of its
/// expiry timestamp, so the unexpired fast-paths only short-circuit when the
/// stored token DIFFERS from `rejected_token` (a concurrent caller genuinely
/// refreshed while we waited); an unexpired copy of the rejected token is
/// refreshed anyway. The failure backoff and per-server single-flight lock
/// still apply.
pub async fn force_refresh_mcp_token(server_name: &str, rejected_token: &str) -> Option<String> {
load_or_refresh_inner(server_name, Some(rejected_token))
.await
.into_token()
}
async fn load_or_refresh_inner(server_name: &str, rejected_token: Option<&str>) -> McpTokenStatus {
let key = mcp_token_key(server_name);
let Some(tokens) = load_oauth_tokens(&key) else {
return McpTokenStatus::NotAuthenticated;
};
if rejected_token.is_none() && Utc::now().timestamp() < tokens.expires_at {
return McpTokenStatus::Token(tokens.access_token);
}
if in_refresh_failure_backoff(server_name) {
debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed");
return McpTokenStatus::RefreshFailed;
}
let lock = refresh_lock(server_name);
let _guard = lock.lock().await;
// A concurrent caller may have refreshed while we waited for the lock. An
// unexpired token is only trusted if it differs from the rejected one:
// the server already proved that exact token bad.
let Some(tokens) = load_oauth_tokens(&key) else {
return McpTokenStatus::NotAuthenticated;
};
if Utc::now().timestamp() < tokens.expires_at
&& rejected_token.is_none_or(|rejected| rejected != tokens.access_token)
{
return McpTokenStatus::Token(tokens.access_token);
}
if in_refresh_failure_backoff(server_name) {
debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed");
return McpTokenStatus::RefreshFailed;
}
match refresh_mcp_token(server_name, &key, &tokens).await {
Ok(access_token) => McpTokenStatus::Token(access_token),
Err(e) => {
note_refresh_failure(server_name);
warn!(
"Failed to refresh OAuth token for MCP server '{server_name}'. \
Run `.mcp auth {server_name}` to re-authenticate."
);
debug!(
"Token refresh error for MCP server '{server_name}': {}",
redact_refresh_error(&e)
);
McpTokenStatus::RefreshFailed
}
}
}
async fn refresh_mcp_token(server_name: &str, key: &str, tokens: &OAuthTokens) -> Result<String> {
if tokens.refresh_token.is_none() {
return Err(anyhow!("no refresh token stored"));
}
let reg =
load_registration(server_name).ok_or_else(|| anyhow!("no cached client registration"))?;
let token_url = reg.token_url.ok_or_else(|| {
anyhow!("cached registration has no token URL (saved by an older version)")
})?;
let resource = reg.resource.ok_or_else(|| {
anyhow!("cached registration has no resource (saved by an older version)")
})?;
let provider = McpOAuthProvider {
client_id: reg.client_id,
authorize_url: String::new(),
token_url,
scopes: String::new(),
fixed_redirect: String::new(),
resource,
};
let client = Client::builder().timeout(REFRESH_HTTP_TIMEOUT).build()?;
let refreshed = refresh_oauth_token(&client, &provider, key, tokens).await?;
Ok(refreshed.access_token)
}
fn refresh_lock(server_name: &str) -> Arc<sync::Mutex<()>> {
static LOCKS: OnceLock<parking_lot::Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
OnceLock::new();
LOCKS
.get_or_init(Default::default)
.lock()
.entry(server_name.to_string())
.or_default()
.clone()
}
fn refresh_failures() -> &'static parking_lot::Mutex<HashMap<String, Instant>> {
static FAILURES: OnceLock<parking_lot::Mutex<HashMap<String, Instant>>> = OnceLock::new();
FAILURES.get_or_init(Default::default)
}
fn note_refresh_failure(server_name: &str) {
refresh_failures()
.lock()
.insert(server_name.to_string(), Instant::now());
}
fn in_refresh_failure_backoff(server_name: &str) -> bool {
refresh_failures()
.lock()
.get(server_name)
.is_some_and(|failed_at| failed_at.elapsed() < REFRESH_FAILURE_BACKOFF)
}
/// Refresh errors may embed the token endpoint's JSON response, which can
/// contain live tokens; strip everything from the first `{` before logging.
fn redact_refresh_error(e: &anyhow::Error) -> String {
let msg = e.to_string();
match msg.find('{') {
Some(idx) => format!("{}<response body redacted>", &msg[..idx]),
None => msg,
} }
} }
@@ -187,7 +374,13 @@ fn load_registration(server_name: &str) -> Option<McpRegistration> {
serde_json::from_str(&content).ok() serde_json::from_str(&content).ok()
} }
fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) -> Result<()> { fn save_registration(
server_name: &str,
client_id: &str,
redirect_uri: &str,
token_url: &str,
resource: &str,
) -> Result<()> {
let dir = paths::oauth_tokens_dir(); let dir = paths::oauth_tokens_dir();
fs::create_dir_all(&dir)?; fs::create_dir_all(&dir)?;
@@ -195,6 +388,8 @@ fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) ->
let reg = McpRegistration { let reg = McpRegistration {
client_id: client_id.to_string(), client_id: client_id.to_string(),
redirect_uri: Some(redirect_uri.to_string()), redirect_uri: Some(redirect_uri.to_string()),
token_url: Some(token_url.to_string()),
resource: Some(resource.to_string()),
}; };
fs::write(path, serde_json::to_string_pretty(&reg)?)?; fs::write(path, serde_json::to_string_pretty(&reg)?)?;
@@ -244,7 +439,12 @@ async fn register_client(endpoint: &str, redirect_uri: &str) -> Result<String> {
response["client_id"] response["client_id"]
.as_str() .as_str()
.ok_or_else(|| anyhow!("Missing client_id in registration response: {response}")) .ok_or_else(|| {
anyhow!(
"Missing client_id in registration response (keys: {})",
token_response_keys(&response)
)
})
.map(|s| s.to_string()) .map(|s| s.to_string())
} }
@@ -421,16 +621,34 @@ fn extract_base_url(url: &str) -> Result<String> {
} }
#[cfg(test)] #[cfg(test)]
mod tests { pub(crate) mod test_support {
use super::*;
use crate::utils::get_env_name; use crate::utils::get_env_name;
use serial_test::serial;
use std::{ use std::{
env, fs, env,
ffi::OsString,
fs,
path::PathBuf,
time::{self, SystemTime}, time::{self, SystemTime},
}; };
fn with_temp_cache<F: FnOnce()>(f: F) { pub(crate) fn with_temp_cache<F: FnOnce()>(f: F) {
struct Restore {
key: String,
prev: Option<OsString>,
root: PathBuf,
}
impl Drop for Restore {
fn drop(&mut self) {
unsafe {
match self.prev.take() {
Some(v) => env::set_var(&self.key, v),
None => env::remove_var(&self.key),
}
}
let _ = fs::remove_dir_all(&self.root);
}
}
let unique = SystemTime::now() let unique = SystemTime::now()
.duration_since(time::UNIX_EPOCH) .duration_since(time::UNIX_EPOCH)
.unwrap() .unwrap()
@@ -442,15 +660,21 @@ mod tests {
unsafe { unsafe {
env::set_var(&env_key, &root); env::set_var(&env_key, &root);
} }
let _restore = Restore {
key: env_key,
prev,
root,
};
f(); f();
unsafe {
match prev {
Some(v) => env::set_var(&env_key, v),
None => env::remove_var(&env_key),
}
}
let _ = fs::remove_dir_all(&root);
} }
}
#[cfg(test)]
mod tests {
use super::test_support::with_temp_cache;
use super::*;
use serial_test::serial;
use std::fs;
#[test] #[test]
fn extract_base_url_strips_path_and_query() { fn extract_base_url_strips_path_and_query() {
@@ -685,12 +909,19 @@ mod tests {
"notion", "notion",
"client-xyz-123", "client-xyz-123",
"http://127.0.0.1:49152/callback", "http://127.0.0.1:49152/callback",
"https://as.example/token",
"https://mcp.example/mcp",
) )
.unwrap(); .unwrap();
let loaded = load_registration("notion"); let loaded = load_registration("notion").unwrap();
assert_eq!(loaded.unwrap().client_id, "client-xyz-123"); assert_eq!(loaded.client_id, "client-xyz-123");
assert_eq!(
loaded.token_url.as_deref(),
Some("https://as.example/token")
);
assert_eq!(loaded.resource.as_deref(), Some("https://mcp.example/mcp"));
}); });
} }
@@ -708,8 +939,22 @@ mod tests {
#[serial] #[serial]
fn registration_second_save_overwrites_first() { fn registration_second_save_overwrites_first() {
with_temp_cache(|| { with_temp_cache(|| {
save_registration("github", "first-id", "http://127.0.0.1:49152/callback").unwrap(); save_registration(
save_registration("github", "second-id", "http://127.0.0.1:49153/callback").unwrap(); "github",
"first-id",
"http://127.0.0.1:49152/callback",
"https://as.example/token",
"https://mcp.example/mcp",
)
.unwrap();
save_registration(
"github",
"second-id",
"http://127.0.0.1:49153/callback",
"https://as.example/token",
"https://mcp.example/mcp",
)
.unwrap();
let loaded = load_registration("github").unwrap(); let loaded = load_registration("github").unwrap();
@@ -737,6 +982,8 @@ mod tests {
assert_eq!(loaded.client_id, "legacy-id"); assert_eq!(loaded.client_id, "legacy-id");
assert_eq!(loaded.redirect_uri, None); assert_eq!(loaded.redirect_uri, None);
assert_eq!(loaded.token_url, None);
assert_eq!(loaded.resource, None);
}); });
} }
@@ -744,7 +991,14 @@ mod tests {
#[serial] #[serial]
fn save_registration_persists_redirect_uri() { fn save_registration_persists_redirect_uri() {
with_temp_cache(|| { with_temp_cache(|| {
save_registration("aws", "client-abc", "http://127.0.0.1:49152/callback").unwrap(); save_registration(
"aws",
"client-abc",
"http://127.0.0.1:49152/callback",
"https://as.example/token",
"https://mcp.example/mcp",
)
.unwrap();
let loaded = load_registration("aws").unwrap(); let loaded = load_registration("aws").unwrap();
@@ -756,6 +1010,251 @@ mod tests {
}); });
} }
#[test]
fn mcp_registration_deserializes_without_new_fields_and_roundtrips() {
let old: McpRegistration = serde_json::from_str(r#"{"client_id":"legacy-id"}"#).unwrap();
assert_eq!(old.client_id, "legacy-id");
assert_eq!(old.token_url, None);
assert_eq!(old.resource, None);
let full = McpRegistration {
client_id: "client-abc".into(),
redirect_uri: Some("http://127.0.0.1:49152/callback".into()),
token_url: Some("https://as.example/token".into()),
resource: Some("https://mcp.example/mcp".into()),
};
let json = serde_json::to_string(&full).unwrap();
let back: McpRegistration = serde_json::from_str(&json).unwrap();
assert_eq!(back.token_url.as_deref(), Some("https://as.example/token"));
assert_eq!(back.resource.as_deref(), Some("https://mcp.example/mcp"));
}
#[test]
#[serial]
fn expired_token_with_old_format_registration_reports_refresh_failed() {
with_temp_cache(|| {
let dir = paths::oauth_tokens_dir();
fs::create_dir_all(&dir).unwrap();
fs::write(
paths::token_file("mcp_legacyref"),
r#"{"access_token":"stale","refresh_token":"refresh-abc","expires_at":0}"#,
)
.unwrap();
fs::write(
dir.join("mcp_legacyref_registration.json"),
r#"{"client_id":"legacy-id"}"#,
)
.unwrap();
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let status = rt.block_on(load_or_refresh_mcp_token("legacyref"));
assert_eq!(status, McpTokenStatus::RefreshFailed);
});
}
#[test]
#[serial]
fn missing_token_file_reports_not_authenticated() {
with_temp_cache(|| {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let status = rt.block_on(load_or_refresh_mcp_token("never-authed"));
assert_eq!(status, McpTokenStatus::NotAuthenticated);
});
}
#[test]
#[serial]
fn force_refresh_returns_concurrently_refreshed_token() {
with_temp_cache(|| {
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
fs::write(
paths::token_file("mcp_force-fresh"),
r#"{"access_token":"fresh-tok","refresh_token":"r","expires_at":9999999999}"#,
)
.unwrap();
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let token = rt.block_on(force_refresh_mcp_token("force-fresh", "rejected-tok"));
assert_eq!(token.as_deref(), Some("fresh-tok"));
});
}
#[test]
#[serial]
fn force_refresh_unexpired_rejected_token_attempts_real_refresh() {
with_temp_cache(|| {
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
fs::write(
paths::token_file("mcp_force-rejected"),
r#"{"access_token":"same-tok","refresh_token":"r","expires_at":9999999999}"#,
)
.unwrap();
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let token = rt.block_on(force_refresh_mcp_token("force-rejected", "same-tok"));
assert_eq!(token, None);
});
}
#[test]
#[serial]
fn force_refresh_missing_token_file_returns_none() {
with_temp_cache(|| {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let token = rt.block_on(force_refresh_mcp_token(
"force-never-authed",
"rejected-tok",
));
assert_eq!(token, None);
});
}
#[test]
#[serial]
fn force_refresh_failed_refresh_returns_none() {
with_temp_cache(|| {
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
fs::write(
paths::token_file("mcp_force-fail"),
r#"{"access_token":"stale","refresh_token":"r","expires_at":0}"#,
)
.unwrap();
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let token = rt.block_on(force_refresh_mcp_token("force-fail", "stale"));
assert_eq!(token, None);
});
}
#[test]
#[serial]
fn force_refresh_concurrent_callers_complete_without_deadlock() {
with_temp_cache(|| {
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
fs::write(
paths::token_file("mcp_force-concurrent"),
r#"{"access_token":"same-tok","refresh_token":"r","expires_at":9999999999}"#,
)
.unwrap();
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let (a, b) = rt.block_on(async {
tokio::join!(
force_refresh_mcp_token("force-concurrent", "same-tok"),
force_refresh_mcp_token("force-concurrent", "same-tok"),
)
});
assert_eq!(a, None);
assert_eq!(b, None);
});
}
#[test]
#[serial]
fn force_refresh_respects_failure_backoff() {
with_temp_cache(|| {
fs::create_dir_all(paths::oauth_tokens_dir()).unwrap();
fs::write(
paths::token_file("mcp_force-backoff"),
r#"{"access_token":"same-tok","refresh_token":"r","expires_at":9999999999}"#,
)
.unwrap();
note_refresh_failure("force-backoff");
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let token = rt.block_on(force_refresh_mcp_token("force-backoff", "same-tok"));
assert_eq!(token, None);
});
}
#[test]
fn token_status_debug_redacts_token() {
assert_eq!(
format!("{:?}", McpTokenStatus::Token("live-secret".into())),
"Token(<redacted>)"
);
assert_eq!(
format!("{:?}", McpTokenStatus::NotAuthenticated),
"NotAuthenticated"
);
assert_eq!(
format!("{:?}", McpTokenStatus::RefreshFailed),
"RefreshFailed"
);
}
#[test]
fn token_status_into_token_extracts_only_token_variant() {
assert_eq!(
McpTokenStatus::Token("tok".into()).into_token(),
Some("tok".to_string())
);
assert_eq!(McpTokenStatus::NotAuthenticated.into_token(), None);
assert_eq!(McpTokenStatus::RefreshFailed.into_token(), None);
}
#[test]
fn refresh_failure_backoff_memoizes_per_server() {
assert!(!in_refresh_failure_backoff("backoff-test-server"));
note_refresh_failure("backoff-test-server");
assert!(in_refresh_failure_backoff("backoff-test-server"));
assert!(!in_refresh_failure_backoff("backoff-other-server"));
}
#[test]
fn redact_refresh_error_strips_response_body() {
let with_body = anyhow!(
"Missing access_token in refresh response: {}",
r#"{"access_token":"live-secret"}"#
);
let without_body = anyhow!("no refresh token stored");
assert_eq!(
redact_refresh_error(&with_body),
"Missing access_token in refresh response: <response body redacted>"
);
assert_eq!(
redact_refresh_error(&without_body),
"no refresh token stored"
);
}
#[test] #[test]
fn cached_redirect_port_matches() { fn cached_redirect_port_matches() {
let port = cached_redirect_port("http://127.0.0.1:49152/callback", "127.0.0.1", None); let port = cached_redirect_port("http://127.0.0.1:49152/callback", "127.0.0.1", None);
+68 -3
View File
@@ -1,11 +1,11 @@
use crate::function::{FunctionDeclaration, JsonSchema}; use crate::function::{self, FunctionDeclaration, JsonSchema};
use anyhow::{Context, Result, bail}; use anyhow::{Context, Result, bail};
use argc::{ChoiceValue, CommandValue, FlagOptionValue}; use argc::{ChoiceValue, CommandValue, FlagOptionValue};
use indexmap::IndexMap; use indexmap::IndexMap;
use std::env;
use std::fs::File; use std::fs::File;
use std::io::Read; use std::io::Read;
use std::path::Path; use std::path::Path;
use std::{env, fs};
pub fn generate_bash_declarations( pub fn generate_bash_declarations(
mut tool_file: File, mut tool_file: File,
@@ -23,7 +23,8 @@ pub fn generate_bash_declarations(
"", "",
env::var("TERM_WIDTH").ok().and_then(|v| v.parse().ok()), env::var("TERM_WIDTH").ok().and_then(|v| v.parse().ok()),
)?; )?;
fs::write(tools_file_path, &build_script) let build_script = allow_empty_required_values(&build_script);
function::write_file_atomic(tools_file_path, &build_script, Some(0o755))
.with_context(|| format!("Failed to write built script to '{tools_file_path:?}'"))?; .with_context(|| format!("Failed to write built script to '{tools_file_path:?}'"))?;
let command_value = argc::export(&build_script, file_name) let command_value = argc::export(&build_script, file_name)
@@ -74,6 +75,20 @@ fn underscore(s: &str) -> String {
s.replace('-', "_") s.replace('-', "_")
} }
/// argc's generated required-param check uses `-z "${!name:-}"`, which
/// conflates "not provided" with "provided but empty", so a required option
/// passed an explicit empty string (e.g. `fs_write --content=''` to create an
/// empty file) is rejected as "required arguments were not provided". The
/// JSON schema we advertise to models treats `required` as *presence*, so
/// rewrite the check to a set-ness test to keep runtime behavior consistent
/// with the schema. Applied post-build so it survives every regeneration.
fn allow_empty_required_values(build_script: &str) -> String {
build_script.replace(
r#"if [[ -z "${!name:-}" ]]; then"#,
r#"if [[ -z "${!name+x}" ]]; then"#,
)
}
fn schema_ty(t: &str) -> JsonSchema { fn schema_ty(t: &str) -> JsonSchema {
JsonSchema { JsonSchema {
type_value: Some(t.to_string()), type_value: Some(t.to_string()),
@@ -147,3 +162,53 @@ fn parse_parameters_schema(flags: &[FlagOptionValue]) -> JsonSchema {
required: Some(required), required: Some(required),
} }
} }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rewrites_argc_required_check_to_setness_test() {
let src = "# @describe test tool\n# @option --content! The content\nmain() { :; }\n";
let built = argc::build(src, "", None).expect("argc build failed");
assert!(
built.contains(r#"if [[ -z "${!name:-}" ]]; then"#),
"argc changed its generated required-param template; update allow_empty_required_values()"
);
let fixed = allow_empty_required_values(&built);
assert!(fixed.contains(r#"if [[ -z "${!name+x}" ]]; then"#));
assert!(!fixed.contains(r#"if [[ -z "${!name:-}" ]]; then"#));
}
#[cfg(unix)]
#[test]
fn second_build_does_not_rewrite_tool_file_test() {
use std::fs;
use std::os::unix::fs::{MetadataExt, PermissionsExt};
let src = "# @describe test tool\n# @option --content! The content\nmain() { :; }\n";
let path = crate::utils::temp_file("bash-build", ".sh");
fs::write(&path, src).expect("failed to write temp script");
let declarations = generate_bash_declarations(File::open(&path).unwrap(), &path, "test")
.expect("first build failed");
assert!(!declarations.is_empty());
let first_ino = fs::metadata(&path).unwrap().ino();
let declarations = generate_bash_declarations(File::open(&path).unwrap(), &path, "test")
.expect("second build failed");
assert!(!declarations.is_empty());
let metadata = fs::metadata(&path).unwrap();
assert_eq!(
metadata.ino(),
first_ino,
"second build rewrote an already-built tool file"
);
assert_eq!(metadata.permissions().mode() & 0o777, 0o755);
fs::remove_file(&path).unwrap();
}
}
+110 -71
View File
@@ -151,7 +151,7 @@ impl Rag {
println!("⚙ Initializing RAG..."); println!("⚙ Initializing RAG...");
let mut data = Self::resolve_init_data(app, config)?; let mut data = Self::resolve_init_data(app, config)?;
data.driver = config.driver.clone().unwrap_or_else(|| "yaml".to_string()); data.driver = config.driver.clone().unwrap_or_else(|| "yaml".to_string());
let mut rag = Self::create(app, name, save_path, data)?; let mut rag = Self::create(app, name, save_path, data).await?;
let loaders = app.document_loaders.clone(); let loaders = app.document_loaders.clone();
let (spinner, spinner_rx) = Spinner::create(""); let (spinner, spinner_rx) = Spinner::create("");
abortable_run_with_spinner_rx( abortable_run_with_spinner_rx(
@@ -267,31 +267,10 @@ impl Rag {
} }
println!("⚙ Initializing RAG..."); println!("⚙ Initializing RAG...");
let (embedding_model, chunk_size, chunk_overlap) = Self::create_config(app)?; let (embedding_model, chunk_size, chunk_overlap) = Self::create_config(app)?;
// Only interactive named-RAG creation offers a driver choice. Temp RAGs and
// agent startup pass `false`; an explicit flag is used rather than inferring
// from the name because the agent path passes the literal name "rag", which is
// indistinguishable from a user creating a RAG genuinely named `rag`.
let driver = if prompt_for_driver { let driver = if prompt_for_driver {
let options = vec![ select_rag_driver()?
"yaml — portable, in-memory HNSW; usable from several Coyote processes at once (default)",
"duckdb — persistent on-disk store; vectors and content survive restarts; HNSW approximate search.",
];
let sel = Select::new("RAG storage driver:", options)
.with_starting_cursor(0)
.prompt()?;
if sel.starts_with("duckdb") {
println!(
"Note: several Coyote processes can query a duckdb RAG at the same time, \
but while one process is ingesting or rebuilding it the others cannot \
read it until that finishes. Changing its driver later means deleting \
and recreating the RAG."
);
"duckdb"
} else {
"yaml"
}
} else { } else {
"yaml" "yaml".to_string()
}; };
let reranker_model = app.rag_reranker_model.clone(); let reranker_model = app.rag_reranker_model.clone();
let top_k = app.rag_top_k; let top_k = app.rag_top_k;
@@ -318,8 +297,8 @@ impl Rag {
graph_hops: Some(graph_hops), graph_hops: Some(graph_hops),
}, },
); );
data.driver = driver.to_string(); data.driver = driver;
let mut rag = Self::create(app, name, save_path, data)?; let mut rag = Self::create(app, name, save_path, data).await?;
let mut paths = doc_paths.to_vec(); let mut paths = doc_paths.to_vec();
if paths.is_empty() { if paths.is_empty() {
paths = add_documents()?; paths = add_documents()?;
@@ -338,12 +317,12 @@ impl Rag {
Ok(rag) Ok(rag)
} }
pub fn load(app: &AppConfig, name: &str, path: &Path) -> Result<Self> { pub async fn load(app: &AppConfig, name: &str, path: &Path) -> Result<Self> {
let err = || format!("Failed to load rag '{name}' at '{}'", path.display()); let err = || format!("Failed to load rag '{name}' at '{}'", path.display());
let content = fs::read_to_string(path).with_context(err)?; let content = fs::read_to_string(path).with_context(err)?;
let data: RagData = serde_yaml::from_str(&content).with_context(err)?; let data: RagData = serde_yaml::from_str(&content).with_context(err)?;
data.validate().with_context(err)?; data.validate().with_context(err)?;
Self::create(app, name, path, data) Self::create(app, name, path, data).await
} }
/// Loads a RAG from a YAML file. External drivers need an async constructor /// Loads a RAG from a YAML file. External drivers need an async constructor
@@ -393,7 +372,7 @@ impl Rag {
last_sources: RwLock::new(None), last_sources: RwLock::new(None),
}) })
} }
_ => Self::load(app, name, path), _ => Self::load(app, name, path).await,
} }
} }
@@ -558,14 +537,22 @@ impl Rag {
Ok(rag) Ok(rag)
} }
pub fn create(app: &AppConfig, name: &str, path: &Path, mut data: RagData) -> Result<Self> { pub async fn create(
app: &AppConfig,
name: &str,
path: &Path,
mut data: RagData,
) -> Result<Self> {
// Deliberately does NOT call rebuild_indexes: both callers construct the Rag // Deliberately does NOT call rebuild_indexes: both callers construct the Rag
// before any documents are added, so rebuilding empty data would be a no-op. // before any documents are added, so rebuilding empty data would be a no-op.
// Actual population happens later via sync_documents. // Actual population happens later via sync_documents.
let (provider, bm25): (Box<dyn RagProvider>, _) = match data.driver.as_str() { let (provider, bm25): (Box<dyn RagProvider>, _) = match data.driver.as_str() {
"duckdb" => { "duckdb" => {
let db_path = providers::duckdb_path_from_yaml(path); let db_path = providers::duckdb_path_from_yaml(path);
let dim = embedding_dim_for_model(&data.embedding_model); let dim = match DuckDbProvider::introspect_dim(&db_path)? {
Some(existing) => existing,
None => probe_embedding_dim(app, &data.embedding_model).await?,
};
let duck = DuckDbProvider::open(&db_path, dim)?; let duck = DuckDbProvider::open(&db_path, dim)?;
// HYDRATE — mandatory, not an optimization. The YAML file for a duckdb // HYDRATE — mandatory, not an optimization. The YAML file for a duckdb
// RAG deliberately omits `vectors`, so `data.vectors` arrives empty from // RAG deliberately omits `vectors`, so `data.vectors` arrives empty from
@@ -586,19 +573,40 @@ impl Rag {
if data.vectors.is_empty() { if data.vectors.is_empty() {
data.vectors = duck.read_all_vectors()?; data.vectors = duck.read_all_vectors()?;
} }
if data.vectors.is_empty() && !data.files.is_empty() {
println!(
"{} RAG '{name}' lists {} indexed file(s), but its vector store \
'{}' holds no vectors, so every search will return nothing. A \
duckdb RAG is two files: bring the .duckdb sidecar along with \
the .yaml, or re-embed with `.rebuild rag`.",
warning_text("WARNING:"),
data.files.len(),
db_path.display()
);
}
// data.files is always populated for duckdb, so build_bm25() is the only // data.files is always populated for duckdb, so build_bm25() is the only
// path; there is no from-DuckDB fallback. // path; there is no from-DuckDB fallback.
let bm25 = data.build_bm25(); let bm25 = data.build_bm25();
(Box::new(duck), bm25) (Box::new(duck), bm25)
} }
"qdrant" => bail!( "qdrant" => bail!(
"Qdrant RAGs cannot be constructed via Rag::create(); \ "RAG '{name}' uses driver 'qdrant' without `attached: true`. \
use Rag::attach() or Rag::load_async() instead" Coyote can currently only READ a pre-existing Qdrant \
collection attach one with `.rag attach`. Writing to a \
Coyote-owned Qdrant collection is not supported yet."
), ),
_ => { "yaml" => {
let bm25 = data.build_bm25(); let bm25 = data.build_bm25();
(Box::new(YamlProvider::from_data(&data)), bm25) (Box::new(YamlProvider::from_data(&data)), bm25)
} }
// Explicitly NOT a catch-all falling through to yaml. A typo'd driver
// used to build a yaml store, pay to embed the whole corpus, persist
// the bad driver, and only fail on the NEXT run, leaving the RAG
// unusable without hand-editing the YAML.
other => bail!(
"Unknown RAG driver '{other}' for RAG '{name}'. \
Valid drivers: yaml, duckdb, qdrant."
),
}; };
let node_to_docs = data.knowledge_graph.build_node_to_docs(); let node_to_docs = data.knowledge_graph.build_node_to_docs();
let embedding_model = let embedding_model =
@@ -856,6 +864,20 @@ impl Rag {
Ok((embeddings, sources, ids)) Ok((embeddings, sources, ids))
} }
pub async fn search_chunks(
&self,
text: &str,
top_k: usize,
rerank_model: Option<&str>,
) -> Result<Vec<(String, String)>> {
let results = self.hybrid_search(text, top_k, rerank_model).await?;
Ok(results
.into_iter()
.map(|(id, content)| (content, self.resolve_source(&id)))
.collect())
}
pub async fn search_with_template( pub async fn search_with_template(
&self, &self,
app: &AppConfig, app: &AppConfig,
@@ -1169,12 +1191,7 @@ impl Rag {
top_k: usize, top_k: usize,
rerank_model: Option<&str>, rerank_model: Option<&str>,
) -> Result<Vec<(DocumentId, String)>> { ) -> Result<Vec<(DocumentId, String)>> {
let vector_search_results = self.vector_search(query, top_k, 0.0).await?; let keyword_leg = async {
debug!("vector_search_results: {vector_search_results:?}",);
let vector_search_ids: Vec<DocumentId> =
vector_search_results.into_iter().map(|(v, _)| v).collect();
let keyword_search_results: Vec<(DocumentId, f32)> =
if self.provider.has_native_keyword_search() { if self.provider.has_native_keyword_search() {
self.provider self.provider
.keyword_search(query, top_k) .keyword_search(query, top_k)
@@ -1185,7 +1202,16 @@ impl Rag {
}) })
} else { } else {
self.keyword_search(query, top_k, 0.0) self.keyword_search(query, top_k, 0.0)
}; }
};
let (vector_search_results, keyword_search_results) =
tokio::join!(self.vector_search(query, top_k, 0.0), keyword_leg);
let vector_search_results = vector_search_results?;
debug!("vector_search_results: {vector_search_results:?}",);
let vector_search_ids: Vec<DocumentId> =
vector_search_results.into_iter().map(|(v, _)| v).collect();
debug!("keyword_search_results: {keyword_search_results:?}",); debug!("keyword_search_results: {keyword_search_results:?}",);
let keyword_search_ids: Vec<DocumentId> = let keyword_search_ids: Vec<DocumentId> =
keyword_search_results.into_iter().map(|(v, _)| v).collect(); keyword_search_results.into_iter().map(|(v, _)| v).collect();
@@ -1842,6 +1868,27 @@ fn select_embedding_model(models: &[&Model]) -> Result<String> {
Ok(result.value) Ok(result.value)
} }
pub(crate) fn select_rag_driver() -> Result<String> {
let options = vec![
"yaml — portable, in-memory HNSW; usable from several Coyote processes at once (default)",
"duckdb — persistent on-disk store; vectors and content survive restarts; HNSW approximate search.",
];
let sel = Select::new("RAG storage driver:", options)
.with_starting_cursor(0)
.prompt()?;
if sel.starts_with("duckdb") {
println!(
"Note: several Coyote processes can query a duckdb RAG at the same time, \
but while one process is ingesting or rebuilding it the others cannot \
read it until that finishes. Changing its driver later means deleting \
and recreating the RAG."
);
Ok("duckdb".to_string())
} else {
Ok("yaml".to_string())
}
}
const EXTRACTOR_SKIP: &str = "Skip"; const EXTRACTOR_SKIP: &str = "Skip";
fn select_extractor_model(app: &AppConfig) -> Result<Option<String>> { fn select_extractor_model(app: &AppConfig) -> Result<Option<String>> {
@@ -2080,21 +2127,25 @@ fn reciprocal_rank_fusion(
.collect() .collect()
} }
/// Map an embedding model id to its vector dimension. async fn probe_embedding_dim(app: &AppConfig, model_id: &str) -> Result<usize> {
/// let model = Model::retrieve_model(app, model_id, ModelType::Embedding)?;
/// The DuckDB `FLOAT[N]` column type and its HNSW index are fixed at schema-creation let client = init_client(&Arc::new(app.clone()), model)?;
/// time, so this value must be decided before the first insert. An unrecognized model let out = client
/// falls back to 1536; if that is wrong, DuckDB raises a dimension-mismatch error on .embeddings(&EmbeddingsData::new(vec!["dimension probe".into()], false))
/// the first insert rather than silently corrupting the schema, and the recovery is to .await
/// delete the sidecar and re-ingest from source. .with_context(|| {
fn embedding_dim_for_model(model_id: &str) -> usize { format!(
match model_id { "Failed to probe the embedding dimension of model '{model_id}'. \
m if m.contains("3-large") => 3072, Creating a duckdb RAG requires one call to the embedding endpoint."
m if m.contains("3-small") || m.contains("ada-002") => 1536, )
m if m.contains("nomic-embed-text") || m.contains("all-minilm") => 768, })?;
m if m.contains("jina-embeddings-v2") => 1024, let dim = out.first().map(|v| v.len()).unwrap_or(0);
_ => 1536,
if dim == 0 {
bail!("Embedding model '{model_id}' returned an empty vector during the dimension probe");
} }
Ok(dim)
} }
/// True only for "the vault does not hold this key". /// True only for "the vault does not hold this key".
@@ -2634,18 +2685,6 @@ mod tests {
assert_eq!(data.attached_source_label(), "[external collection]"); assert_eq!(data.attached_source_label(), "[external collection]");
} }
#[test]
fn embedding_dim_for_model_maps_known_models() {
assert_eq!(embedding_dim_for_model("text-embedding-3-large"), 3072);
assert_eq!(embedding_dim_for_model("text-embedding-3-small"), 1536);
assert_eq!(embedding_dim_for_model("text-embedding-ada-002"), 1536);
assert_eq!(embedding_dim_for_model("nomic-embed-text"), 768);
assert_eq!(embedding_dim_for_model("all-minilm"), 768);
assert_eq!(embedding_dim_for_model("jina-embeddings-v2-base-en"), 1024);
// Unknown models fall back to the OpenAI-compatible default.
assert_eq!(embedding_dim_for_model("some-unknown-model"), 1536);
}
#[test] #[test]
fn document_id_round_trip() { fn document_id_round_trip() {
let id = DocumentId::new(5, 17); let id = DocumentId::new(5, 17);
@@ -3005,7 +3044,7 @@ mod tests {
#[test] #[test]
fn reciprocal_rank_fusion_empty_lists() { fn reciprocal_rank_fusion_empty_lists() {
let result = super::reciprocal_rank_fusion(vec![], vec![], 5); let result = reciprocal_rank_fusion(vec![], vec![], 5);
assert!(result.is_empty(), "empty input should produce empty output"); assert!(result.is_empty(), "empty input should produce empty output");
} }
@@ -3013,7 +3052,7 @@ mod tests {
fn reciprocal_rank_fusion_deduplicates_across_signals() { fn reciprocal_rank_fusion_deduplicates_across_signals() {
let doc_a = DocumentId::new(0, 0); let doc_a = DocumentId::new(0, 0);
let doc_b = DocumentId::new(0, 1); let doc_b = DocumentId::new(0, 1);
let result = super::reciprocal_rank_fusion( let result = reciprocal_rank_fusion(
vec![vec![doc_a, doc_b], vec![doc_a, doc_b]], vec![vec![doc_a, doc_b], vec![doc_a, doc_b]],
vec![1.0, 1.0], vec![1.0, 1.0],
5, 5,
@@ -3030,7 +3069,7 @@ mod tests {
#[test] #[test]
fn reciprocal_rank_fusion_respects_top_k() { fn reciprocal_rank_fusion_respects_top_k() {
let docs: Vec<DocumentId> = (0..10).map(|i| DocumentId::new(0, i)).collect(); let docs: Vec<DocumentId> = (0..10).map(|i| DocumentId::new(0, i)).collect();
let result = super::reciprocal_rank_fusion(vec![docs], vec![1.0], 3); let result = reciprocal_rank_fusion(vec![docs], vec![1.0], 3);
assert_eq!(result.len(), 3, "result should be capped at top_k=3"); assert_eq!(result.len(), 3, "result should be capped at top_k=3");
} }
@@ -3038,7 +3077,7 @@ mod tests {
fn reciprocal_rank_fusion_weights_affect_ranking() { fn reciprocal_rank_fusion_weights_affect_ranking() {
let doc_a = DocumentId::new(0, 0); let doc_a = DocumentId::new(0, 0);
let doc_b = DocumentId::new(0, 1); let doc_b = DocumentId::new(0, 1);
let result = super::reciprocal_rank_fusion( let result = reciprocal_rank_fusion(
vec![vec![doc_a, doc_b], vec![doc_b, doc_a]], vec![vec![doc_a, doc_b], vec![doc_b, doc_a]],
vec![10.0, 1.0], vec![10.0, 1.0],
2, 2,
+211 -9
View File
@@ -5,7 +5,7 @@ use std::collections::HashMap;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Result, anyhow, bail};
use async_trait::async_trait; use async_trait::async_trait;
use duckdb::types::Value; use duckdb::types::Value;
use duckdb::{AccessMode, Config, Connection}; use duckdb::{AccessMode, Config, Connection, OptionalExt};
use indexmap::IndexMap; use indexmap::IndexMap;
use log::warn; use log::warn;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
@@ -30,6 +30,11 @@ pub(crate) fn duckdb_path_from_yaml(yaml_path: &Path) -> PathBuf {
yaml_path.with_extension("duckdb") yaml_path.with_extension("duckdb")
} }
fn parse_float_array_dim(data_type: &str) -> Option<usize> {
let inner = data_type.strip_prefix("FLOAT[")?.strip_suffix(']')?;
inner.parse::<usize>().ok().filter(|&n| n > 0)
}
/// The shared connection together with the access mode it was opened with. /// The shared connection together with the access mode it was opened with.
/// ///
/// `conn` is an `Option` only so that an upgrade can DROP the read-only connection /// `conn` is an `Option` only so that an upgrade can DROP the read-only connection
@@ -42,6 +47,11 @@ struct ConnHandle {
/// `duplicate()` clone sharing the `Arc` observes an upgrade performed through any /// `duplicate()` clone sharing the `Arc` observes an upgrade performed through any
/// other handle instead of keeping its own stale copy of the mode. /// other handle instead of keeping its own stale copy of the mode.
writable: bool, writable: bool,
/// Embedding dimension the `FLOAT[N]` column was opened (or last rebuilt) with.
/// Shared for the same reason as `writable`: a self-healing rebuild through one
/// handle updates the width, and every `duplicate()` clone must cast with the new
/// width instead of erroring on a healthy store with its stale copy.
dim: usize,
} }
impl ConnHandle { impl ConnHandle {
@@ -68,8 +78,6 @@ impl ConnHandle {
pub struct DuckDbProvider { pub struct DuckDbProvider {
path: PathBuf, path: PathBuf,
conn: Arc<Mutex<ConnHandle>>, conn: Arc<Mutex<ConnHandle>>,
/// Embedding dimension; fixed at open time because the `FLOAT[N]` column depends on it.
dim: usize,
/// True once an FTS index has been built on `documents`. Until then /// True once an FTS index has been built on `documents`. Until then
/// `fts_main_documents.match_bm25` does not exist and any keyword query would /// `fts_main_documents.match_bm25` does not exist and any keyword query would
/// fail with a DuckDB catalog error. Backs `has_native_keyword_search`. /// fail with a DuckDB catalog error. Backs `has_native_keyword_search`.
@@ -96,8 +104,8 @@ impl DuckDbProvider {
conn: Arc::new(Mutex::new(ConnHandle { conn: Arc::new(Mutex::new(ConnHandle {
conn: Some(conn), conn: Some(conn),
writable, writable,
dim,
})), })),
dim,
fts_ready: AtomicBool::new(fts_exists), fts_ready: AtomicBool::new(fts_exists),
}) })
} }
@@ -164,6 +172,35 @@ impl DuckDbProvider {
Ok(conn) Ok(conn)
} }
pub fn introspect_dim(db_path: &Path) -> Result<Option<usize>> {
if !db_path.exists() {
return Ok(None);
}
let conn = Self::open_read_only(db_path).with_context(|| {
format!(
"Cannot inspect the existing RAG store at '{}'",
db_path.display()
)
})?;
let ty: Option<String> = conn
.query_row(
"SELECT data_type FROM duckdb_columns() \
WHERE table_name = 'vectors' AND column_name = 'embedding'",
[],
|r| r.get(0),
)
.optional()
.with_context(|| {
format!(
"Failed to introspect the embedding dimension of the DuckDB store \
at '{}'",
db_path.display()
)
})?;
Ok(ty.and_then(|t| parse_float_array_dim(&t)))
}
/// Open the store read-write and make sure its schema exists. Exactly one process /// Open the store read-write and make sure its schema exists. Exactly one process
/// may hold such a handle, and no reader from another process may hold it meanwhile. /// may hold such a handle, and no reader from another process may hold it meanwhile.
fn open_read_write(db_path: &Path, dim: usize) -> Result<Connection> { fn open_read_write(db_path: &Path, dim: usize) -> Result<Connection> {
@@ -245,7 +282,7 @@ impl DuckDbProvider {
return Ok(()); return Ok(());
} }
drop(handle.conn.take()); drop(handle.conn.take());
match Self::open_read_write(&self.path, self.dim) { match Self::open_read_write(&self.path, handle.dim) {
Ok(conn) => { Ok(conn) => {
handle.conn = Some(conn); handle.conn = Some(conn);
handle.writable = true; handle.writable = true;
@@ -428,13 +465,33 @@ impl RagProvider for DuckDbProvider {
if embedding.iter().any(|f| !f.is_finite()) { if embedding.iter().any(|f| !f.is_finite()) {
bail!("Query embedding contains a non-finite value (NaN or infinity)"); bail!("Query embedding contains a non-finite value (NaN or infinity)");
} }
let handle = self.lock_conn()?;
let dim = handle.dim;
if embedding.len() != dim {
let rows: i64 = handle
.conn()?
.query_row("SELECT count(*) FROM vectors", [], |r| r.get(0))
.context("Failed to count vectors before a dimension-mismatch query")?;
if rows == 0 {
// A never-synced store answers "nothing", not a cast error.
return Ok(Vec::new());
}
bail!(
"RAG store at '{}' was built with {dim}-dim embeddings, but the \
embedding model now returns {}-dim vectors. The embedding model \
changed since ingestion. Re-embed the documents, or delete the \
sidecar file and re-ingest.",
self.path.display(),
embedding.len()
);
}
let vals: String = embedding let vals: String = embedding
.iter() .iter()
.map(|f| f.to_string()) .map(|f| f.to_string())
.collect::<Vec<_>>() .collect::<Vec<_>>()
.join(", "); .join(", ");
let dim = self.dim;
let handle = self.lock_conn()?;
// array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[]. // array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[].
// ORDER BY distance ASC is required for the planner to use hnsw_idx; the // ORDER BY distance ASC is required for the planner to use hnsw_idx; the
// similarity form (DESC) does NOT trigger the ANN index. Distance is converted // similarity form (DESC) does NOT trigger the ANN index. Distance is converted
@@ -554,7 +611,28 @@ impl RagProvider for DuckDbProvider {
); );
} }
} }
let dim = self.dim; let dim = match data.vectors.first() {
None => self.lock_conn()?.dim,
Some((_, first)) => {
let dim = first.len();
if let Some((doc_id, other)) = data.vectors.iter().find(|(_, e)| e.len() != dim) {
let matching = data.vectors.values().filter(|e| e.len() == dim).count();
bail!(
"Refusing to rebuild the RAG store at '{}': the rebuild batch \
mixes {dim}-dim and {}-dim vectors ({matching} vs {} vectors; \
first mismatch: document {}). Re-embed the documents, or \
delete the sidecar file and re-ingest.",
self.path.display(),
other.len(),
data.vectors.len() - matching,
doc_id.0
);
}
dim
}
};
// THE write path. Everything above this line only reads, so the upgrade happens // THE write path. Everything above this line only reads, so the upgrade happens
// here, after both guards have had their say: a rebuild that is going to be // here, after both guards have had their say: a rebuild that is going to be
// refused must not first take the exclusive lock away from other processes. // refused must not first take the exclusive lock away from other processes.
@@ -664,6 +742,8 @@ impl RagProvider for DuckDbProvider {
// fall back to local BM25, which is also empty, and therefore correct. // fall back to local BM25, which is also empty, and therefore correct.
self.fts_ready.store(doc_count > 0, Ordering::Relaxed); self.fts_ready.store(doc_count > 0, Ordering::Relaxed);
handle.dim = dim;
Ok(()) Ok(())
} }
@@ -729,10 +809,11 @@ impl RagProvider for DuckDbProvider {
// Sharing the Arc also shares the ACCESS MODE, which lives inside the ConnHandle // Sharing the Arc also shares the ACCESS MODE, which lives inside the ConnHandle
// rather than beside it: when one handle upgrades itself to read-write, every // rather than beside it: when one handle upgrades itself to read-write, every
// clone is upgraded with it and none is left holding a stale "read-only" belief. // clone is upgraded with it and none is left holding a stale "read-only" belief.
// The embedding dimension lives there too, so a self-healing rebuild through
// one handle updates the width every clone casts with.
Box::new(DuckDbProvider { Box::new(DuckDbProvider {
path: self.path.clone(), path: self.path.clone(),
conn: Arc::clone(&self.conn), conn: Arc::clone(&self.conn),
dim: self.dim,
fts_ready: AtomicBool::new(self.fts_ready.load(Ordering::Relaxed)), fts_ready: AtomicBool::new(self.fts_ready.load(Ordering::Relaxed)),
}) })
} }
@@ -1074,6 +1155,127 @@ mod tests {
.expect("a fresh RAG with nothing indexed must rebuild cleanly"); .expect("a fresh RAG with nothing indexed must rebuild cleanly");
} }
#[test]
fn parse_float_array_dim_handles_arrays_lists_and_scalars() {
assert_eq!(parse_float_array_dim("FLOAT[768]"), Some(768));
assert_eq!(parse_float_array_dim("FLOAT[]"), None);
assert_eq!(parse_float_array_dim("VARCHAR"), None);
}
#[test]
fn introspect_dim_round_trips_the_open_dim() {
let db = TempDb::new("introspect");
assert_eq!(
DuckDbProvider::introspect_dim(&db.path).unwrap(),
None,
"a file that does not exist has no dim"
);
{
let _provider = DuckDbProvider::open(&db.path, 5).unwrap();
}
assert_eq!(DuckDbProvider::introspect_dim(&db.path).unwrap(), Some(5));
}
#[test]
fn introspect_dim_propagates_an_unopenable_existing_file() {
let db = TempDb::new("introspectgarbage");
fs::write(&db.path, b"not a duckdb database").unwrap();
let err = DuckDbProvider::introspect_dim(&db.path).unwrap_err();
assert!(
format!("{err:#}").contains(&format!(
"Cannot inspect the existing RAG store at '{}'",
db.path.display()
)),
"an existing-but-unopenable file must be an error naming the path; got: {err:#}"
);
}
#[tokio::test]
async fn rebuild_indexes_self_heals_dim_from_the_vectors_it_writes() {
let db = TempDb::new("selfheal");
let mut provider = DuckDbProvider::open(&db.path, 5).unwrap();
let mut data = minimal_rag_data();
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
provider.rebuild_indexes(&data, true).await.unwrap();
let results = provider
.vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
.await
.unwrap();
assert_eq!(results.len(), 1);
drop(provider);
assert_eq!(DuckDbProvider::introspect_dim(&db.path).unwrap(), Some(3));
}
#[tokio::test]
async fn a_self_healed_dim_is_visible_through_duplicate_clones() {
let db = TempDb::new("dimdup");
let mut provider = DuckDbProvider::open(&db.path, 5).unwrap();
let dup = provider.duplicate(&minimal_rag_data());
let mut data = minimal_rag_data();
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
provider.rebuild_indexes(&data, true).await.unwrap();
let results = dup.vector_search(&[0.1, 0.2, 0.3], 5, 0.0).await.unwrap();
assert_eq!(
results.len(),
1,
"a duplicate() clone must observe the dim written by a rebuild through \
the original"
);
}
#[tokio::test]
async fn rebuild_indexes_rejects_mixed_dim_vectors() {
let db = TempDb::new("mixeddim");
let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
let mut data = minimal_rag_data();
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
data.vectors.insert(DocumentId(1), vec![0.1, 0.2, 0.3, 0.4]);
let err = provider.rebuild_indexes(&data, true).await.unwrap_err();
assert!(
err.to_string().contains("rebuild batch mixes"),
"got: {err}"
);
}
#[tokio::test]
async fn vector_search_dim_mismatch_on_empty_store_returns_nothing() {
let db = TempDb::new("dimempty");
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
let results = provider.vector_search(&[0.1, 0.2], 5, 0.0).await.unwrap();
assert!(
results.is_empty(),
"a never-synced store must answer 'nothing', not a cast error"
);
}
#[tokio::test]
async fn vector_search_dim_mismatch_on_populated_store_errors() {
let db = TempDb::new("dimfull");
let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
let mut data = minimal_rag_data();
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
provider.rebuild_indexes(&data, true).await.unwrap();
let err = provider
.vector_search(&[0.1, 0.2], 5, 0.0)
.await
.unwrap_err();
assert!(err.to_string().contains("was built with"), "got: {err}");
}
#[tokio::test] #[tokio::test]
async fn duplicate_shares_the_same_connection() { async fn duplicate_shares_the_same_connection() {
let db = TempDb::new("dup"); let db = TempDb::new("dup");
+96 -11
View File
@@ -9,6 +9,7 @@ use reqwest::{Client, Response, StatusCode};
use serde_json::Value; use serde_json::Value;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use url::{Host, Url};
/// Marks a `DocumentId` that stands in for a point id Coyote cannot carry /// Marks a `DocumentId` that stands in for a point id Coyote cannot carry
/// directly. Qdrant accepts UUID strings as point ids, and that is what /// directly. Qdrant accepts UUID strings as point ids, and that is what
@@ -104,7 +105,7 @@ fn parse_search_hits(
let score = pt["score"].as_f64()? as f32; let score = pt["score"].as_f64()? as f32;
Some((interner.document_id(&pt["id"])?, score)) Some((interner.document_id(&pt["id"])?, score))
}) })
.filter(|(_, score)| *score > min_score) .filter(|(_, score)| min_score <= 0.0 || *score > min_score)
.collect()) .collect())
} }
@@ -183,7 +184,22 @@ pub struct QdrantProvider {
} }
impl QdrantProvider { impl QdrantProvider {
fn make_client(api_key: Option<&str>) -> Result<Client> { fn skips_proxy(base_url: &str) -> bool {
let Ok(url) = Url::parse(base_url) else {
return false;
};
match url.host() {
Some(Host::Domain(name)) => {
name == "localhost" || name.ends_with(".localhost") || name.ends_with(".local")
}
Some(Host::Ipv4(ip)) => ip.is_loopback() || ip.is_private() || ip.is_link_local(),
// No stable is_unique_local, so fc00::/7 is matched directly.
Some(Host::Ipv6(ip)) => ip.is_loopback() || ip.segments()[0] & 0xfe00 == 0xfc00,
None => false,
}
}
fn make_client(base_url: &str, api_key: Option<&str>) -> Result<Client> {
let mut headers = HeaderMap::new(); let mut headers = HeaderMap::new();
if let Some(key) = api_key { if let Some(key) = api_key {
let mut value = let mut value =
@@ -191,10 +207,11 @@ impl QdrantProvider {
value.set_sensitive(true); value.set_sensitive(true);
headers.insert("api-key", value); headers.insert("api-key", value);
} }
Client::builder() let mut builder = Client::builder().default_headers(headers);
.default_headers(headers) if Self::skips_proxy(base_url) {
.build() builder = builder.no_proxy();
.context("Failed to build reqwest client") }
builder.build().context("Failed to build reqwest client")
} }
pub(crate) fn normalize_base_url(host: &str) -> String { pub(crate) fn normalize_base_url(host: &str) -> String {
@@ -219,7 +236,7 @@ impl QdrantProvider {
api_key: Option<&str>, api_key: Option<&str>,
) -> Result<Value> { ) -> Result<Value> {
let base_url = Self::normalize_base_url(host); let base_url = Self::normalize_base_url(host);
let client = Self::make_client(api_key)?; let client = Self::make_client(&base_url, api_key)?;
let resp = client let resp = client
.get(format!("{base_url}/collections/{collection}")) .get(format!("{base_url}/collections/{collection}"))
.send() .send()
@@ -237,7 +254,7 @@ impl QdrantProvider {
pub async fn new(host: &str, collection: &str, api_key: Option<&str>) -> Result<Self> { pub async fn new(host: &str, collection: &str, api_key: Option<&str>) -> Result<Self> {
let base_url = Self::normalize_base_url(host); let base_url = Self::normalize_base_url(host);
let client = Self::make_client(api_key)?; let client = Self::make_client(&base_url, api_key)?;
let resp = client let resp = client
.get(format!("{base_url}/collections/{collection}")) .get(format!("{base_url}/collections/{collection}"))
.send() .send()
@@ -260,7 +277,7 @@ impl QdrantProvider {
pub async fn list_collections(host: &str, api_key: Option<&str>) -> Result<Vec<String>> { pub async fn list_collections(host: &str, api_key: Option<&str>) -> Result<Vec<String>> {
let base_url = Self::normalize_base_url(host); let base_url = Self::normalize_base_url(host);
let client = Self::make_client(api_key)?; let client = Self::make_client(&base_url, api_key)?;
let resp = client let resp = client
.get(format!("{base_url}/collections")) .get(format!("{base_url}/collections"))
.send() .send()
@@ -310,7 +327,7 @@ impl QdrantProvider {
api_key: Option<&str>, api_key: Option<&str>,
) -> Result<Option<String>> { ) -> Result<Option<String>> {
let base_url = Self::normalize_base_url(host); let base_url = Self::normalize_base_url(host);
let client = Self::make_client(api_key)?; let client = Self::make_client(&base_url, api_key)?;
let url = format!("{base_url}/collections/{collection}/points/scroll"); let url = format!("{base_url}/collections/{collection}/points/scroll");
let body = serde_json::json!({ "limit": 1, "with_payload": false }); let body = serde_json::json!({ "limit": 1, "with_payload": false });
@@ -353,7 +370,8 @@ impl RagProvider for QdrantProvider {
// `score_threshold` is deliberately NOT sent. It is metric-aware: on Cosine // `score_threshold` is deliberately NOT sent. It is metric-aware: on Cosine
// collections 0.0 means "no floor" as expected, but Euclid collections score // collections 0.0 means "no floor" as expected, but Euclid collections score
// by negative distance, where 0.0 filters everything out. The attach wizard // by negative distance, where 0.0 filters everything out. The attach wizard
// does not pin the distance metric, so filter locally instead. // does not pin the distance metric, so filter locally instead; i.e. where a
// 0.0 floor is correctly treated as "no floor" (see `parse_search_hits`).
let body = serde_json::json!({ let body = serde_json::json!({
"vector": embedding, "vector": embedding,
"limit": top_k, "limit": top_k,
@@ -576,6 +594,73 @@ mod tests {
assert!(provider.fetch_content(&[]).await.unwrap().is_empty()); assert!(provider.fetch_content(&[]).await.unwrap().is_empty());
} }
#[test]
fn local_and_private_hosts_skip_the_proxy() {
for host in [
"http://localhost:6333",
"http://127.0.0.1:6333",
"http://192.168.0.56:6333",
"http://10.1.2.3:6333",
"http://172.16.4.5:6333",
"http://qdrant.local:6333",
"http://[::1]:6333",
] {
assert!(
QdrantProvider::skips_proxy(host),
"{host} should not be proxied"
);
}
}
#[test]
fn public_hosts_still_honour_the_environment() {
for host in [
"https://qdrant.example.com",
"http://8.8.8.8:6333",
"https://xyz.eu-central.aws.cloud.qdrant.io:6333",
"http://172.32.0.1:6333",
] {
assert!(
!QdrantProvider::skips_proxy(host),
"{host} must keep the environment's proxy"
);
}
}
/// Euclid collections score by NEGATIVE distance, so the 0.0 the caller
/// passes must mean "no floor". Filtering on it drops every hit — the exact
/// bug that keeps Qdrant's own `score_threshold` off the wire.
#[test]
fn a_zero_floor_keeps_negative_euclid_scores() {
let mut interner = PointIdInterner::default();
let search = serde_json::json!({
"result": [
{"id": 1, "score": -0.12},
{"id": 2, "score": -8.5},
]
});
let hits = parse_search_hits(&mut interner, &search, 0.0).unwrap();
assert_eq!(hits.len(), 2, "a 0.0 floor must not drop negative scores");
}
#[test]
fn a_positive_floor_still_filters() {
let mut interner = PointIdInterner::default();
let search = serde_json::json!({
"result": [
{"id": 1, "score": 0.9},
{"id": 2, "score": 0.2},
]
});
let hits = parse_search_hits(&mut interner, &search, 0.5).unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].0, DocumentId(1));
}
/// A UUID-keyed collection has to survive the whole `vector_search` → /// A UUID-keyed collection has to survive the whole `vector_search` →
/// `fetch_content` round trip, and the fetch must ask Qdrant for the ORIGINAL /// `fetch_content` round trip, and the fetch must ask Qdrant for the ORIGINAL
/// string id. Parsing ids with `as_u64()` used to drop these hits inside a /// string id. Parsing ids with `as_u64()` used to drop these hits inside a
+34 -2
View File
@@ -39,8 +39,10 @@ static INLINE_CODE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"`([^`\n]+
static IMAGE_RE: LazyLock<Regex> = static IMAGE_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"!\[([^\]]*)\]\(([^)]+)\)").unwrap()); LazyLock::new(|| Regex::new(r"!\[([^\]]*)\]\(([^)]+)\)").unwrap());
static LINK_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\[([^\]]+)\]\(([^)]+)\)").unwrap()); static LINK_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\[([^\]]+)\]\(([^)]+)\)").unwrap());
static BOLD_AST_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\*\*([^*\n]+)\*\*").unwrap()); static BOLD_AST_RE: LazyLock<Regex> =
static BOLD_US_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"__([^_\n]+)__").unwrap()); LazyLock::new(|| Regex::new(r"\*\*((?:[^*\n]|\*(?!\*))+?)\*\*").unwrap());
static BOLD_US_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"__((?:[^_\n]|_(?!_))+?)__").unwrap());
static ITALIC_AST_RE: LazyLock<Regex> = static ITALIC_AST_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?<![*\w])\*(?!\s)([^*\n]+?)(?<!\s)\*(?!\*)").unwrap()); LazyLock::new(|| Regex::new(r"(?<![*\w])\*(?!\s)([^*\n]+?)(?<!\s)\*(?!\*)").unwrap());
static ITALIC_US_RE: LazyLock<Regex> = static ITALIC_US_RE: LazyLock<Regex> =
@@ -2419,6 +2421,36 @@ std::error::Error>> {
assert!(result.contains("\x1b[9m"), "strikethrough SGR: {result:?}"); assert!(result.contains("\x1b[9m"), "strikethrough SGR: {result:?}");
} }
#[test]
fn bold_asterisk_wraps_italic_asterisk() {
let styles = test_styles();
let result = apply_inline("**loud *soft* loud**", &styles);
assert!(
!result.contains("**"),
"outer bold markers stripped: {result:?}"
);
assert!(result.contains("\x1b[1m"), "bold SGR present: {result:?}");
assert!(result.contains("\x1b[3m"), "italic SGR present: {result:?}");
assert!(result.contains("soft"));
}
#[test]
fn bold_underscore_wraps_italic_underscore() {
let styles = test_styles();
let result = apply_inline("__loud _soft_ loud__", &styles);
assert!(
!result.contains("__"),
"outer bold markers stripped: {result:?}"
);
assert!(result.contains("\x1b[1m"), "bold SGR present: {result:?}");
assert!(result.contains("\x1b[3m"), "italic SGR present: {result:?}");
assert!(result.contains("soft"));
}
#[test] #[test]
fn bold_wraps_inline_code() { fn bold_wraps_inline_code() {
let styles = test_styles(); let styles = test_styles();
+7 -3
View File
@@ -19,8 +19,8 @@ use crate::config::{AssetCategory, paths};
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail}; use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
use crate::render::render_error; use crate::render::render_error;
use crate::utils::{ use crate::utils::{
AbortSignal, SHELL, abortable_run_with_spinner, create_abort_signal, dimmed_text, run_command, AbortSignal, SHELL, abortable_run_with_spinner, create_abort_signal, dimmed_text,
set_text, temp_file, drain_stale_tty_input, run_command, set_text, temp_file,
}; };
use crate::sandbox::SANDBOX_ENV_FLAG; use crate::sandbox::SANDBOX_ENV_FLAG;
@@ -411,6 +411,10 @@ Type ".help" for additional help.
} }
} }
// Discard any stray terminal-query reply bytes (e.g. late colorsaurus
// OSC 11 / DA1 responses) so they don't get injected into the prompt.
drain_stale_tty_input();
loop { loop {
if self.abort_signal.aborted_ctrld() { if self.abort_signal.aborted_ctrld() {
break; break;
@@ -892,7 +896,7 @@ pub async fn run_repl_command(
".rag" => match split_first_arg(args) { ".rag" => match split_first_arg(args) {
Some(("attach", rest)) => match rest { Some(("attach", rest)) => match rest {
Some(name) if !name.trim().is_empty() => { Some(name) if !name.trim().is_empty() => {
ctx.attach_rag(name.trim()).await?; ctx.attach_rag(name.trim(), abort_signal.clone()).await?;
} }
_ => println!("Usage: .rag attach <name>"), _ => println!("Usage: .rag attach <name>"),
}, },
+296 -15
View File
@@ -80,6 +80,8 @@ pub(crate) struct CredentialSpec {
pub env_var: String, pub env_var: String,
pub proxy_managed: bool, pub proxy_managed: bool,
pub inject: Vec<InjectRule>, pub inject: Vec<InjectRule>,
pub custom_hosts: Vec<String>,
pub custom_allow_entries: Vec<String>,
pub servers: Vec<String>, pub servers: Vec<String>,
} }
@@ -110,8 +112,13 @@ pub(crate) fn collect_credentials(
} }
let mut by_secret: BTreeMap<String, Aggregate> = BTreeMap::new(); let mut by_secret: BTreeMap<String, Aggregate> = BTreeMap::new();
let mut hosts_by_server: BTreeMap<String, ServerSecretHosts> = BTreeMap::new();
for (server_name, config) in servers { for (server_name, config) in servers {
for occurrence in collect_server_occurrences(config)? { let occurrences = collect_server_occurrences(config)?;
if !occurrences.is_empty() {
hosts_by_server.insert(server_name.clone(), server_secret_hosts(config));
}
for occurrence in occurrences {
let agg = by_secret let agg = by_secret
.entry(occurrence.secret_name) .entry(occurrence.secret_name)
.or_insert_with(|| Aggregate { .or_insert_with(|| Aggregate {
@@ -152,8 +159,9 @@ pub(crate) fn collect_credentials(
eprintln!( eprintln!(
"MCP secrets {} all target the '{header}' header for '{domain}'. The \ "MCP secrets {} all target the '{header}' header for '{domain}'. The \
sandbox proxy cannot tell which one a given request needs, so these \ sandbox proxy cannot tell which one a given request needs, so these \
secrets will be resolved from environment variables inside the \ secrets will be provisioned as placeholder-based custom secrets \
sandbox instead.", instead: each env var holds a unique placeholder that the proxy \
swaps for the real value in outbound request headers.",
quoted_list(secrets) quoted_list(secrets)
); );
slot_conflicted.extend(secrets.iter().cloned()); slot_conflicted.extend(secrets.iter().cloned());
@@ -191,6 +199,21 @@ pub(crate) fn collect_credentials(
} }
let proxy_managed = agg.all_injectable && !slot_conflicted.contains(&secret_name); let proxy_managed = agg.all_injectable && !slot_conflicted.contains(&secret_name);
let (custom_hosts, custom_allow_entries) = if proxy_managed {
(Vec::new(), Vec::new())
} else {
let mut targets: BTreeSet<String> = BTreeSet::new();
let mut allow: BTreeSet<String> = BTreeSet::new();
for hosts in agg
.servers
.iter()
.filter_map(|server| hosts_by_server.get(server))
{
targets.extend(hosts.targets.iter().cloned());
allow.extend(hosts.allow_entries.iter().cloned());
}
(targets.into_iter().collect(), allow.into_iter().collect())
};
credentials.push(CredentialSpec { credentials.push(CredentialSpec {
secret_name, secret_name,
service_id, service_id,
@@ -201,6 +224,8 @@ pub(crate) fn collect_credentials(
} else { } else {
Vec::new() Vec::new()
}, },
custom_hosts,
custom_allow_entries,
servers: agg.servers.into_iter().collect(), servers: agg.servers.into_iter().collect(),
}); });
} }
@@ -365,6 +390,98 @@ fn parse_https_domain(raw: &str) -> Option<String> {
} }
} }
#[derive(Debug, Default, PartialEq, Eq)]
pub(crate) struct ServerSecretHosts {
pub targets: BTreeSet<String>,
pub allow_entries: BTreeSet<String>,
}
pub(crate) fn server_secret_hosts(server: &Value) -> ServerSecretHosts {
let mut hosts = ServerSecretHosts::default();
scrape_hosts_value(server, &mut hosts);
if let Some(host) = server
.get("url")
.and_then(Value::as_str)
.and_then(|raw| Url::parse(raw).ok())
.and_then(|url| url.host_str().map(str::to_string))
.filter(|host| is_usable_host(host))
{
hosts.targets.insert(host);
}
hosts
}
fn scrape_hosts_value(value: &Value, out: &mut ServerSecretHosts) {
match value {
Value::String(s) => {
for url in urls_in_text(s) {
let Some(host) = url.host_str().filter(|h| is_usable_host(h)) else {
continue;
};
out.targets.insert(host.to_string());
if let Some(entry) = allow_entry_for_parsed(&url) {
out.allow_entries.insert(entry);
}
}
}
Value::Object(map) => {
for v in map.values() {
scrape_hosts_value(v, out);
}
}
Value::Array(arr) => {
for v in arr {
scrape_hosts_value(v, out);
}
}
_ => {}
}
}
/// A host usable in the sbx target and allow grammars: non-empty, not a
/// bracketed IPv6 literal (no spelling in either grammar), and not an
/// unresolved `{{placeholder}}` fragment (would register garbage targets
/// and could fail kit validation at create).
fn is_usable_host(host: &str) -> bool {
!host.is_empty() && !host.starts_with('[') && !host.contains('{') && !host.contains('}')
}
/// Every http(s) URL embedded in `text`. Tokens end at whitespace, quotes,
/// or URL-hostile punctuation so comma- or bracket-separated lists don't
/// bleed into one another.
fn urls_in_text(text: &str) -> Vec<Url> {
const TERMINATORS: &[char] = &['"', '\'', ',', ';', '(', ')', '[', ']', '{', '}', '<', '>'];
let mut out = Vec::new();
for (idx, _) in text.match_indices("http") {
let candidate = &text[idx..];
if !candidate.starts_with("http://") && !candidate.starts_with("https://") {
continue;
}
let end = candidate
.find(|c: char| c.is_whitespace() || TERMINATORS.contains(&c))
.unwrap_or(candidate.len());
if let Ok(url) = Url::parse(&candidate[..end]) {
out.push(url);
}
}
out
}
/// Extracts the host of every http(s) URL embedded in `text`, ports stripped.
/// Bracketed IPv6 hosts and placeholder fragments are skipped.
pub(crate) fn hosts_in_text(text: &str) -> BTreeSet<String> {
urls_in_text(text)
.iter()
.filter_map(|url| url.host_str())
.filter(|host| is_usable_host(host))
.map(str::to_string)
.collect()
}
/// Collects a network allow-list entry for every remote MCP server `url` /// Collects a network allow-list entry for every remote MCP server `url`
/// (http/https), so user-configured servers are reachable regardless of how, /// (http/https), so user-configured servers are reachable regardless of how,
/// or whether, their credentials are provisioned. Https on the default port /// or whether, their credentials are provisioned. Https on the default port
@@ -394,13 +511,14 @@ pub(crate) fn allow_entry_for_url(raw: &str) -> Option<String> {
return None; return None;
} }
let host = url.host_str()?; allow_entry_for_parsed(&url)
if host.is_empty() || host.starts_with('[') { }
return None;
} fn allow_entry_for_parsed(url: &Url) -> Option<String> {
let host = url.host_str().filter(|h| is_usable_host(h))?;
match url.port_or_known_default() { match url.port_or_known_default() {
Some(443) if scheme == "https" => Some(host.to_string()), Some(443) if url.scheme() == "https" => Some(host.to_string()),
Some(port) => Some(format!("{host}:{port}")), Some(port) => Some(format!("{host}:{port}")),
None => None, None => None,
} }
@@ -489,12 +607,17 @@ pub(crate) fn render_mixin_document(
serde_yaml::to_string(&mixin).context("Failed to serialize generated sandbox mixin") serde_yaml::to_string(&mixin).context("Failed to serialize generated sandbox mixin")
} }
/// Renders the generated `coyote-mcp` mixin. Only proxy-managed credentials
/// are declared. sbx does not materialize env vars for `proxyManaged: false`
/// mixin credentials, so the rest are provisioned as custom secrets outside
/// the mixin, and only their target hosts join the network allow list here.
pub(crate) fn render_mixin_yaml( pub(crate) fn render_mixin_yaml(
credentials: &[CredentialSpec], credentials: &[CredentialSpec],
server_allow_entries: &[String], server_allow_entries: &[String],
) -> Result<String> { ) -> Result<String> {
let entries = credentials let entries = credentials
.iter() .iter()
.filter(|c| c.proxy_managed)
.map(|c| CredentialEntry { .map(|c| CredentialEntry {
service: c.service_id.clone(), service: c.service_id.clone(),
description: format!( description: format!(
@@ -510,14 +633,20 @@ pub(crate) fn render_mixin_yaml(
}) })
.collect(); .collect();
let mut allow_entries: Vec<String> = server_allow_entries.to_vec();
for credential in credentials.iter().filter(|c| !c.proxy_managed) {
allow_entries.extend(credential.custom_allow_entries.iter().cloned());
}
render_mixin_document( render_mixin_document(
MCP_MIXIN_NAME, MCP_MIXIN_NAME,
"Auto-generated by Coyote at launch: allows network egress to the user's remote MCP \ "Auto-generated by Coyote at launch: allows network egress to the user's remote MCP \
servers and declares their credentials so Docker Sandboxes binds them (bindings are \ servers and declares their credentials so Docker Sandboxes binds them (bindings are \
approved on first interactive run). Values are pre-seeded from Coyote's vault via \ approved on first interactive run). Proxy-injectable values are pre-seeded from \
`sbx secret set`.", Coyote's vault via `sbx secret set`; the remaining secrets are provisioned as \
placeholder-based custom secrets via `sbx secret set-custom`.",
entries, entries,
server_allow_entries, &allow_entries,
) )
} }
@@ -603,6 +732,11 @@ mod tests {
assert_eq!(cred.service_id, "github-pat"); assert_eq!(cred.service_id, "github-pat");
assert_eq!(cred.env_var, "COYOTE_SECRET_GITHUB_PAT"); assert_eq!(cred.env_var, "COYOTE_SECRET_GITHUB_PAT");
assert!(cred.proxy_managed); assert!(cred.proxy_managed);
assert!(
cred.custom_hosts.is_empty(),
"proxy-managed credentials are provisioned via `sbx secret set`, \
not set-custom, so they carry no custom hosts"
);
assert_eq!( assert_eq!(
cred.inject, cred.inject,
vec![InjectRule { vec![InjectRule {
@@ -839,6 +973,12 @@ mod tests {
"secrets sharing an inject domain must both fall back to env" "secrets sharing an inject domain must both fall back to env"
); );
assert!(creds.iter().all(|c| c.inject.is_empty())); assert!(creds.iter().all(|c| c.inject.is_empty()));
assert!(
creds
.iter()
.all(|c| c.custom_hosts == vec!["api.githubcopilot.com".to_string()]),
"demoted secrets must derive their set-custom targets from the server url"
);
} }
#[test] #[test]
@@ -1018,21 +1158,162 @@ mod tests {
} }
#[test] #[test]
fn rendered_mixin_omits_inject_and_permissions_for_env_based_secrets() { fn rendered_mixin_drops_env_based_credentials() {
let servers = servers(json!({ let servers = servers(json!({
"local": { "command": "run", "env": { "KEY": "{{NOTION_TOKEN}}" } } "local": { "command": "run", "env": { "KEY": "{{NOTION_TOKEN}}" } }
})); }));
let creds = collect_credentials(&servers).unwrap(); let creds = collect_credentials(&servers).unwrap();
assert_eq!(creds.len(), 1);
assert!(!creds[0].proxy_managed, "precondition");
assert!(
creds[0].custom_hosts.is_empty(),
"a stdio server with no URLs anywhere yields no derivable hosts"
);
let yaml = render_mixin_yaml(&creds, &[]).unwrap(); let yaml = render_mixin_yaml(&creds, &[]).unwrap();
let value: serde_yaml::Value = serde_yaml::from_str(&yaml).unwrap(); let value: serde_yaml::Value = serde_yaml::from_str(&yaml).unwrap();
let cred = &value["credentials"][0]; assert!(
assert_eq!(cred["apiKey"]["proxyManaged"].as_bool(), Some(false)); value.get("credentials").is_none(),
assert!(cred["apiKey"].get("inject").is_none()); "sbx does not materialize env vars for proxyManaged:false mixin \
credentials; declaring them would be dead config"
);
assert!(value.get("permissions").is_none()); assert!(value.get("permissions").is_none());
} }
#[test]
fn hosts_in_text_extracts_hosts_and_strips_ports() {
assert_eq!(
hosts_in_text("https://api.example.com:8443/v1"),
BTreeSet::from(["api.example.com".to_string()])
);
assert_eq!(
hosts_in_text("see http://a.example.com/x and https://b.example.com/y"),
BTreeSet::from(["a.example.com".to_string(), "b.example.com".to_string()])
);
assert_eq!(
hosts_in_text("https://api.example.com/mcp?key={{KEY}}"),
BTreeSet::from(["api.example.com".to_string()])
);
assert!(hosts_in_text("ws://sock.example.com/mcp").is_empty());
assert!(hosts_in_text("qdrant.example.com:6333").is_empty());
assert!(hosts_in_text("https://[::1]:8443/mcp").is_empty());
assert!(hosts_in_text("httpserver is not a scheme").is_empty());
}
#[test]
fn hosts_in_text_terminates_urls_at_punctuation() {
assert_eq!(
hosts_in_text("https://a.example.com,https://b.example.com;https://c.example.com"),
BTreeSet::from([
"a.example.com".to_string(),
"b.example.com".to_string(),
"c.example.com".to_string()
])
);
assert_eq!(
hosts_in_text("(see https://docs.example.com)"),
BTreeSet::from(["docs.example.com".to_string()])
);
}
#[test]
fn server_secret_hosts_unions_url_host_and_scraped_urls() {
let server = json!({
"command": "run",
"args": ["--endpoint", "https://api.vendor.example/v2"],
"env": { "BASE_URL": "http://internal.example.com:8080/api" }
});
let hosts = server_secret_hosts(&server);
assert_eq!(
hosts.targets,
BTreeSet::from([
"api.vendor.example".to_string(),
"internal.example.com".to_string()
]),
"set-custom targets are port-stripped"
);
assert_eq!(
hosts.allow_entries,
BTreeSet::from([
"api.vendor.example".to_string(),
"internal.example.com:8080".to_string()
]),
"allow entries keep non-default ports so the hosts stay reachable"
);
}
#[test]
fn server_secret_hosts_includes_non_http_url_host() {
let server = json!({ "url": "ws://sock.example.com/mcp" });
let hosts = server_secret_hosts(&server);
assert_eq!(
hosts.targets,
BTreeSet::from(["sock.example.com".to_string()])
);
assert!(
hosts.allow_entries.is_empty(),
"non-http(s) urls have no spelling in the allow grammar"
);
}
#[test]
fn server_secret_hosts_skips_placeholder_hosts() {
let server = json!({
"url": "https://{{TENANT}}.example.com/mcp",
"env": { "API_URL": "https://{{REGION}}.api.example.com/v1" }
});
let hosts = server_secret_hosts(&server);
assert!(
hosts.targets.is_empty() && hosts.allow_entries.is_empty(),
"a host containing an unresolved placeholder must never reach \
set-custom argv or the mixin allow list: {hosts:?}"
);
}
#[test]
fn demoted_credential_hosts_are_unioned_into_allow() {
let servers = servers(json!({
"local": {
"command": "run",
"env": {
"KEY": "{{TOKEN}}",
"API_URL": "https://api.internal.example.com:8443/v1"
}
}
}));
let creds = collect_credentials(&servers).unwrap();
assert_eq!(creds.len(), 1);
assert!(!creds[0].proxy_managed, "precondition");
assert_eq!(
creds[0].custom_hosts,
vec!["api.internal.example.com".to_string()]
);
assert_eq!(
creds[0].custom_allow_entries,
vec!["api.internal.example.com:8443".to_string()],
"allow entries keep the port that set-custom targets must strip"
);
let yaml = render_mixin_yaml(&creds, &[]).unwrap();
let value: serde_yaml::Value = serde_yaml::from_str(&yaml).unwrap();
assert!(value.get("credentials").is_none());
assert_eq!(
value["permissions"]["network"]["allow"][0].as_str(),
Some("api.internal.example.com:8443"),
"a demoted credential's derived hosts must still be reachable"
);
}
#[test] #[test]
fn rendered_mixin_is_deterministic() { fn rendered_mixin_is_deterministic() {
let servers = servers(json!({ let servers = servers(json!({
+318 -7
View File
@@ -1,9 +1,10 @@
use std::env; use std::env;
use std::fs; use std::fs;
use std::fs::{read_dir, read_to_string}; use std::fs::{read_dir, read_to_string};
use std::io;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use anyhow::{Context, Result}; use anyhow::{Context, Result, anyhow, bail};
use serde_yaml::Value; use serde_yaml::Value;
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
@@ -12,6 +13,7 @@ use crate::config::paths;
const SBX_MIXIN_FILE_NAME: &str = "sbx-mixin.yaml"; const SBX_MIXIN_FILE_NAME: &str = "sbx-mixin.yaml";
const SBX_MIXIN_FILE_SUFFIX: &str = ".sbx-mixin.yaml"; const SBX_MIXIN_FILE_SUFFIX: &str = ".sbx-mixin.yaml";
const KIT_SPEC_FILE_NAME: &str = "spec.yaml"; const KIT_SPEC_FILE_NAME: &str = "spec.yaml";
const MIXIN_FILES_DIR_NAME: &str = "files";
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct DiscoveredMixin { pub struct DiscoveredMixin {
@@ -34,33 +36,152 @@ impl DiscoveredMixin {
pub fn wrap_mixin_as_kit(mixin_path: &Path) -> Result<PathBuf> { pub fn wrap_mixin_as_kit(mixin_path: &Path) -> Result<PathBuf> {
let bytes = fs::read(mixin_path) let bytes = fs::read(mixin_path)
.with_context(|| format!("Failed to read sbx mixin {}", mixin_path.display()))?; .with_context(|| format!("Failed to read sbx mixin {}", mixin_path.display()))?;
wrap_mixin_bytes_as_kit(&bytes, &mixin_path.display().to_string()) let label = mixin_path.display().to_string();
let files = mixin_path
.parent()
.map(|p| p.join(MIXIN_FILES_DIR_NAME))
.filter(|p| p.is_dir())
.map(|dir| collect_staged_files(&dir))
.transpose()?
.unwrap_or_default();
stage_kit(&bytes, &files, &label)
} }
pub fn wrap_mixin_bytes_as_kit(bytes: &[u8], label: &str) -> Result<PathBuf> { pub fn wrap_mixin_bytes_as_kit(bytes: &[u8], label: &str) -> Result<PathBuf> {
stage_kit(bytes, &[], label)
}
struct StagedFile {
relpath: PathBuf,
mode: u32,
bytes: Vec<u8>,
}
fn stage_kit(spec_bytes: &[u8], files: &[StagedFile], label: &str) -> Result<PathBuf> {
let mut hasher = Sha256::new(); let mut hasher = Sha256::new();
hasher.update(bytes); hasher.update(spec_bytes);
for f in files {
let rel_str = f.relpath.to_str().ok_or_else(|| {
anyhow!(
"Non-UTF-8 path inside mixin {MIXIN_FILES_DIR_NAME}/: {}",
f.relpath.display()
)
})?;
hasher.update(b"\0COYOTE_MIXIN_FILE\0");
hasher.update((rel_str.len() as u64).to_le_bytes());
hasher.update(rel_str.as_bytes());
hasher.update(f.mode.to_le_bytes());
hasher.update((f.bytes.len() as u64).to_le_bytes());
hasher.update(&f.bytes);
}
let hash = format!("{:x}", hasher.finalize()); let hash = format!("{:x}", hasher.finalize());
let kit_dir = paths::sbx_mixin_kits_dir().join(&hash); let kit_dir = paths::sbx_mixin_kits_dir().join(&hash);
let spec_path = kit_dir.join(KIT_SPEC_FILE_NAME); let spec_path = kit_dir.join(KIT_SPEC_FILE_NAME);
let files_dst = kit_dir.join(MIXIN_FILES_DIR_NAME);
if let Ok(existing) = fs::read(&spec_path) let spec_matches = fs::read(&spec_path).is_ok_and(|existing| existing == spec_bytes);
&& existing == bytes let files_ready = files.is_empty() || files_dst.is_dir();
{ if spec_matches && files_ready {
return Ok(kit_dir); return Ok(kit_dir);
} }
fs::create_dir_all(&kit_dir) fs::create_dir_all(&kit_dir)
.with_context(|| format!("Failed to create mixin kit dir {}", kit_dir.display()))?; .with_context(|| format!("Failed to create mixin kit dir {}", kit_dir.display()))?;
fs::write(&spec_path, bytes) fs::write(&spec_path, spec_bytes)
.with_context(|| format!("Failed to write {}", spec_path.display()))?; .with_context(|| format!("Failed to write {}", spec_path.display()))?;
if !files.is_empty() {
if files_dst.exists() {
fs::remove_dir_all(&files_dst).with_context(|| {
format!(
"Failed to clear stale mixin files at {}",
files_dst.display()
)
})?;
}
for f in files {
let dst = files_dst.join(&f.relpath);
if let Some(parent) = dst.parent() {
fs::create_dir_all(parent)
.with_context(|| format!("Failed to create dir {}", parent.display()))?;
}
fs::write(&dst, &f.bytes)
.with_context(|| format!("Failed to write staged mixin file {}", dst.display()))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
fs::set_permissions(&dst, fs::Permissions::from_mode(f.mode))
.with_context(|| format!("Failed to set mode on {}", dst.display()))?;
}
}
}
debug!("Wrapped mixin {label} as kit at {}", kit_dir.display()); debug!("Wrapped mixin {label} as kit at {}", kit_dir.display());
Ok(kit_dir) Ok(kit_dir)
} }
fn collect_staged_files(root: &Path) -> Result<Vec<StagedFile>> {
let mut out = Vec::new();
walk_staged_files(root, Path::new(""), &mut out)?;
Ok(out)
}
fn walk_staged_files(abs_dir: &Path, rel_dir: &Path, out: &mut Vec<StagedFile>) -> Result<()> {
let rd = fs::read_dir(abs_dir)
.with_context(|| format!("Failed to read mixin files dir {}", abs_dir.display()))?;
let mut entries: Vec<_> = rd
.collect::<io::Result<Vec<_>>>()
.with_context(|| format!("Failed to iterate mixin files dir {}", abs_dir.display()))?;
entries.sort_by_key(|e| e.file_name());
for entry in entries {
let file_type = entry
.file_type()
.with_context(|| format!("Failed to stat {}", entry.path().display()))?;
let abs = entry.path();
let rel = rel_dir.join(entry.file_name());
if file_type.is_symlink() {
bail!(
"Symlinks are not allowed inside a mixin {MIXIN_FILES_DIR_NAME}/ tree: {}",
abs.display()
);
}
if file_type.is_dir() {
walk_staged_files(&abs, &rel, out)?;
} else if file_type.is_file() {
let bytes = fs::read(&abs)
.with_context(|| format!("Failed to read staged mixin file {}", abs.display()))?;
let mode = staged_file_mode(&entry)?;
out.push(StagedFile {
relpath: rel,
mode,
bytes,
});
}
}
Ok(())
}
#[cfg(unix)]
fn staged_file_mode(entry: &fs::DirEntry) -> Result<u32> {
use std::os::unix::fs::PermissionsExt;
let meta = entry
.metadata()
.with_context(|| format!("Failed to stat {}", entry.path().display()))?;
Ok(meta.permissions().mode() & 0o777)
}
#[cfg(not(unix))]
fn staged_file_mode(_entry: &fs::DirEntry) -> Result<u32> {
Ok(0o644)
}
pub fn discover() -> Result<Vec<DiscoveredMixin>> { pub fn discover() -> Result<Vec<DiscoveredMixin>> {
let mut out = Vec::new(); let mut out = Vec::new();
@@ -556,6 +677,196 @@ network:
"kit_path should not return the original file path" "kit_path should not return the original file path"
); );
} }
fn write_staged_file(mixin: &Path, rel: &str, content: &[u8]) {
let dst = mixin.parent().unwrap().join(MIXIN_FILES_DIR_NAME).join(rel);
fs::create_dir_all(dst.parent().unwrap()).unwrap();
fs::write(&dst, content).unwrap();
}
#[test]
#[serial]
fn wrap_mixin_as_kit_copies_sibling_files_tree_into_kit() {
let _guard = TestCacheDirGuard::new();
let mixin = write_mixin("files-copy", "kind: mixin\nname: probe\n");
write_staged_file(&mixin, "home/hello.md", b"# hello\n");
write_staged_file(&mixin, "home/nested/deep.txt", b"deep\n");
let kit_dir = wrap_mixin_as_kit(&mixin).unwrap();
assert!(kit_dir.join("spec.yaml").exists());
let files_root = kit_dir.join(MIXIN_FILES_DIR_NAME);
assert!(files_root.is_dir(), "kit dir must contain a files/ tree");
assert_eq!(
fs::read(files_root.join("home/hello.md")).unwrap(),
b"# hello\n"
);
assert_eq!(
fs::read(files_root.join("home/nested/deep.txt")).unwrap(),
b"deep\n"
);
}
#[test]
#[serial]
fn wrap_mixin_as_kit_hash_changes_when_a_staged_file_is_edited() {
let _guard = TestCacheDirGuard::new();
let mixin = write_mixin("files-hash-content", "kind: mixin\nname: probe\n");
write_staged_file(&mixin, "home/note.md", b"before\n");
let kit_before = wrap_mixin_as_kit(&mixin).unwrap();
write_staged_file(&mixin, "home/note.md", b"after\n");
let kit_after = wrap_mixin_as_kit(&mixin).unwrap();
assert_ne!(
kit_before, kit_after,
"editing a staged file must invalidate the kit hash"
);
assert_eq!(
fs::read(kit_after.join("files/home/note.md")).unwrap(),
b"after\n"
);
}
#[test]
#[serial]
fn wrap_mixin_as_kit_hash_changes_when_a_staged_file_is_added() {
let _guard = TestCacheDirGuard::new();
let mixin = write_mixin("files-hash-added", "kind: mixin\nname: probe\n");
write_staged_file(&mixin, "home/one.md", b"one\n");
let kit_before = wrap_mixin_as_kit(&mixin).unwrap();
write_staged_file(&mixin, "home/two.md", b"two\n");
let kit_after = wrap_mixin_as_kit(&mixin).unwrap();
assert_ne!(
kit_before, kit_after,
"adding a staged file must invalidate the kit hash"
);
}
#[test]
#[serial]
fn wrap_mixin_as_kit_hash_unchanged_when_no_files_dir() {
let _guard = TestCacheDirGuard::new();
let content = "kind: mixin\nname: legacy\n";
let mixin = write_mixin("legacy-no-files", content);
let with_helper = wrap_mixin_as_kit(&mixin).unwrap();
let bytes_only = wrap_mixin_bytes_as_kit(content.as_bytes(), "legacy").unwrap();
assert_eq!(
with_helper, bytes_only,
"mixins without a sibling files/ must keep the legacy bytes-only hash to reuse existing cache dirs"
);
}
#[test]
#[serial]
fn wrap_mixin_as_kit_ignores_sibling_files_that_is_not_a_directory() {
let _guard = TestCacheDirGuard::new();
let content = "kind: mixin\nname: probe\n";
let mixin = write_mixin("files-not-a-dir", content);
fs::write(mixin.parent().unwrap().join(MIXIN_FILES_DIR_NAME), b"decoy").unwrap();
let wrapped = wrap_mixin_as_kit(&mixin).unwrap();
let bytes_only = wrap_mixin_bytes_as_kit(content.as_bytes(), "probe").unwrap();
assert_eq!(
wrapped, bytes_only,
"a regular file named files must be ignored, not staged"
);
assert!(!wrapped.join(MIXIN_FILES_DIR_NAME).exists());
}
#[test]
#[serial]
fn wrap_mixin_as_kit_rebuilds_files_when_cache_dir_missing_files_tree() {
let _guard = TestCacheDirGuard::new();
let mixin = write_mixin("files-rebuild", "kind: mixin\nname: probe\n");
write_staged_file(&mixin, "home/hello.md", b"hi\n");
let kit_dir = wrap_mixin_as_kit(&mixin).unwrap();
let files_dst = kit_dir.join(MIXIN_FILES_DIR_NAME);
fs::remove_dir_all(&files_dst).unwrap();
assert!(!files_dst.exists());
let kit_again = wrap_mixin_as_kit(&mixin).unwrap();
assert_eq!(kit_again, kit_dir, "kit path is content-addressed");
assert!(
files_dst.is_dir(),
"a partial cache (spec present, files/ missing) must be rebuilt"
);
assert_eq!(fs::read(files_dst.join("home/hello.md")).unwrap(), b"hi\n");
}
#[test]
#[serial]
fn wrap_mixin_as_kit_deterministic_with_staged_files() {
let _guard = TestCacheDirGuard::new();
let content = "kind: mixin\nname: probe\n";
let mixin_one = write_mixin("determ-1", content);
write_staged_file(&mixin_one, "home/note.md", b"same\n");
let mixin_two = write_mixin("determ-2", content);
write_staged_file(&mixin_two, "home/note.md", b"same\n");
let kit_a = wrap_mixin_as_kit(&mixin_one).unwrap();
let kit_b = wrap_mixin_as_kit(&mixin_two).unwrap();
assert_eq!(
kit_a, kit_b,
"identical spec+files must produce the same content-addressed kit dir"
);
}
#[cfg(unix)]
#[test]
#[serial]
fn wrap_mixin_as_kit_rejects_symlinks_inside_files_tree() {
use std::os::unix::fs::symlink;
let _guard = TestCacheDirGuard::new();
let mixin = write_mixin("files-symlink", "kind: mixin\nname: probe\n");
let files_dir = mixin.parent().unwrap().join(MIXIN_FILES_DIR_NAME);
fs::create_dir_all(&files_dir).unwrap();
let target = files_dir.join("target.txt");
fs::write(&target, b"real").unwrap();
symlink(&target, files_dir.join("link.txt")).unwrap();
let err = wrap_mixin_as_kit(&mixin).unwrap_err();
let msg = format!("{err:#}");
assert!(
msg.contains("Symlinks are not allowed"),
"expected symlink rejection, got: {msg}"
);
}
#[cfg(unix)]
#[test]
#[serial]
fn wrap_mixin_as_kit_preserves_executable_bit() {
use std::os::unix::fs::PermissionsExt;
let _guard = TestCacheDirGuard::new();
let mixin = write_mixin("files-exec", "kind: mixin\nname: probe\n");
write_staged_file(&mixin, "bin/run.sh", b"#!/bin/sh\necho hi\n");
let src = mixin
.parent()
.unwrap()
.join(MIXIN_FILES_DIR_NAME)
.join("bin/run.sh");
fs::set_permissions(&src, fs::Permissions::from_mode(0o755)).unwrap();
let kit_dir = wrap_mixin_as_kit(&mixin).unwrap();
let dst = kit_dir.join("files/bin/run.sh");
let mode = fs::metadata(&dst).unwrap().permissions().mode() & 0o777;
assert_eq!(
mode, 0o755,
"executable bit must survive the copy into the kit dir"
);
}
} }
#[test] #[test]
+899 -70
View File
File diff suppressed because it is too large Load Diff
+29 -1
View File
@@ -35,6 +35,7 @@ use nu_ansi_term::Color;
use serde_json::Value; use serde_json::Value;
use std::borrow::Cow; use std::borrow::Cow;
use std::collections::VecDeque; use std::collections::VecDeque;
use std::io;
use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicBool;
use std::sync::{LazyLock, Mutex, OnceLock}; use std::sync::{LazyLock, Mutex, OnceLock};
use std::{cmp, env, path::PathBuf, process}; use std::{cmp, env, path::PathBuf, process};
@@ -45,7 +46,7 @@ pub static CODE_BLOCK_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?ms)```\w*(.*)```").unwrap()); LazyLock::new(|| Regex::new(r"(?ms)```\w*(.*)```").unwrap());
pub static THINK_TAG_RE: LazyLock<Regex> = pub static THINK_TAG_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)^\s*<think>.*?</think>(\s*|$)").unwrap()); LazyLock::new(|| Regex::new(r"(?s)^\s*<think>.*?</think>(\s*|$)").unwrap());
pub static IS_STDOUT_TERMINAL: LazyLock<bool> = LazyLock::new(|| std::io::stdout().is_terminal()); pub static IS_STDOUT_TERMINAL: LazyLock<bool> = LazyLock::new(|| io::stdout().is_terminal());
pub static HEADLESS: AtomicBool = AtomicBool::new(false); pub static HEADLESS: AtomicBool = AtomicBool::new(false);
pub static ACP_SERVER: AtomicBool = AtomicBool::new(false); pub static ACP_SERVER: AtomicBool = AtomicBool::new(false);
@@ -133,6 +134,33 @@ pub fn parse_bool(value: &str) -> Option<bool> {
} }
} }
pub fn drain_stale_tty_input() {
use crossterm::event::{poll, read};
use std::time::{Duration, Instant};
if !io::stdin().is_terminal() {
return;
}
if crossterm::terminal::enable_raw_mode().is_err() {
return;
}
let deadline = Instant::now() + Duration::from_millis(100);
while Instant::now() < deadline {
match poll(Duration::from_millis(10)) {
Ok(true) => {
if read().is_err() {
break;
}
}
_ => break,
}
}
let _ = crossterm::terminal::disable_raw_mode();
}
pub fn estimate_token_length(text: &str) -> usize { pub fn estimate_token_length(text: &str) -> usize {
let weighted: usize = text.chars().map(|c| if c.is_ascii() { 1 } else { 2 }).sum(); let weighted: usize = text.chars().map(|c| if c.is_ascii() { 1 } else { 2 }).sum();
weighted.div_ceil(4) weighted.div_ceil(4)
-14
View File
@@ -54,20 +54,6 @@ pub async fn expand_glob_paths<T: AsRef<str>>(
Ok(new_paths) Ok(new_paths)
} }
pub fn clear_dir(dir: &Path) -> Result<()> {
for entry in fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
if path.is_dir() {
fs::remove_dir_all(&path)?;
} else {
fs::remove_file(&path)?;
}
}
Ok(())
}
pub fn list_file_names<T: AsRef<Path>>(dir: T, ext: &str) -> Vec<String> { pub fn list_file_names<T: AsRef<Path>>(dir: T, ext: &str) -> Vec<String> {
match fs::read_dir(dir.as_ref()) { match fs::read_dir(dir.as_ref()) {
Ok(rd) => { Ok(rd) => {
+3
View File
@@ -12,6 +12,7 @@ pub use utils::prompt_provider_choice;
use crate::cli::Cli; use crate::cli::Cli;
use crate::config::AppConfig; use crate::config::AppConfig;
use crate::utils::drain_stale_tty_input;
use crate::vault::utils::ensure_password_file_initialized; use crate::vault::utils::ensure_password_file_initialized;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Result, anyhow, bail};
use fancy_regex::Regex; use fancy_regex::Regex;
@@ -151,6 +152,7 @@ impl Vault {
"Vault management is disabled in sandbox mode. Use `coyote --add-secret` on your host." "Vault management is disabled in sandbox mode. Use `coyote --add-secret` on your host."
); );
} }
drain_stale_tty_input();
let secret_value = Password::new("Enter the secret value:") let secret_value = Password::new("Enter the secret value:")
.with_validator(required!()) .with_validator(required!())
.with_display_mode(PasswordDisplayMode::Masked) .with_display_mode(PasswordDisplayMode::Masked)
@@ -190,6 +192,7 @@ impl Vault {
"Vault management is disabled in sandbox mode. Use `coyote --add-secret` on your host." "Vault management is disabled in sandbox mode. Use `coyote --add-secret` on your host."
); );
} }
drain_stale_tty_input();
let secret_value = Password::new("Enter the secret value:") let secret_value = Password::new("Enter the secret value:")
.with_validator(required!()) .with_validator(required!())
.with_display_mode(PasswordDisplayMode::Masked) .with_display_mode(PasswordDisplayMode::Masked)
+4
View File
@@ -1,5 +1,6 @@
use crate::config::ensure_parent_exists; use crate::config::ensure_parent_exists;
use crate::sandbox::{SANDBOX_ENV_FLAG, sandbox_secret_env_var}; use crate::sandbox::{SANDBOX_ENV_FLAG, sandbox_secret_env_var};
use crate::utils::drain_stale_tty_input;
use crate::vault::{SECRET_RE, Vault}; use crate::vault::{SECRET_RE, Vault};
use anyhow::Result; use anyhow::Result;
use anyhow::anyhow; use anyhow::anyhow;
@@ -68,6 +69,7 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
} }
} }
drain_stale_tty_input();
let ans = Confirm::new( let ans = Confirm::new(
format!( format!(
"The configured password file '{}' is empty. Create a password?", "The configured password file '{}' is empty. Create a password?",
@@ -107,6 +109,7 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
} }
} }
} else { } else {
drain_stale_tty_input();
let ans = Confirm::new("No password file configured. Do you want to create one now?") let ans = Confirm::new("No password file configured. Do you want to create one now?")
.with_default(true) .with_default(true)
.prompt()?; .prompt()?;
@@ -185,6 +188,7 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
} }
pub fn prompt_provider_choice() -> Result<Option<SupportedProvider>> { pub fn prompt_provider_choice() -> Result<Option<SupportedProvider>> {
drain_stale_tty_input();
let choices = vec![ let choices = vec![
"local - encrypted file on this machine", "local - encrypted file on this machine",
"aws_secrets_manager - AWS Secrets Manager", "aws_secrets_manager - AWS Secrets Manager",