Compare commits

...
Author SHA1 Message Date
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
25 changed files with 2193 additions and 168 deletions
+3
View File
@@ -5,3 +5,6 @@
.idea/ .idea/
/coyote.iml /coyote.iml
/.idea/ /.idea/
.coyote/**
.sisyphus/**
.coyote-project.json
+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.
+10 -2
View File
@@ -33,7 +33,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 +45,11 @@ 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}"
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
+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());
}
} }
+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);
}
}
+324 -41
View File
@@ -2,13 +2,13 @@ use super::access_token::{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,
@@ -624,15 +744,35 @@ pub async fn prepare_oauth_access_token(
}; };
let tokens = if Utc::now().timestamp() >= tokens.expires_at { let tokens = if Utc::now().timestamp() >= tokens.expires_at {
match provider.flow() { let guard = refresh_guard(client_name);
OAuthFlow::Pkce | OAuthFlow::DeviceCode => { let _guard = guard.lock().await;
refresh_oauth_token(client, provider, client_name, &tokens).await?
} // A concurrent caller may have refreshed while we waited for the
OAuthFlow::ClientCredentials => { // lock; a valid in-memory token means the winner already populated
run_client_credentials_flow(provider, client_name).await?; // the cache.
load_oauth_tokens(client_name) if is_valid_access_token(client_name) {
.ok_or_else(|| anyhow!("Token file missing after client_credentials refresh"))? 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 {
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
@@ -886,11 +1026,54 @@ 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::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 +1651,104 @@ 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]
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"));
}
} }
+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 =
+1 -1
View File
@@ -103,7 +103,7 @@ impl McpFactory {
} }
let bearer_token = if spec.is_remote() { let bearer_token = if spec.is_remote() {
oauth::load_valid_mcp_token(name) oauth::load_or_refresh_mcp_token(name).await
} else { } else {
None None
}; };
+8 -1
View File
@@ -437,6 +437,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 +898,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"
+18 -5
View File
@@ -16,8 +16,9 @@ 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,
@@ -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)
}); });
} }
@@ -3467,6 +3472,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 +4147,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 +4183,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 +4204,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(())
} }
+21
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,6 +24,7 @@ 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};
@@ -495,6 +497,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(
@@ -1252,6 +1264,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()
+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,
}))
}
+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)?;
+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(())
}
+4 -1
View File
@@ -1,3 +1,4 @@
pub(crate) mod manage;
pub(crate) mod oauth; pub(crate) mod oauth;
mod sse_transport; mod sse_transport;
@@ -62,6 +63,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")]
@@ -324,7 +327,7 @@ impl McpRegistry {
.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 bearer_token = if spec.is_remote() {
oauth::load_valid_mcp_token(&id) oauth::load_or_refresh_mcp_token(&id).await
} else { } else {
None None
}; };
+300 -28
View File
@@ -1,15 +1,25 @@
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::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 +44,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 +138,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 +162,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 +175,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 +202,113 @@ 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> { pub async fn load_or_refresh_mcp_token(server_name: &str) -> Option<String> {
let tokens = load_oauth_tokens(&mcp_token_key(server_name))?; let key = mcp_token_key(server_name);
let tokens = load_oauth_tokens(&key)?;
if Utc::now().timestamp() < tokens.expires_at { if Utc::now().timestamp() < tokens.expires_at {
Some(tokens.access_token) return Some(tokens.access_token);
} else { }
None
if in_refresh_failure_backoff(server_name) {
debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed");
return None;
}
let lock = refresh_lock(server_name);
let _guard = lock.lock().await;
// A concurrent caller may have refreshed while we waited for the lock.
let tokens = load_oauth_tokens(&key)?;
if Utc::now().timestamp() < tokens.expires_at {
return Some(tokens.access_token);
}
if in_refresh_failure_backoff(server_name) {
debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed");
return None;
}
match refresh_mcp_token(server_name, &key, &tokens).await {
Ok(access_token) => Some(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)
);
None
}
}
}
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 +322,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 +336,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 +387,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())
} }
@@ -426,11 +574,31 @@ mod tests {
use crate::utils::get_env_name; use crate::utils::get_env_name;
use serial_test::serial; 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) { 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,14 +610,12 @@ 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);
} }
#[test] #[test]
@@ -685,12 +851,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 +881,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 +924,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 +933,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 +952,82 @@ 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_returns_none() {
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 token = rt.block_on(load_or_refresh_mcp_token("legacyref"));
assert_eq!(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);
+77 -38
View File
@@ -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,7 +297,7 @@ 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)?;
let mut paths = doc_paths.to_vec(); let mut paths = doc_paths.to_vec();
if paths.is_empty() { if paths.is_empty() {
@@ -586,19 +565,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 +856,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 +1183,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 +1194,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 +1860,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>> {
@@ -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,
+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();
+1 -1
View File
@@ -892,7 +892,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>"),
}, },
+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]
+106 -20
View File
@@ -337,29 +337,20 @@ fn inject_rag_secrets(vault: &Vault, registered: &HashSet<String>) -> Result<()>
if !data.attached { if !data.attached {
continue; continue;
} }
let Some(placeholder) = data.driver_config.get("api_key") else { let secret_names = driver_config_secret_names(&data);
let Some((primary, extra)) = secret_names.split_first() else {
continue; continue;
}; };
let service_id = mcp_credentials::secret_service_id(&stem);
if service_id.is_empty() || registered.contains(&service_id) {
continue;
}
let secret_name = placeholder
.trim_start_matches("{{")
.trim_end_matches("}}")
.trim();
match vault.get_secret(secret_name, false) { let service_id = mcp_credentials::secret_service_id(&stem);
Ok(secret_value) => { if !service_id.is_empty() && !registered.contains(&service_id) {
sbx_secret_set(&service_id, &secret_value) bind_rag_secret(vault, &service_id, primary, &stem)?;
.context("Failed to register RAG secret with sbx")?; }
}
Err(e) => { for name in extra {
eprintln!( let id = mcp_credentials::secret_service_id(name);
"Warning: could not load secret '{secret_name}' for RAG '{stem}': {e}. \ if !id.is_empty() && !registered.contains(&id) {
Queries to this RAG will fail inside the sandbox. \ bind_rag_secret(vault, &id, name, &stem)?;
Run `coyote --add-secret {secret_name}` to fix."
);
} }
} }
} }
@@ -367,6 +358,43 @@ fn inject_rag_secrets(vault: &Vault, registered: &HashSet<String>) -> Result<()>
Ok(()) Ok(())
} }
fn driver_config_secret_names(data: &RagData) -> Vec<String> {
let mut names: Vec<String> = Vec::new();
for value in data.driver_config.values() {
let trimmed = value.trim();
let Ok(Some(caps)) = SECRET_RE.captures(trimmed) else {
continue;
};
if caps.get(0).map(|m| m.as_str()) != Some(trimmed) {
continue;
}
let Some(name) = caps.get(1).map(|m| m.as_str().trim()) else {
continue;
};
if !name.is_empty() && !names.iter().any(|n| n == name) {
names.push(name.to_string());
}
}
names
}
fn bind_rag_secret(vault: &Vault, service_id: &str, secret_name: &str, stem: &str) -> Result<()> {
match vault.get_secret(secret_name, false) {
Ok(secret_value) => {
sbx_secret_set(service_id, &secret_value)
.context("Failed to register RAG secret with sbx")?;
}
Err(e) => {
eprintln!(
"Warning: could not load secret '{secret_name}' for RAG '{stem}': {e}. \
Queries to this RAG will fail inside the sandbox. \
Run `coyote --add-secret {secret_name}` to fix."
);
}
}
Ok(())
}
fn provider_to_sbx_service(provider_type: &str, client_name: Option<&str>) -> String { fn provider_to_sbx_service(provider_type: &str, client_name: Option<&str>) -> String {
match provider_type { match provider_type {
"claude" => "anthropic".to_string(), "claude" => "anthropic".to_string(),
@@ -646,6 +674,64 @@ fn chown_agent_recursive(sandbox: &str, path: &str) -> Result<()> {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
fn rag_with(driver_config: &[(&str, &str)]) -> RagData {
let mut data = RagData::new("m".into(), 1024, 50, None, 5, None, Default::default());
data.driver = "qdrant".to_string();
data.attached = true;
for (k, v) in driver_config {
data.driver_config.insert(k.to_string(), v.to_string());
}
data
}
#[test]
fn secret_names_are_found_whatever_the_field_is_called() {
let data = rag_with(&[
("host", "qdrant.example.com:6333"),
("collection", "docs"),
("token", "{{SOME_TOKEN}}"),
]);
assert_eq!(driver_config_secret_names(&data), vec!["SOME_TOKEN"]);
}
#[test]
fn a_literal_credential_is_not_treated_as_a_secret_name() {
let data = rag_with(&[("api_key", "sk-a-real-looking-key")]);
assert!(driver_config_secret_names(&data).is_empty());
}
#[test]
fn plain_values_are_never_mistaken_for_secrets() {
let data = rag_with(&[("host", "localhost:6333"), ("collection", "docs")]);
assert!(driver_config_secret_names(&data).is_empty());
}
#[test]
fn a_partial_placeholder_is_not_a_credential() {
let data = rag_with(&[("api_key", "Bearer {{KEY}}")]);
assert!(driver_config_secret_names(&data).is_empty());
}
#[test]
fn several_secrets_are_all_found_and_deduped() {
let data = rag_with(&[
("api_key", "{{QDRANT_KEY}}"),
("host", "localhost:6333"),
("token", "{{ OTHER_TOKEN }}"),
("fallback_key", "{{QDRANT_KEY}}"),
]);
assert_eq!(
driver_config_secret_names(&data),
vec!["QDRANT_KEY", "OTHER_TOKEN"],
"order follows driver_config, and a repeat is not registered twice"
);
}
use std::time::{SystemTime, UNIX_EPOCH}; use std::time::{SystemTime, UNIX_EPOCH};
#[test] #[test]