Compare commits

..
Author SHA1 Message Date
Dark-Alex-17 ae96a7e031 fmt: applied formatting 2026-08-12 19:35:27 -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
Dark-Alex-17 d6c114fe58 feat: simplified the duckdb selection prompt 2026-08-11 16:50:50 -06:00
Dark-Alex-17 912e00a627 docs(rag): correct the duckdb concurrency note in the driver prompt
The driver picker still told users a duckdb RAG can only be open in one Coyote
process at a time. That stopped being true once the store began opening
read-only for queries: any number of processes can now query it concurrently.

The restriction that remains is narrower and only bites while writing, so the
prompt now states that instead — several processes can query at once, but an
ingest or rebuild locks the others out until it finishes.
2026-08-11 16:38:55 -06:00
Dark-Alex-17 e006e29ff1 feat(rag): let several Coyote processes query one duckdb RAG at once
The DuckDB store was always opened read-write, which takes an exclusive file
lock, so a second Coyote process could not even read the RAG. Querying does not
write, and DuckDB permits many concurrent readers as long as no writer is
attached, so the store is now opened read-only whenever it already carries a
complete schema.

Creating or initializing the store still writes, as does rebuilding, so those
paths take the exclusive handle. The rebuild path upgrades a read-only
connection in place, which every clone of the handle observes because the mode
lives behind the shared mutex rather than beside it. Extension loads and the
HNSW persistence setting are per-connection and are re-established on the
upgraded connection.

An upgrade that loses the race for the write lock reports that another process
holds the RAG and that nothing was written, then reopens read-only so the
session can keep querying. A read-only handle also refuses writes outright, so a
missed upgrade cannot silently discard an ingest.
2026-08-11 14:58:18 -06:00
Dark-Alex-17 c0067d387c refactor(rag): drop the hardcoded embedding model hint from attach
The attach wizard mapped a collection's vector dimension to a hardcoded list of
model names and printed them as likely candidates. The list was never checked
against the models the user actually has configured, so it could recommend a
model they cannot select, and one entry was a parenthetical note rather than a
model id and so could never match anything. Any list like this rots as models
are released.

The dimension itself comes from the server and is worth stating, so it is still
printed, as is the warning that a mismatched embedding model returns bad
results. Deriving real candidates would need a dimension recorded against each
configured model, which the model config does not carry today.
2026-08-11 14:06:57 -06:00
Dark-Alex-17 118c346345 feat(rag): support string and UUID Qdrant point IDs
Point ids were read with as_u64() inside a filter_map, so a string id was
silently dropped and a UUID-keyed collection returned zero hits with no error.
The attach wizard therefore refused such collections and told the user to
rebuild with integer ids, which defeats the purpose of attaching to a
collection someone else already built. LangChain, a common way to populate
Qdrant, uses UUIDs by default.

The integer id was never load-bearing for this driver. DocumentId is a packed
(file, chunk) pair used positionally by the local drivers, but an attached RAG
holds no local files or vectors and every positional consumer already returns
early on it, so the id only has to survive the round trip from search back to
the content fetch. Ids that cannot make that trip as a u64 are interned behind
a synthetic handle and restored when the fetch is issued, leaving collections
that already use integer ids on exactly the path they used before.
2026-08-11 14:05:38 -06:00
Dark-Alex-17 dc677a2529 fix(sandbox): discover agent-scoped RAG mixin sidecars
An agent-scoped RAG writes its config to <data>/agents/<agent>/<rag>.yaml, so
its sbx mixin sidecar lands beside it as <rag>.sbx-mixin.yaml. Discovery scanned
the agents directory only for a file named exactly sbx-mixin.yaml, and scanned
for suffixed sidecars only in the top-level rags directory, so a RAG attached
while an agent was active contributed no network allow rule and no credential to
the sandbox. The failure was silent: the sandbox launched and the RAG was simply
unreachable from inside it.

The two collectors differed only in the filename shape they matched, so they are
now one scan that takes the set of layouts to look for. The agents directory
asks for both its own sbx-mixin.yaml and the suffixed sidecars one level in,
which is the shape that was missing. Discovery order is unchanged, and it is
load-bearing: each mixin becomes a --kit in list order and later ones layer over
earlier ones, so the workspace mixin must stay last.
2026-08-11 14:05:31 -06:00
Dark-Alex-17 5e2b9c98ad fix(rag): stop the attach wizard from silently accepting an empty collection
`sample_point_id` returns `None` for a collection with no points, so the
UUID guard's `if let Some(..)` fell straight through and the wizard attached
happily. The result is a RAG that answers every query with zero hits and
never says why.

Sample once, then check for emptiness explicitly. This warns and asks rather
than hard-failing: an empty collection is not necessarily a mistake, since
another tool may be about to populate it, and none of the wizard's remaining
probes can distinguish that from a misconfiguration. The confirmation
defaults to "no" so it cannot be walked past by accident, and `attach`
already refuses to run non-interactively, so no unattended path reaches it.
2026-08-11 13:55:59 -06:00
Dark-Alex-17 c458ca93a9 feat(rag): create the API key secret inline in the attach wizard
The wizard hard-errored with "Secret 'X' not found in vault. Run
`coyote --add-secret X` first.", throwing away every answer the user had
already given it. Offer to create the secret in place instead, deferring to
`Vault::add_secret` for the masked prompt, the provider write and the
confirmation line, then read it back.

Only a genuine `SecretError::NotFound` triggers the offer. An auth failure,
a provider outage, or the vault being disabled inside a sandbox all
propagate with their own message, because prompting for a value that cannot
be stored would fail one step later and bury the real cause. Declining the
offer fails with both ways out spelled: add the secret up front, or answer
"no" to the API-key question.
2026-08-11 13:50:33 -06:00
Dark-Alex-17 860566bf50 feat: let workflow rag nodes select a RAG driver
`RagNode` gains an optional `driver`, forwarded into `RagInitConfig` so a
graph node can build its knowledge base on duckdb instead of yaml. Nodes
that name no driver forward `None`, which still resolves to yaml, so
existing workflows are unaffected.

An unknown driver is rejected up front rather than at construction time.
`Rag::create` dispatches unknown drivers to its yaml catch-all, so a typo
would otherwise embed every document and persist the bogus string, after
which every subsequent load fails validation and the agent cannot start.
The check asks `RagData::validate()` through a probe value instead of
restating the list of valid drivers, so the two cannot drift.
2026-08-11 13:46:59 -06:00
Dark-Alex-17 7f90710427 refactor(rag): interpolate every driver_config value, not just api_key
Only `driver_config["api_key"]` was interpolated, so any credential-bearing
driver field added later would have shipped its raw `{{PLACEHOLDER}}` to the
server. Resolve every value instead, via `resolve_driver_config`.

Resolution still happens into a function-local copy and never touches
`RagData`: `save()` serializes `self.data` and is called by `.set rag_top_k`
and friends, so a resolved credential parked there would be written to the
RAG's YAML in plaintext. The literal `{{NAME}}` also has to survive on disk
because sandbox credential provisioning parses it back out to learn which
vault secret to bind. Scope stays `driver_config` deliberately: the rest of
a RAG file is ingested document text, where `{{...}}` is ordinary content.

Covered by a test that saves after a load and asserts the placeholder, not
the secret, is what reaches the file.
2026-08-11 13:46:32 -06:00
Dark-Alex-17 93a934439b fix(rag): fail loudly when a RAG's vault secret is missing
`interpolate_secrets` does not error on a secret the vault cannot resolve:
it substitutes the empty string and returns the name in its second tuple
element. `load_async` discarded that vec, so a typo'd or deleted vault
secret produced `api_key = ""` and an unexplained 401 from Qdrant.

Bail instead, naming the RAG and the missing secrets, matching what global
config loading already does.
2026-08-11 13:44:00 -06:00
21 changed files with 3068 additions and 272 deletions
+1
View File
@@ -5,3 +5,4 @@
.idea/
/coyote.iml
/.idea/
.coyote
+15
View File
@@ -36,6 +36,21 @@ RUN set -euo pipefail; \
install -m 0755 "$TMPDIR/usql_static" /usr/local/bin/usql; \
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
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`
* **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`
* [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,
etc., and they are used within agents and tools.
+161 -1
View File
@@ -13,6 +13,20 @@ use is_terminal::IsTerminal;
use std::collections::HashSet;
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)]
#[command(author, version, about, long_about = None)]
#[command(
@@ -41,10 +55,15 @@ use std::io::{Read, stdin};
"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 {
/// Input text
#[arg(trailing_var_arg = true)]
#[arg(allow_hyphen_values = true)]
text: Vec<String>,
/// 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))]
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
#[arg(long, value_name = "NAME", help_heading = "Sandbox")]
pub sandbox: Option<Option<String>>,
@@ -254,6 +324,15 @@ pub struct Cli {
/// Generate static shell completion scripts
#[arg(long, value_name = "SHELL", value_enum, help_heading = "Shell")]
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 {
@@ -633,4 +712,85 @@ mod tests {
fn parse_sandbox_is_exclusive() {
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 super::access_token::get_access_token;
@@ -368,17 +369,32 @@ pub fn claude_build_chat_completions_body(
]
} else {
// 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 assistant_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() {
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(
json!({ "role": "assistant", "content": assistant_parts }),
);
messages.push(json!({ "role": "user", "content": user_parts }));
assistant_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 {
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() {
thinking.push(ThinkingBlock::Thinking {
thinking: v.to_string(),
signature: item["signature"]
.as_str()
.unwrap_or_default()
.to_string(),
signature: item["signature"].as_str().unwrap_or_default().to_string(),
});
}
}
@@ -535,3 +548,100 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
};
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);
}
}
+82 -16
View File
@@ -4,6 +4,7 @@ use crate::{
client::Model,
config::memory,
function::{Functions, run_llm_function},
graph, rag,
};
use super::rag_cache::RagKey;
@@ -12,6 +13,7 @@ use crate::config::prompts::{
DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS,
DEFAULT_USER_INTERACTION_INSTRUCTIONS,
};
use crate::graph::types::RagNode;
use crate::graph::{Graph, GraphParser, NodeType};
use crate::rag::RagInitConfig;
use crate::vault::SECRET_RE;
@@ -184,7 +186,7 @@ impl Agent {
&rag_path_clone,
&document_paths,
abort,
false,
true,
)
.await
})
@@ -246,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);
Ok(Self {
@@ -952,6 +958,30 @@ fn resolve_document_paths(
Ok(document_paths)
}
/// How a graph rag node describes the knowledge base it wants built.
///
/// `driver` is forwarded as-is: `None` means the node did not ask for one, which
/// `RagInitConfig` resolves to yaml, so workflows written before drivers existed
/// keep their current storage.
///
/// Every field is now named explicitly, so adding one to `RagInitConfig` breaks
/// this literal. That is deliberate: the new field then gets a decision about
/// whether a rag node can drive it, instead of silently taking its default.
fn rag_init_config(rag_node: &RagNode) -> RagInitConfig {
RagInitConfig {
embedding_model: rag_node.embedding_model.clone(),
chunk_size: rag_node.chunk_size,
chunk_overlap: rag_node.chunk_overlap,
reranker_model: rag_node.reranker_model.clone(),
top_k: rag_node.top_k,
batch_size: rag_node.batch_size,
extractor_model: rag_node.extractor_model.clone(),
extractor_prompt: rag_node.extractor_prompt.clone(),
graph_hops: rag_node.graph_hops,
driver: rag_node.driver.clone(),
}
}
#[allow(clippy::too_many_arguments)]
async fn init_graph_rags(
app: &AppConfig,
@@ -989,21 +1019,18 @@ async fn init_graph_rags(
})
.await?
} else {
let config = RagInitConfig {
embedding_model: rag_node.embedding_model.clone(),
chunk_size: rag_node.chunk_size,
chunk_overlap: rag_node.chunk_overlap,
reranker_model: rag_node.reranker_model.clone(),
top_k: rag_node.top_k,
batch_size: rag_node.batch_size,
extractor_model: rag_node.extractor_model.clone(),
extractor_prompt: rag_node.extractor_prompt.clone(),
graph_hops: rag_node.graph_hops,
// Graph-node RAGs are yaml-only: `RagNode` has no `driver` field, so
// there is nothing to forward. The rest-pattern also keeps this literal
// from breaking on future `RagInitConfig` additions.
..Default::default()
};
// Checked before anything is built: an unknown driver would otherwise
// fall through `Rag::create`'s catch-all to a yaml store, embed every
// document, and persist the bogus driver string. The RAG would then be
// rejected on every subsequent load, leaving the agent unstartable.
// Graph validation catches this too, but it is skipped when
// `validate_before_run` is off, so this guard is the load-bearing one.
if let Some(driver) = &rag_node.driver
&& let Some(message) = graph::validator::rag_driver_error(driver)
{
bail!("rag node '{node_id}': {message}");
}
let mut config = rag_init_config(rag_node);
let fully_specified = config.embedding_model.is_some()
&& config.chunk_size.is_some()
&& config.chunk_overlap.is_some();
@@ -1029,6 +1056,10 @@ async fn init_graph_rags(
initialized. RAG initialization is required for this agent."
);
}
if config.driver.is_none() {
config.driver = Some(rag::select_rag_driver()?);
}
}
let document_paths =
@@ -1337,4 +1368,39 @@ version: "1.0"
assert_eq!(meta.description, "");
}
#[test]
fn rag_init_config_forwards_an_explicit_driver() {
let node: RagNode =
serde_yaml::from_str("documents: [\"./docs\"]\ndriver: duckdb\n").unwrap();
assert_eq!(rag_init_config(&node).driver.as_deref(), Some("duckdb"));
}
/// A node that names no driver must forward `None`, which `RagInitConfig`
/// documents as "yaml". Existing workflows therefore keep their yaml store.
#[test]
fn rag_init_config_leaves_the_driver_unset_by_default() {
let node: RagNode = serde_yaml::from_str("documents: [\"./docs\"]\n").unwrap();
assert_eq!(rag_init_config(&node).driver, None);
}
/// The driver must ride alongside the rest of the node's settings, not
/// replace them.
#[test]
fn rag_init_config_forwards_the_other_settings_too() {
let node: RagNode = serde_yaml::from_str(
"documents: [\"./docs\"]\ndriver: duckdb\nchunk_size: 512\nchunk_overlap: 64\ntop_k: 7\nembedding_model: some:model\n",
)
.unwrap();
let config = rag_init_config(&node);
assert_eq!(config.driver.as_deref(), Some("duckdb"));
assert_eq!(config.chunk_size, Some(512));
assert_eq!(config.chunk_overlap, Some(64));
assert_eq!(config.top_k, Some(7));
assert_eq!(config.embedding_model.as_deref(), Some("some:model"));
}
}
+8 -1
View File
@@ -437,6 +437,10 @@ pub(crate) fn remove_rag_sidecars(dir: &Path, name: &str) -> Result<()> {
if duckdb_path.exists() {
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"));
if mixin_path.exists() {
remove_file(&mixin_path).with_context(|| {
@@ -894,16 +898,19 @@ mod tests {
}
#[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 duckdb = root.join("docs.duckdb");
let wal = root.join("docs.duckdb.wal");
let mixin = root.join("docs.sbx-mixin.yaml");
fs::write(&duckdb, "db").unwrap();
fs::write(&wal, "wal").unwrap();
fs::write(&mixin, "mixin").unwrap();
remove_rag_sidecars(&root, "docs").unwrap();
assert!(!duckdb.exists(), "the .duckdb sidecar must be removed");
assert!(!wal.exists(), "the .duckdb.wal sidecar must be removed");
assert!(
!mixin.exists(),
"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::function::{
FunctionDeclaration, Functions, ToolCallTracker, ToolResult, memory::MEMORY_FUNCTION_PREFIX,
skill::SKILL_FUNCTION_PREFIX, supervisor::SUPERVISOR_FUNCTION_PREFIX,
todo::TODO_FUNCTION_PREFIX, user_interaction::USER_FUNCTION_PREFIX,
rag_query::RAG_FUNCTION_PREFIX, skill::SKILL_FUNCTION_PREFIX,
supervisor::SUPERVISOR_FUNCTION_PREFIX, todo::TODO_FUNCTION_PREFIX,
user_interaction::USER_FUNCTION_PREFIX,
};
use crate::mcp::{
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<()> {
self.rag.take();
self.tool_scope.functions.remove_rag_query_functions();
Ok(())
}
@@ -1137,6 +1139,7 @@ impl RequestContext {
&& !v.name.starts_with("agent__")
&& !v.name.starts_with("memory__")
&& !v.name.starts_with("skill__")
&& !v.name.starts_with("rag__")
})
.map(|v| v.name.clone())
.collect()
@@ -1957,7 +1960,8 @@ impl RequestContext {
|| (!matches!(role.skills_enabled(), Some(false))
&& v.name.starts_with(SKILL_FUNCTION_PREFIX))
|| (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)
})
.cloned()
@@ -1987,6 +1991,7 @@ impl RequestContext {
|| v.name.starts_with(TODO_FUNCTION_PREFIX)
|| v.name.starts_with(SUPERVISOR_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() {
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();
self.tool_scope = ToolScope {
@@ -4136,7 +4147,7 @@ impl RequestContext {
super::TEMP_RAG_NAME,
&rag_path,
&[],
abort_signal,
abort_signal.clone(),
false,
)
.await?,
@@ -4172,10 +4183,11 @@ impl RequestContext {
};
self.rag = Some(rag);
self.rag_key = rag_key;
self.refresh_tool_scope(abort_signal).await?;
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);
if rag_path.exists() {
bail!(
@@ -4192,6 +4204,7 @@ impl RequestContext {
self.rag_cache().insert(key.clone(), &rag);
self.rag = Some(rag);
self.rag_key = Some(key);
self.refresh_tool_scope(abort_signal).await?;
Ok(())
}
+21
View File
@@ -1,4 +1,5 @@
pub(crate) mod memory;
pub(crate) mod rag_query;
pub(crate) mod skill;
pub(crate) mod supervisor;
pub(crate) mod todo;
@@ -23,6 +24,7 @@ use futures_util::future;
use indexmap::IndexMap;
use indoc::formatdoc;
use memory::MEMORY_FUNCTION_PREFIX;
use rag_query::RAG_FUNCTION_PREFIX;
use rust_embed::Embed;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
@@ -495,6 +497,16 @@ impl Functions {
.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>) {
let mut invoke_function_properties = IndexMap::new();
invoke_function_properties.insert(
@@ -1252,6 +1264,15 @@ impl ToolCall {
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) {
Ok(Some(contents)) => serde_json::from_str(&contents)
.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,
}))
}
+103
View File
@@ -367,6 +367,13 @@ pub struct RagNode {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub graph_hops: Option<usize>,
/// Storage driver for this node's knowledge base ("yaml", "duckdb"). `None`
/// means "yaml". Only honored when the knowledge base is first built;
/// changing it afterwards has no effect until the RAG is deleted and
/// re-initialized.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub driver: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub state_updates: Option<HashMap<String, String>>,
@@ -1152,4 +1159,100 @@ nodes:
assert!(triage.next.as_ref().unwrap().is_fan_out());
assert_eq!(triage.next.as_ref().unwrap().as_slice().len(), 2);
}
fn rag_node_of(graph: &Graph, id: &str) -> RagNode {
match &graph.get_node(id).unwrap().node_type {
NodeType::Rag(r) => r.clone(),
other => panic!("expected a rag node, got {other:?}"),
}
}
#[test]
fn rag_node_deserializes_an_explicit_driver() {
let yaml = r#"
name: kb
start: research
nodes:
research:
type: rag
documents: ["./docs"]
driver: duckdb
next: done
done:
type: end
output: ok
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
assert_eq!(
rag_node_of(&graph, "research").driver.as_deref(),
Some("duckdb")
);
}
/// Workflows written before drivers existed must keep parsing, and must keep
/// asking for nothing, so `RagInitConfig` resolves them to the yaml default.
#[test]
fn rag_node_without_a_driver_stays_unset() {
let yaml = r#"
name: kb
start: research
nodes:
research:
type: rag
documents: ["./docs"]
next: done
done:
type: end
output: ok
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
assert_eq!(rag_node_of(&graph, "research").driver, None);
}
#[test]
fn rag_node_driver_survives_a_serialize_round_trip() {
let yaml = r#"
name: kb
start: research
nodes:
research:
type: rag
documents: ["./docs"]
driver: duckdb
next: done
done:
type: end
output: ok
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
let reparsed: Graph =
serde_yaml::from_str(&serde_yaml::to_string(&graph).unwrap()).unwrap();
assert_eq!(
rag_node_of(&reparsed, "research").driver.as_deref(),
Some("duckdb")
);
}
/// `skip_serializing_if` must keep `driver:` out of graphs that never set it.
#[test]
fn rag_node_without_a_driver_omits_the_key_when_serialized() {
let yaml = r#"
name: kb
start: research
nodes:
research:
type: rag
documents: ["./docs"]
next: done
done:
type: end
output: ok
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
assert!(!serde_yaml::to_string(&graph).unwrap().contains("driver"));
}
}
+101
View File
@@ -2,6 +2,7 @@ use super::state::template_root_keys;
use super::types::{Graph, Node, NodeType};
use crate::client::{Model, ModelType};
use crate::config::{Agent, AppConfig, paths};
use crate::rag::{GraphRagConfig, RagData};
use anyhow::{Result, bail};
use std::collections::{BTreeMap, HashSet, VecDeque};
use std::path::PathBuf;
@@ -96,6 +97,51 @@ pub struct GraphValidator {
skill_exists: fn(&str) -> bool,
}
/// A minimal `RagData` whose only interesting field is `driver`. The numeric
/// arguments are the smallest values that satisfy `validate()`'s unrelated
/// floors (top_k >= 1, and chunk_size >= 1 with chunk_overlap < chunk_size for
/// a non-attached RAG). `RagData::new` sets `attached: false`, which is the
/// correct shape here: a graph rag node always builds its own local knowledge
/// base from `documents` and can never be attached.
fn rag_driver_probe(driver: &str) -> RagData {
let mut data = RagData::new(
String::new(),
1,
0,
None,
1,
None,
GraphRagConfig::default(),
);
data.driver = driver.to_string();
data
}
/// `Some(message)` when `driver` is one that `RagData::validate()` would reject.
///
/// The set of valid drivers is defined in exactly one place, `RagData::validate()`,
/// so this asks that function rather than restating the list here.
///
/// Fails open on purpose: the first probe below uses the default driver, which is
/// valid by definition. If even that one is rejected, `validate()` has grown a
/// precondition the probe fixture no longer satisfies, and every verdict from here
/// would be a false positive that rejects working graphs. In that case we decline
/// to judge and leave enforcement to RAG construction. The
/// `rag_driver_probe_fixture_is_accepted` test turns that silent degradation into a
/// loud failure. Both `validate()` calls are load-bearing; neither is redundant.
pub(crate) fn rag_driver_error(driver: &str) -> Option<String> {
if rag_driver_probe(&RagData::default().driver)
.validate()
.is_err()
{
return None;
}
rag_driver_probe(driver)
.validate()
.err()
.map(|err| err.to_string())
}
impl GraphValidator {
pub fn new(base_dir: impl Into<PathBuf>) -> Self {
Self {
@@ -154,6 +200,11 @@ impl GraphValidator {
not be written to state",
));
}
if let Some(driver) = &r.driver
&& let Some(message) = rag_driver_error(driver)
{
result.error(ValidationError::with_node(node_id, message));
}
}
}
}
@@ -1031,6 +1082,7 @@ mod tests {
extractor_model: None,
extractor_prompt: None,
graph_hops: None,
driver: None,
state_updates,
timeout: None,
}),
@@ -1385,6 +1437,55 @@ mod tests {
);
}
/// Guards the fail-open branch in `rag_driver_error`. If this fails,
/// `RagData::validate()` grew a precondition the probe fixture no longer
/// satisfies and rag-node driver validation has silently switched itself off.
/// Repair the fixture in `rag_driver_probe`; do not delete this test.
#[test]
fn rag_driver_probe_fixture_is_accepted() {
let default_driver = RagData::default().driver;
assert!(
rag_driver_probe(&default_driver).validate().is_ok(),
"probe fixture rejected for the default driver '{default_driver}'"
);
}
#[test]
fn rag_driver_error_defers_to_ragdata_validate() {
assert_eq!(rag_driver_error("yaml"), None);
assert_eq!(rag_driver_error("duckdb"), None);
let message = rag_driver_error("duckdbb").expect("unknown driver must be rejected");
assert!(message.contains("duckdbb"), "got: {message}");
}
#[test]
fn rag_node_with_unknown_driver_errors_naming_the_node() {
let mut node = rag_node("kb", &["./docs"], true);
if let NodeType::Rag(ref mut r) = node.node_type {
r.driver = Some("postgres".into());
}
let graph = graph_with(vec![("kb", node), ("end", end_node("end"))], "kb");
let result = validator().validate(&graph);
assert!(!result.is_valid());
let err = result.into_result().unwrap_err().to_string();
assert!(err.contains("[kb]"), "must name the node: {err}");
assert!(err.contains("postgres"), "must name the driver: {err}");
}
#[test]
fn rag_node_with_duckdb_driver_produces_no_findings() {
let mut node = rag_node("kb", &["./docs"], true);
if let NodeType::Rag(ref mut r) = node.node_type {
r.driver = Some("duckdb".into());
}
let graph = graph_with(vec![("kb", node), ("end", end_node("end"))], "kb");
assert!(validator().validate(&graph).is_valid());
}
fn agent_node(id: &str, agent: &str, next: Option<&str>) -> Node {
Node {
id: id.into(),
+12
View File
@@ -196,6 +196,18 @@ async fn main() -> Result<()> {
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 {
let cfg = Config::load_with_interpolation(true).await?;
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(())
}
+3
View File
@@ -1,3 +1,4 @@
pub(crate) mod manage;
pub(crate) mod oauth;
mod sse_transport;
@@ -62,6 +63,8 @@ pub(crate) struct McpServersConfig {
pub(crate) struct McpOAuthConfig {
#[serde(rename = "clientId", skip_serializing_if = "Option::is_none")]
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")]
pub callback_port: Option<u16>,
#[serde(rename = "redirectHost", skip_serializing_if = "Option::is_none")]
+360 -89
View File
@@ -18,6 +18,7 @@ use crate::vault::{Vault, interpolate_secrets};
use anyhow::{Context, Result, anyhow, bail};
use bm25::{Language, SearchEngine, SearchEngineBuilder};
use gman::SecretError;
use hnsw_rs::prelude::*;
use indexmap::{IndexMap, IndexSet};
use inquire::{Confirm, Select, Text, required, validator::Validation};
@@ -266,29 +267,10 @@ impl Rag {
}
println!("⚙ Initializing RAG...");
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 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. Can only be open in ONE Coyote process at a time, and its driver cannot be changed later without recreating the RAG",
];
let sel = Select::new("RAG storage driver:", options)
.with_starting_cursor(0)
.prompt()?;
if sel.starts_with("duckdb") {
println!(
"Note: a duckdb RAG can only be open in one Coyote process at a time, \
and changing its driver later means deleting and recreating the RAG."
);
"duckdb"
} else {
"yaml"
}
select_rag_driver()?
} else {
"yaml"
"yaml".to_string()
};
let reranker_model = app.rag_reranker_model.clone();
let top_k = app.rag_top_k;
@@ -315,7 +297,7 @@ impl Rag {
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 paths = doc_paths.to_vec();
if paths.is_empty() {
@@ -355,38 +337,27 @@ impl Rag {
let raw_content = fs::read_to_string(path).with_context(err)?;
// Parsed WITHOUT secret interpolation, so `driver_config` keeps its
// `{{...}}` placeholders. Interpolating here would bake the resolved API
// key into `self.data`, which `save()` then writes back to disk in
// plaintext.
// `{{...}}` placeholders in `self.data`. Resolution happens below, into a
// function-local copy only — see `resolve_driver_config` for why the
// resolved values must never travel back into `data`.
let data: RagData = serde_yaml::from_str(&raw_content).with_context(err)?;
data.validate().with_context(err)?;
match data.driver.as_str() {
"qdrant" => {
let host = data
.driver_config
let driver_config = resolve_driver_config(&data.driver_config, vault, name)?;
let host = driver_config
.get("host")
.context("qdrant driver requires 'host' in driver_config")?
.clone();
let collection = data
.driver_config
let collection = driver_config
.get("collection")
.context("qdrant driver requires 'collection' in driver_config")?
.clone();
let api_key = driver_config.get("api_key").map(String::as_str);
let api_key: Option<String> = match data.driver_config.get("api_key") {
Some(placeholder) => {
let (resolved, _) =
interpolate_secrets(placeholder, vault).with_context(|| {
format!("Failed to resolve api_key secret for RAG '{name}'")
})?;
Some(resolved)
}
None => None,
};
let provider = QdrantProvider::new(&host, &collection, api_key.as_deref()).await?;
let provider = QdrantProvider::new(&host, &collection, api_key).await?;
let embedding_model =
Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?;
Ok(Rag {
@@ -440,12 +411,7 @@ impl Rag {
.with_default("QDRANT_API_KEY")
.with_validator(required!("This field is required"))
.prompt()?;
let resolved = vault.get_secret(&secret_name, false).with_context(|| {
format!(
"Secret '{secret_name}' not found in vault. \
Run `coyote --add-secret {secret_name}` first."
)
})?;
let resolved = resolve_or_create_api_key_secret(vault, &secret_name)?;
Some((secret_name, resolved))
} else {
None
@@ -468,18 +434,28 @@ impl Rag {
let collection = Select::new("Select collection:", collections).prompt()?;
// Point IDs are read with `as_u64()`, which yields None for a JSON string.
// A UUID-keyed collection would therefore return zero hits with no error,
// so refuse it here instead of attaching something silently broken.
if let Some(raw_id) = QdrantProvider::sample_point_id(&host, &collection, api_key).await?
&& raw_id.starts_with('"')
{
bail!(
"Collection '{collection}' uses string (UUID) point IDs. \
Coyote requires integer point IDs. Rebuild the collection with integer IDs \
(e.g. LangChain: pass ids=list(range(len(docs))) to add_documents())."
let sample_id = QdrantProvider::sample_point_id(&host, &collection, api_key).await?;
// `None` means the scroll came back with no points at all: the collection
// is empty. Attaching is not necessarily wrong — another tool may be about
// to fill it — but accepting it silently yields a RAG that answers every
// query with nothing and never explains why, and none of the checks below
// can tell that apart from a misconfiguration. Ask, defaulting to no, so it
// cannot happen by accident. (`attach` already refuses to run
// non-interactively, so there is no unattended path through this prompt.)
if sample_id.is_none() {
println!(
"⚠️ Collection '{collection}' contains no points. Queries will return \
nothing until something writes to it."
);
let attach_anyway = Confirm::new("Attach to this empty collection anyway?")
.with_default(false)
.prompt()?;
if !attach_anyway {
bail!("Collection '{collection}' is empty; nothing to attach to.");
}
}
println!(" This collection must store document text in a 'page_content' payload field.");
let dim = QdrantProvider::get_vector_dimension(&host, &collection, api_key)
@@ -496,13 +472,7 @@ impl Rag {
);
}
if dim > 0 {
let candidates = embedding_model_candidates_for_dimension(dim);
if !candidates.is_empty() {
println!(
"Collection uses {dim}-dim vectors. Likely models: {}",
candidates.join(", ")
);
}
println!("Collection uses {dim}-dim vectors.");
}
println!(
"⚠️ If the embedding model doesn't match what built this collection, \
@@ -595,19 +565,40 @@ impl Rag {
if data.vectors.is_empty() {
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
// path; there is no from-DuckDB fallback.
let bm25 = data.build_bm25();
(Box::new(duck), bm25)
}
"qdrant" => bail!(
"Qdrant RAGs cannot be constructed via Rag::create(); \
use Rag::attach() or Rag::load_async() instead"
"RAG '{name}' uses driver 'qdrant' without `attached: true`. \
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();
(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 embedding_model =
@@ -865,6 +856,20 @@ impl Rag {
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(
&self,
app: &AppConfig,
@@ -1178,12 +1183,7 @@ impl Rag {
top_k: usize,
rerank_model: Option<&str>,
) -> Result<Vec<(DocumentId, String)>> {
let vector_search_results = self.vector_search(query, top_k, 0.0).await?;
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)> =
let keyword_leg = async {
if self.provider.has_native_keyword_search() {
self.provider
.keyword_search(query, top_k)
@@ -1194,7 +1194,16 @@ impl Rag {
})
} else {
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:?}",);
let keyword_search_ids: Vec<DocumentId> =
keyword_search_results.into_iter().map(|(v, _)| v).collect();
@@ -1839,21 +1848,6 @@ fn driver_auth_header(driver: &str) -> (&'static str, &'static str) {
}
}
/// Embedding models known to produce a given vector dimension, used to hint the
/// user toward a model compatible with the collection they just picked.
fn embedding_model_candidates_for_dimension(dim: u64) -> Vec<&'static str> {
match dim {
1536 => vec!["text-embedding-3-small", "text-embedding-ada-002"],
3072 => vec!["text-embedding-3-large"],
768 => vec!["nomic-embed-text", "all-minilm-l6-v2"],
1024 => vec![
"text-embedding-3-small (matryoshka-1024)",
"jina-embeddings-v2-base",
],
_ => vec![],
}
}
fn select_embedding_model(models: &[&Model]) -> Result<String> {
let max_width = models.iter().map(|v| v.id().len()).max().unwrap_or(0);
let models: Vec<_> = models
@@ -1866,6 +1860,27 @@ fn select_embedding_model(models: &[&Model]) -> Result<String> {
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";
fn select_extractor_model(app: &AppConfig) -> Result<Option<String>> {
@@ -2121,6 +2136,123 @@ fn embedding_dim_for_model(model_id: &str) -> usize {
}
}
/// True only for "the vault does not hold this key".
///
/// Everything else — an auth failure, a provider outage, or the vault being
/// disabled because Coyote is running inside a sandbox — must NOT be treated as
/// a missing secret. Offering to create one in those cases would prompt for a
/// value that cannot be stored and bury the real reason.
fn is_missing_secret(err: &anyhow::Error) -> bool {
matches!(
err.downcast_ref::<SecretError>(),
Some(SecretError::NotFound { .. })
)
}
/// Reads `secret_name` out of the vault, offering to create it in place when the
/// vault simply does not hold it yet.
///
/// Sending the user off to run `coyote --add-secret` mid-wizard discarded every
/// answer they had already given. `Vault::add_secret` does the masked prompt,
/// the provider write and the confirmation line, so this defers to it rather
/// than collecting or storing the value itself.
fn resolve_or_create_api_key_secret(vault: &Vault, secret_name: &str) -> Result<String> {
let read_err = match vault.get_secret(secret_name, false) {
Ok(secret) => return Ok(secret),
Err(err) => err,
};
if !is_missing_secret(&read_err) {
return Err(read_err)
.with_context(|| format!("Cannot read secret '{secret_name}' from the vault"));
}
let create = Confirm::new(&format!(
"Secret '{secret_name}' is not in the vault. Create it now?"
))
.with_default(true)
.prompt()?;
if !create {
bail!(
"This instance needs an API key, so '{secret_name}' has to exist before \
attaching. Add it with `coyote --add-secret {secret_name}` and re-run, or \
re-run and answer 'no' when asked whether the instance requires an API key."
);
}
vault
.add_secret(secret_name)
.with_context(|| format!("Failed to store secret '{secret_name}' in the vault"))?;
vault
.get_secret(secret_name, false)
.with_context(|| format!("Secret '{secret_name}' is unreadable after being stored"))
}
/// Resolves `{{SECRET}}` placeholders in every `driver_config` value against the
/// vault, returning a DETACHED copy.
///
/// Three properties this must preserve, each of which has already bitten:
///
/// 1. The resolved values never go back into `RagData`. `Rag::save()`
/// serializes `self.data`, and `.set rag_top_k`, `.set rag_reranker_model`
/// and every post-sync save call it — so a resolved credential parked in
/// `data.driver_config` gets written to the RAG's YAML file in plaintext the
/// next time the user changes any setting.
/// 2. The literal `{{NAME}}` text survives in `data` and on disk. Sandbox
/// credential provisioning parses that placeholder back out of the file to
/// learn which vault secret to bind into the sandbox; resolve it away and
/// provisioning silently finds nothing to register.
/// 3. Only `driver_config` is interpolated, never the whole file. The rest of a
/// RAG file is ingested document text and vectors — where `{{...}}` is
/// ordinary content (Jinja, Mustache, Vue, Go templates) that would be read
/// as a secret reference, blanked to `""`, and persisted on the next save.
/// `driver_config` is small and is the only place credentials live.
fn resolve_driver_config(
driver_config: &IndexMap<String, String>,
vault: &Vault,
rag_name: &str,
) -> Result<IndexMap<String, String>> {
resolve_driver_config_with(driver_config, rag_name, |value| {
interpolate_secrets(value, vault)
})
}
/// Interpolation core, taking the resolver as an argument so it can be exercised
/// without a vault. Mirrors `interpolate_secrets` / `interpolate_secrets_with`.
fn resolve_driver_config_with<F>(
driver_config: &IndexMap<String, String>,
rag_name: &str,
mut interpolate: F,
) -> Result<IndexMap<String, String>>
where
F: FnMut(&str) -> Result<(String, Vec<String>)>,
{
let mut resolved = IndexMap::with_capacity(driver_config.len());
let mut missing: Vec<String> = Vec::new();
for (key, value) in driver_config {
let (value, value_missing) = interpolate(value).with_context(|| {
format!("Failed to resolve '{key}' in driver_config for RAG '{rag_name}'")
})?;
missing.extend(value_missing);
resolved.insert(key.clone(), value);
}
// A secret the vault does not hold is NOT an error inside
// `interpolate_secrets`: it substitutes the empty string and only reports the
// name. Accepting that ships an empty credential, and the user sees an
// unexplained 401 from the server instead of the typo they made.
if !missing.is_empty() {
missing.sort();
missing.dedup();
bail!(
"RAG '{rag_name}' references secrets that are missing from the vault: {}. \
Add them with `coyote --add-secret <name>`, then try again.",
missing.join(", ")
);
}
Ok(resolved)
}
#[cfg(test)]
mod tests {
use super::*;
@@ -2171,6 +2303,145 @@ mod tests {
assert!(yaml.contains("{{QDRANT_API_KEY}}"));
}
const FAKE_SECRET: &str = "sk-live-fake-value-for-tests";
fn attached_qdrant_data() -> RagData {
let mut data = RagData {
driver: "qdrant".to_string(),
attached: true,
embedding_model: "text-embedding-3-small".to_string(),
top_k: 5,
..Default::default()
};
data.driver_config
.insert("host".into(), "localhost:6333".into());
data.driver_config.insert("collection".into(), "c".into());
data.driver_config
.insert("api_key".into(), "{{QDRANT_API_KEY}}".into());
data
}
/// THE invariant behind `resolve_driver_config` returning a detached copy.
///
/// `save()` serializes `self.data`, and `.set rag_top_k`, `.set
/// rag_reranker_model` and every post-sync save call it. If load ever bakes
/// the resolved credential into `data.driver_config`, the next trivial
/// setting change writes the user's plaintext API key into the RAG's YAML
/// file. The literal placeholder must also survive, because sandbox
/// credential provisioning parses it back off disk.
#[test]
fn a_save_after_load_writes_the_placeholder_not_the_resolved_secret() {
let dir = TempDir::new("driver-config-secret");
let path = dir.path.join("kb.yaml");
let data = attached_qdrant_data();
// Exactly what `load_async` does with the parsed data.
let resolved = resolve_driver_config_with(&data.driver_config, "kb", |value| {
Ok((value.replace("{{QDRANT_API_KEY}}", FAKE_SECRET), vec![]))
})
.unwrap();
assert_eq!(
resolved["api_key"], FAKE_SECRET,
"the live client still has to receive the real key"
);
assert_eq!(
data.driver_config["api_key"], "{{QDRANT_API_KEY}}",
"resolution must not mutate the RagData that save() serializes"
);
let rag = Rag {
app_config: Arc::new(AppConfig::default()),
name: "kb".to_string(),
path: path.display().to_string(),
embedding_model: Model::new("openai", "text-embedding-3-small"),
bm25: data.build_bm25(),
provider: Box::new(YamlProvider::from_data(&data)),
node_to_docs: IndexMap::new(),
data,
last_sources: RwLock::new(None),
};
assert!(rag.save().unwrap());
let on_disk = fs::read_to_string(&path).unwrap();
assert!(
on_disk.contains("{{QDRANT_API_KEY}}"),
"sandbox provisioning parses this placeholder back off disk: {on_disk}"
);
assert!(
!on_disk.contains(FAKE_SECRET),
"a save after load leaked the plaintext secret to {}",
path.display()
);
}
/// Every value is interpolated, not just `api_key` — a credential-bearing
/// field added later must not ship its raw placeholder to the server.
#[test]
fn resolution_covers_every_driver_config_value() {
let mut driver_config = IndexMap::new();
driver_config.insert("host".to_string(), "{{QDRANT_HOST}}".to_string());
driver_config.insert("collection".to_string(), "c".to_string());
driver_config.insert("api_key".to_string(), "{{QDRANT_API_KEY}}".to_string());
let resolved = resolve_driver_config_with(&driver_config, "kb", |value| {
let out = value
.replace("{{QDRANT_HOST}}", "qdrant.internal:6333")
.replace("{{QDRANT_API_KEY}}", FAKE_SECRET);
Ok((out, vec![]))
})
.unwrap();
assert_eq!(resolved["host"], "qdrant.internal:6333");
assert_eq!(resolved["collection"], "c");
assert_eq!(resolved["api_key"], FAKE_SECRET);
}
/// Missing secrets are reported together, deduplicated, and name the RAG.
#[test]
fn missing_secrets_fail_the_load_instead_of_resolving_to_empty() {
let mut driver_config = IndexMap::new();
driver_config.insert("host".to_string(), "{{QDRANT_HOST}}".to_string());
driver_config.insert("api_key".to_string(), "{{QDRANT_API_KEY}}".to_string());
let err = resolve_driver_config_with(&driver_config, "kb", |value| {
// What `interpolate_secrets` really does for an absent secret: blank it
// out and report the name rather than returning Err.
Ok((
String::new(),
vec![value.trim_matches(['{', '}']).to_string()],
))
})
.expect_err("an empty API key must not be accepted as a successful load");
let msg = err.to_string();
assert!(msg.contains("kb"), "the RAG must be named: {msg}");
assert!(msg.contains("QDRANT_HOST"), "got: {msg}");
assert!(msg.contains("QDRANT_API_KEY"), "got: {msg}");
}
/// Only a genuine NotFound may trigger the attach wizard's "create it now?"
/// offer. The vault is disabled wholesale inside a sandbox, where creating a
/// secret is impossible — misreading that as "missing" would prompt for a
/// value that cannot be stored and hide why.
#[test]
fn only_a_not_found_error_counts_as_a_missing_secret() {
let not_found = anyhow::Error::new(SecretError::NotFound {
key: "QDRANT_API_KEY".to_string(),
provider: "local",
});
assert!(is_missing_secret(&not_found));
let auth_failed = anyhow::Error::new(SecretError::AuthFailed {
provider: "local",
source: anyhow!("bad vault password"),
});
assert!(!is_missing_secret(&auth_failed));
// What `Vault::get_secret` returns in sandbox mode: a plain anyhow error.
let sandboxed = anyhow!("Vault management is disabled in sandbox mode.");
assert!(!is_missing_secret(&sandboxed));
}
/// A qdrant RAG's vectors MUST survive serialization.
///
/// `save()` omits vectors only for `driver == "duckdb"`. Qdrant must not join
+561 -51
View File
@@ -4,8 +4,8 @@ use std::collections::HashMap;
use anyhow::{Context, Result, anyhow, bail};
use async_trait::async_trait;
use duckdb::Connection;
use duckdb::types::Value;
use duckdb::{AccessMode, Config, Connection};
use indexmap::IndexMap;
use log::warn;
use std::path::{Path, PathBuf};
@@ -17,8 +17,12 @@ use std::sync::{Arc, Mutex, MutexGuard};
/// `~/.duckdb/extensions/...`. Two threads installing the same extension at once
/// both perform that move; on Windows the loser's move targets a file the winner
/// already holds open and fails with "Access is denied", where POSIX would let the
/// replacement through. Guards nothing but the install step, so it is never held
/// across a `DuckDbProvider::conn` guard and cannot invert lock order.
/// replacement through. Guards nothing but the install step.
///
/// Lock order is `DuckDbProvider::conn` -> INSTALL_LOCK, never the reverse:
/// `ensure_writable` reopens the connection, and so may install, while holding the
/// `ConnHandle` guard, whereas nothing ever acquires a `ConnHandle` guard while
/// holding this lock. The cycle that would deadlock cannot form.
static INSTALL_LOCK: Mutex<()> = Mutex::new(());
/// Derive the DuckDB sidecar path from a RAG's YAML path: `docs.yaml` -> `docs.duckdb`.
@@ -26,9 +30,44 @@ pub(crate) fn duckdb_path_from_yaml(yaml_path: &Path) -> PathBuf {
yaml_path.with_extension("duckdb")
}
/// The shared connection together with the access mode it was opened with.
///
/// `conn` is an `Option` only so that an upgrade can DROP the read-only connection
/// before asking DuckDB for a read-write one. It is `Some` at every point an outside
/// caller can observe, and is never left `None` on a path that returns `Ok`.
struct ConnHandle {
conn: Option<Connection>,
/// True when `conn` was opened READ_WRITE. This lives behind the same mutex as the
/// connection itself rather than next to it in `DuckDbProvider`, so that a
/// `duplicate()` clone sharing the `Arc` observes an upgrade performed through any
/// other handle instead of keeping its own stale copy of the mode.
writable: bool,
}
impl ConnHandle {
fn conn(&self) -> Result<&Connection> {
self.conn.as_ref().ok_or_else(Self::lost)
}
fn conn_mut(&mut self) -> Result<&mut Connection> {
self.conn.as_mut().ok_or_else(Self::lost)
}
/// Only reachable when a read-write upgrade failed AND reopening read-only failed
/// too. Returning an error beats panicking inside a locked scope, which would
/// poison the mutex for the remaining life of the process.
fn lost() -> anyhow::Error {
anyhow!(
"The DuckDB connection was lost: upgrading it to read-write failed and the \
store could not be reopened read-only afterwards. Another process is \
holding the file; retry once it has released it."
)
}
}
pub struct DuckDbProvider {
path: PathBuf,
conn: Arc<Mutex<Connection>>,
conn: Arc<Mutex<ConnHandle>>,
/// Embedding dimension; fixed at open time because the `FLOAT[N]` column depends on it.
dim: usize,
/// True once an FTS index has been built on `documents`. Until then
@@ -40,28 +79,140 @@ pub struct DuckDbProvider {
impl DuckDbProvider {
/// Open (or create) the DuckDB file. `dim` is the embedding vector dimension,
/// supplied by the caller who knows the model.
///
/// Opens READ-ONLY whenever the file already carries a complete schema, so that any
/// number of Coyote processes can query the same RAG at the same time. DuckDB allows
/// many concurrent readers XOR exactly one writer, so the exclusive read-write handle
/// is taken only when there is actually something to write: when the store has to be
/// created or initialized here, or lazily through `ensure_writable` on the rebuild
/// path.
pub fn open(db_path: &Path, dim: usize) -> Result<Self> {
let conn = Connection::open(db_path).with_context(|| {
let (conn, writable) = Self::open_for_workload(db_path, dim)?;
// A reopened file may already carry a live FTS index from a previous session,
// in which case keyword search works immediately.
let fts_exists = Self::probe_fts_index(&conn);
Ok(Self {
path: db_path.to_path_buf(),
conn: Arc::new(Mutex::new(ConnHandle {
conn: Some(conn),
writable,
})),
dim,
fts_ready: AtomicBool::new(fts_exists),
})
}
/// Pick the weakest access mode that can serve this store, returning the connection
/// and whether it came back writable.
fn open_for_workload(db_path: &Path, dim: usize) -> Result<(Connection, bool)> {
if db_path.exists()
&& let Ok(conn) = Self::open_read_only(db_path)
&& Self::store_is_initialized(&conn)
{
return Ok((conn, false));
}
// Three cases land here: the file does not exist yet, it could not be opened
// read-only (another process holds it read-write), or it carries no usable
// schema. All of them need a read-write handle, and the read-write attempt is
// also what produces the actionable lock error for the middle case.
let conn = Self::open_read_write(db_path, dim)?;
Ok((conn, true))
}
/// Is this file already a fully initialized Coyote store?
///
/// This gate decides whether a read-only open is viable, so it must be exact: every
/// statement in `init_schema` is rejected outright on a read-only handle, INCLUDING
/// `CREATE TABLE IF NOT EXISTS` against a table that already exists, which DuckDB
/// refuses rather than treating as a no-op. Anything missing therefore forces a
/// read-write open. The HNSW index is part of the check because a store whose tables
/// survived but whose index did not would otherwise be opened read-only and silently
/// serve every `vector_search` from a full scan.
fn store_is_initialized(conn: &Connection) -> bool {
let tables: i64 = conn
.query_row(
"SELECT count(*) FROM duckdb_tables() \
WHERE table_name IN ('vectors', 'documents')",
[],
|r| r.get(0),
)
.unwrap_or(0);
if tables < 2 {
return false;
}
conn.query_row(
"SELECT count(*) FROM duckdb_indexes() WHERE index_name = 'hnsw_idx'",
[],
|r| r.get::<_, i64>(0),
)
.map(|n| n > 0)
.unwrap_or(false)
}
/// Open the store read-only. Many processes may hold such a handle at once.
fn open_read_only(db_path: &Path) -> Result<Connection> {
let config = Config::default()
.access_mode(AccessMode::ReadOnly)
.context("Failed to build a read-only DuckDB configuration")?;
let conn = Connection::open_with_flags(db_path, config).with_context(|| {
format!(
"Failed to open the DuckDB store at '{}'. If another Coyote process (or \
another window) has this RAG open, close it and retry — a duckdb RAG can \
only be open in ONE process at a time. Unlike the yaml driver, its data \
lives in a single file with an exclusive lock.",
"Failed to open the DuckDB store at '{}' read-only",
db_path.display()
)
})?;
// Statement order is load-bearing. `hnsw_enable_experimental_persistence` is
// registered BY the vss extension, so setting it before `LOAD vss` fails with
// "Setting with name ... is not in the catalog, but it exists in the vss
// extension". Omitting it entirely makes CREATE INDEX ... USING HNSW on a
// file-backed database fail with "HNSW index persistence is not yet supported
// by default". ensure vss (installing it if missing) -> ensure fts -> SET ->
// CREATE INDEX.
Self::ensure_extension(&conn, "vss")?;
Self::ensure_extension(&conn, "fts")?;
Self::establish_session(&conn)?;
Ok(conn)
}
/// Open the store read-write and make sure its schema exists. Exactly one process
/// may hold such a handle, and no reader from another process may hold it meanwhile.
fn open_read_write(db_path: &Path, dim: usize) -> Result<Connection> {
let config = Config::default()
.access_mode(AccessMode::ReadWrite)
.context("Failed to build a read-write DuckDB configuration")?;
let conn = Connection::open_with_flags(db_path, config).with_context(|| {
format!(
"Failed to open the DuckDB store at '{}' for writing. Another Coyote \
process (or another window) has this RAG open: a duckdb RAG supports MANY \
concurrent READERS, but only ONE writer at a time, and a writer excludes \
readers in other processes. Close that process, or wait for its sync to \
finish, and retry.",
db_path.display()
)
})?;
Self::establish_session(&conn)?;
Self::init_schema(&conn, dim)?;
Ok(conn)
}
/// Install the per-connection session state that every connection needs, whatever
/// its access mode.
///
/// Extension `LOAD`s and `SET` are per-CONNECTION, not per-database: a connection
/// opened later — an upgrade, in particular — starts with none of this and must run
/// it again. None of these statements write to the database, so they all succeed on
/// a read-only handle.
///
/// Statement order is load-bearing. `hnsw_enable_experimental_persistence` is
/// registered BY the vss extension, so setting it before `LOAD vss` fails with
/// "Setting with name ... is not in the catalog, but it exists in the vss
/// extension". ensure vss (installing it if missing) -> ensure fts -> SET.
fn establish_session(conn: &Connection) -> Result<()> {
Self::ensure_extension(conn, "vss")?;
Self::ensure_extension(conn, "fts")?;
conn.execute_batch("SET hnsw_enable_experimental_persistence = true;")
.context("Failed to enable DuckDB HNSW index persistence")
}
/// Create the tables and the vector index. Every statement here WRITES, so this only
/// ever runs on a read-write connection.
///
/// Must be preceded by `establish_session`: without the `SET` it performs, a
/// CREATE INDEX ... USING HNSW on a file-backed database fails with "HNSW index
/// persistence is not yet supported by default".
fn init_schema(conn: &Connection, dim: usize) -> Result<()> {
conn.execute_batch(&format!(
"SET hnsw_enable_experimental_persistence = true;
CREATE TABLE IF NOT EXISTS vectors (
"CREATE TABLE IF NOT EXISTS vectors (
doc_id UBIGINT PRIMARY KEY,
embedding FLOAT[{dim}]
);
@@ -73,16 +224,44 @@ impl DuckDbProvider {
page_content TEXT NOT NULL
);"
))
.context("Failed to initialize DuckDB schema")?;
// A reopened file may already carry a live FTS index from a previous session,
// in which case keyword search works immediately.
let fts_exists = Self::probe_fts_index(&conn);
Ok(Self {
path: db_path.to_path_buf(),
conn: Arc::new(Mutex::new(conn)),
dim,
fts_ready: AtomicBool::new(fts_exists),
})
.context("Failed to initialize DuckDB schema")
}
/// Guarantee the shared connection is read-write, upgrading it in place if it is not.
/// EVERY write path must call this before touching the store.
///
/// The upgrade replaces the `Connection` INSIDE the shared `Arc<Mutex<..>>`, so
/// `duplicate()` clones, which share that `Arc`, see it too. The read-only connection
/// is dropped before the read-write open because DuckDB tracks the file lock per
/// database instance and the old handle still holds one.
///
/// On failure the store is reopened read-only so that queries keep working, and the
/// error is propagated so the caller aborts instead of writing. A failed upgrade must
/// leave the provider degraded, never bricked, and never silently read-only-with-a-
/// caller-that-thinks-it-wrote.
fn ensure_writable(&self) -> Result<()> {
let mut handle = self.lock_conn()?;
if handle.writable {
return Ok(());
}
drop(handle.conn.take());
match Self::open_read_write(&self.path, self.dim) {
Ok(conn) => {
handle.conn = Some(conn);
handle.writable = true;
Ok(())
}
Err(e) => {
handle.conn = Self::open_read_only(&self.path).ok();
Err(e.context(format!(
"Cannot write to the DuckDB RAG at '{}': it is open read-only and could \
not be upgraded to read-write, because another Coyote process has this \
RAG open. NOTHING WAS WRITTEN. Close the other process, or wait for it \
to finish, and retry.",
self.path.display()
)))
}
}
}
/// Make a DuckDB extension available on `conn`, installing it if this machine does
@@ -136,8 +315,10 @@ impl DuckDbProvider {
/// `rebuild_indexes` does `CREATE OR REPLACE TABLE` and writes exactly what it is
/// given, so a thinned map is committed as the new truth on the next sync.
pub(crate) fn read_all_vectors(&self) -> Result<IndexMap<DocumentId, Vec<f32>>> {
let conn = self.lock_conn()?;
let mut stmt = conn.prepare("SELECT doc_id, embedding FROM vectors")?;
let handle = self.lock_conn()?;
let mut stmt = handle
.conn()?
.prepare("SELECT doc_id, embedding FROM vectors")?;
let raw: Vec<(u64, Vec<f32>)> = stmt
.query_map([], |row| {
let id: u64 = row.get(0)?;
@@ -226,7 +407,7 @@ impl DuckDbProvider {
/// Never `.lock().unwrap()` here: a panic anywhere inside a locked scope poisons the
/// mutex permanently, and an unwrap would then turn every subsequent RAG query into
/// a panic for the remaining life of the process.
fn lock_conn(&self) -> Result<MutexGuard<'_, Connection>> {
fn lock_conn(&self) -> Result<MutexGuard<'_, ConnHandle>> {
self.conn
.lock()
.map_err(|e| anyhow!("DuckDB connection mutex was poisoned: {e}"))
@@ -253,7 +434,7 @@ impl RagProvider for DuckDbProvider {
.collect::<Vec<_>>()
.join(", ");
let dim = self.dim;
let conn = self.lock_conn()?;
let handle = self.lock_conn()?;
// array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[].
// ORDER BY distance ASC is required for the planner to use hnsw_idx; the
// similarity form (DESC) does NOT trigger the ANN index. Distance is converted
@@ -263,7 +444,7 @@ impl RagProvider for DuckDbProvider {
array_cosine_distance(embedding, [{vals}]::FLOAT[{dim}]) AS distance \
FROM vectors ORDER BY distance ASC LIMIT {top_k}"
);
let mut stmt = conn.prepare(&sql)?;
let mut stmt = handle.conn()?.prepare(&sql)?;
let results = stmt
.query_map([], |row| {
let id: u64 = row.get(0)?;
@@ -300,12 +481,12 @@ impl RagProvider for DuckDbProvider {
if ids.is_empty() {
return Ok(vec![]);
}
let conn = self.lock_conn()?;
let handle = self.lock_conn()?;
let placeholders = ids.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let sql =
format!("SELECT doc_id, page_content FROM documents WHERE doc_id IN ({placeholders})");
let params: Vec<Value> = ids.iter().map(|id| Value::UBigInt(id.0 as u64)).collect();
let mut stmt = conn.prepare(&sql)?;
let mut stmt = handle.conn()?.prepare(&sql)?;
let mut rows: Vec<(DocumentId, String)> = stmt
.query_map(duckdb::params_from_iter(params.iter()), |row| {
let id: u64 = row.get(0)?;
@@ -345,8 +526,10 @@ impl RagProvider for DuckDbProvider {
// Scoped: the guard MUST be dropped before `lock_conn()` is taken again
// below. `Mutex` is not reentrant; holding both self-deadlocks at
// runtime, with no compile error.
let conn = self.lock_conn()?;
conn.query_row("SELECT count(*) FROM vectors", [], |r| r.get(0))
let handle = self.lock_conn()?;
handle
.conn()?
.query_row("SELECT count(*) FROM vectors", [], |r| r.get(0))
.context("Failed to count existing vectors before rebuild")?
};
if existing > 0 {
@@ -372,7 +555,16 @@ impl RagProvider for DuckDbProvider {
}
}
let dim = self.dim;
let mut conn = self.lock_conn()?;
// THE write path. Everything above this line only reads, so the upgrade happens
// here, after both guards have had their say: a rebuild that is going to be
// refused must not first take the exclusive lock away from other processes.
//
// This is also the point that makes a silently-dropped write impossible. If the
// upgrade fails, `?` aborts the rebuild before a single statement is issued and
// the caller gets the error. Nothing below can run on a read-only connection.
self.ensure_writable()?;
let mut handle = self.lock_conn()?;
let conn = handle.conn_mut()?;
let tx = conn
.transaction()
.context("Failed to begin DuckDB transaction")?;
@@ -476,9 +668,9 @@ impl RagProvider for DuckDbProvider {
}
async fn keyword_search(&self, query: &str, top_k: usize) -> Result<Vec<(DocumentId, f32)>> {
let conn = self.lock_conn()?;
let handle = self.lock_conn()?;
// match_bm25 returns NULL for non-matching rows; WHERE filters them out.
let mut stmt = conn.prepare(
let mut stmt = handle.conn()?.prepare(
"SELECT doc_id, fts_main_documents.match_bm25(doc_id, ?) AS score
FROM documents
WHERE score IS NOT NULL
@@ -533,6 +725,10 @@ impl RagProvider for DuckDbProvider {
// disk and a rebuild through one handle is immediately visible to the other.
// That is unavoidable for any on-disk store and is handled by the discipline
// documented on `Rag`'s Clone impl, the pre-clone instance must be discarded.
//
// Sharing the Arc also shares the ACCESS MODE, which lives inside the ConnHandle
// rather than beside it: when one handle upgrades itself to read-write, every
// clone is upgraded with it and none is left holding a stale "read-only" belief.
Box::new(DuckDbProvider {
path: self.path.clone(),
conn: Arc::clone(&self.conn),
@@ -547,6 +743,7 @@ mod tests {
use super::*;
use crate::rag::provider::RagProvider;
use crate::rag::{RagDocument, RagFile};
use std::process::Command;
use std::time::{SystemTime, UNIX_EPOCH};
use std::{env, fs};
@@ -625,7 +822,8 @@ mod tests {
async fn open_creates_schema() {
let db = TempDb::new("schema");
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
let v: i64 = conn
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
@@ -643,7 +841,8 @@ mod tests {
let db = TempDb::new("vsearch");
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
{
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
// The ::FLOAT[3] cast is REQUIRED: a bare [0.1, 0.2, 0.3] literal infers
// DOUBLE[], which does not match the FLOAT[N] ARRAY column type.
conn.execute(
@@ -668,7 +867,8 @@ mod tests {
let db = TempDb::new("fetch");
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
{
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
conn.execute(
"INSERT INTO documents (doc_id, page_content) VALUES (42, 'hello world')",
[],
@@ -712,7 +912,8 @@ mod tests {
.await
.expect("second rebuild must not violate the primary key constraint");
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
let count: i64 = conn
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
.unwrap();
@@ -744,7 +945,8 @@ mod tests {
reloaded.vectors.insert(DocumentId(2), vec![0.7, 0.8, 0.9]);
provider.rebuild_indexes(&reloaded, false).await.unwrap();
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
let count: i64 = conn
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
.unwrap();
@@ -768,7 +970,8 @@ mod tests {
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
provider.rebuild_indexes(&data, true).await.unwrap();
{
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
let docs: i64 = conn
.query_row("SELECT COUNT(*) FROM documents", [], |r| r.get(0))
.unwrap();
@@ -849,7 +1052,8 @@ mod tests {
"got: {err}"
);
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
let count: i64 = conn
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
.unwrap();
@@ -877,7 +1081,8 @@ mod tests {
let dup = provider.duplicate(&minimal_rag_data());
{
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
conn.execute(
"INSERT INTO documents (doc_id, page_content) VALUES (7, 'shared row')",
[],
@@ -900,7 +1105,8 @@ mod tests {
let db = TempDb::new("nonfinite");
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
{
let conn = provider.conn.lock().unwrap();
let handle = provider.conn.lock().unwrap();
let conn = handle.conn().unwrap();
conn.execute(
"INSERT INTO vectors (doc_id, embedding) VALUES (0, [0.1, 0.2, 0.3]::FLOAT[3])",
[],
@@ -923,6 +1129,310 @@ mod tests {
);
}
/// Was the shared connection opened read-write? Reads the flag that lives inside the
/// shared handle, which is the same one `ensure_writable` flips.
fn is_writable(provider: &DuckDbProvider) -> bool {
provider.conn.lock().unwrap().writable
}
/// Build a fully initialized store, then let the read-write handle go so the file is
/// unlocked for the next opener.
async fn seed_store(path: &Path) {
let mut provider = DuckDbProvider::open(path, 3).unwrap();
assert!(is_writable(&provider), "a fresh file must open read-write");
let mut data = populated_rag_data();
data.vectors
.insert(DocumentId::new(0, 0), vec![0.1, 0.2, 0.3]);
provider.rebuild_indexes(&data, true).await.unwrap();
}
#[tokio::test]
async fn a_fresh_file_is_opened_read_write() {
let db = TempDb::new("freshrw");
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
assert!(
is_writable(&provider),
"the schema has to be created, which writes, so a missing file must open \
read-write"
);
}
#[tokio::test]
async fn an_initialized_store_is_reopened_read_only() {
let db = TempDb::new("reopenro");
seed_store(&db.path).await;
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
assert!(
!is_writable(&provider),
"a store that needs no schema work must open read-only, so that other Coyote \
processes can query it at the same time"
);
}
#[tokio::test]
async fn a_file_without_the_coyote_schema_is_opened_read_write() {
let db = TempDb::new("noschema");
{
// A valid DuckDB file that is not one of ours. Opening it read-only would
// strand it forever: the init batch is refused on a read-only handle.
let conn = Connection::open(&db.path).unwrap();
conn.execute_batch("CREATE TABLE unrelated (x INTEGER);")
.unwrap();
}
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
assert!(
is_writable(&provider),
"a schema-less file must open read-write"
);
let handle = provider.conn.lock().unwrap();
let tables: i64 = handle
.conn()
.unwrap()
.query_row(
"SELECT count(*) FROM duckdb_tables() \
WHERE table_name IN ('vectors', 'documents')",
[],
|r| r.get(0),
)
.unwrap();
assert_eq!(tables, 2, "the schema must have been created");
}
#[tokio::test]
async fn a_read_only_store_still_serves_vector_and_keyword_search() {
let db = TempDb::new("rosearch");
seed_store(&db.path).await;
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
assert!(!is_writable(&provider), "precondition: opened read-only");
let hits = provider
.vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
.await
.unwrap();
assert_eq!(
hits.len(),
1,
"the persisted HNSW index must be queryable read-only"
);
assert!(hits[0].1 > 0.99);
assert!(
provider.has_native_keyword_search(),
"the FTS index built by the previous session must still be detected on a \
read-only handle"
);
let kw = provider.keyword_search("alpha", 5).await.unwrap();
assert_eq!(kw.len(), 1, "keyword search must work read-only");
let docs = provider
.fetch_content(&[DocumentId::new(0, 0)])
.await
.unwrap();
assert_eq!(docs[0].1, "alpha keyword");
}
#[tokio::test]
async fn a_read_only_handle_refuses_a_direct_write() {
let db = TempDb::new("rorefuse");
seed_store(&db.path).await;
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
assert!(!is_writable(&provider), "precondition: opened read-only");
let handle = provider.conn.lock().unwrap();
let err = handle
.conn()
.unwrap()
.execute(
"INSERT INTO documents (doc_id, page_content) VALUES (99, 'nope')",
[],
)
.unwrap_err();
// The backstop, not the primary defence: `rebuild_indexes` upgrades first and
// never reaches a write on a read-only handle. It matters anyway because DuckDB
// lets `transaction()` open and `commit()` return Ok on a read-only connection,
// so a write that slipped through would look like it had succeeded.
assert!(
err.to_string().contains("read-only mode"),
"a read-only handle must refuse writes loudly; got: {err}"
);
}
#[tokio::test]
async fn rebuild_indexes_upgrades_a_read_only_connection() {
let db = TempDb::new("upgrade");
seed_store(&db.path).await;
let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
assert!(!is_writable(&provider), "precondition: opened read-only");
let mut data = populated_rag_data();
data.vectors = provider.read_all_vectors().unwrap();
data.vectors
.insert(DocumentId::new(1, 0), vec![0.4, 0.5, 0.6]);
provider.rebuild_indexes(&data, false).await.unwrap();
assert!(
is_writable(&provider),
"the write path must have upgraded the connection in place"
);
// The upgraded connection is a NEW connection, so the per-connection session
// state has to have been re-established on it. Without the re-run `SET`, the
// CREATE INDEX ... USING HNSW inside rebuild_indexes would already have failed.
let persisted: String = {
let handle = provider.conn.lock().unwrap();
handle
.conn()
.unwrap()
.query_row(
"SELECT CAST(current_setting('hnsw_enable_experimental_persistence') \
AS VARCHAR)",
[],
|r| r.get(0),
)
.unwrap()
};
assert_eq!(
persisted, "true",
"the upgraded connection must re-run the SET; session state does not carry \
over from the dropped read-only connection"
);
drop(provider);
let reopened = DuckDbProvider::open(&db.path, 3).unwrap();
let all = reopened.read_all_vectors().unwrap();
assert_eq!(all.len(), 2, "the upgraded write must have reached disk");
}
#[tokio::test]
async fn an_upgrade_is_visible_through_duplicate_clones() {
let db = TempDb::new("upgradedup");
seed_store(&db.path).await;
let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
assert!(!is_writable(&provider), "precondition: opened read-only");
let dup = provider.duplicate(&minimal_rag_data());
let mut data = populated_rag_data();
data.vectors = provider.read_all_vectors().unwrap();
provider.rebuild_indexes(&data, false).await.unwrap();
// `duplicate()` shares the Arc, and the access mode lives inside it, so the clone
// must observe the upgrade rather than keep believing it is read-only.
let via_dup = dup.fetch_content(&[DocumentId::new(0, 0)]).await.unwrap();
assert_eq!(
via_dup.len(),
1,
"the clone must still read after an upgrade"
);
assert!(
is_writable(&provider),
"the shared handle must report writable to every clone"
);
}
/// Two REAL OS processes reading one store at the same time.
///
/// Ignored by default because it re-executes the test binary as a child process,
/// which is heavier and more environment-dependent than the rest of the suite. Run it
/// with:
/// cargo test --all -- --ignored duckdb_store_is_shared_across_processes
///
/// It cannot be written as an ordinary in-process test: DuckDB keeps ONE database
/// instance per process, so a second open in the same process bypasses the file lock
/// entirely (a read-write open succeeds even while this process holds a read-only
/// one). Only separate processes exercise the lock this feature exists to avoid.
#[tokio::test]
#[ignore = "spawns a second OS process; run explicitly with --ignored"]
async fn duckdb_store_is_shared_across_processes() {
const CHILD_DB: &str = "COYOTE_DUCKDB_MULTIPROC_DB";
const CHILD_EXPECT: &str = "COYOTE_DUCKDB_MULTIPROC_EXPECT";
const TEST_NAME: &str =
"rag::providers::duckdb::tests::duckdb_store_is_shared_across_processes";
if let Ok(path) = env::var(CHILD_DB) {
let expect = env::var(CHILD_EXPECT).unwrap_or_default();
let opened = DuckDbProvider::open(Path::new(&path), 3);
match expect.as_str() {
"readable" => {
let provider = opened.expect(
"a second process must be able to open a store that another \
process holds READ-ONLY",
);
assert!(!is_writable(&provider), "the child must land read-only");
let hits = provider
.vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
.await
.unwrap();
assert_eq!(hits.len(), 1, "the child must read the seeded vector");
}
"blocked" => {
let err = opened.err().expect(
"a second process must NOT be able to open a store that another \
process holds READ-WRITE",
);
let msg = format!("{err:#}");
assert!(
msg.contains("concurrent READERS") && msg.contains("only ONE writer"),
"the lock error must explain the reader/writer rule; got: {msg}"
);
}
other => panic!("unknown child expectation {other:?}"),
}
return;
}
let db = TempDb::new("multiproc");
seed_store(&db.path).await;
let run_child = |expect: &str| {
Command::new(env::current_exe().unwrap())
.args(["--exact", "--ignored", "--nocapture", TEST_NAME])
.env(CHILD_DB, &db.path)
.env(CHILD_EXPECT, expect)
.output()
.expect("failed to spawn the child test process")
};
// Phase 1: this process holds a READ-ONLY handle. The child must get one too.
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
assert!(!is_writable(&provider), "precondition: parent is read-only");
let out = run_child("readable");
assert!(
out.status.success(),
"child could not share the read-only store:\n{}\n{}",
String::from_utf8_lossy(&out.stdout),
String::from_utf8_lossy(&out.stderr)
);
let hits = provider
.vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
.await
.unwrap();
assert_eq!(
hits.len(),
1,
"the parent must still read after the child ran"
);
// Phase 2: upgrade this process to READ-WRITE. The child must now be refused,
// with the message that explains why.
provider.ensure_writable().unwrap();
let out = run_child("blocked");
assert!(
out.status.success(),
"a writer must exclude other processes, with an actionable error:\n{}\n{}",
String::from_utf8_lossy(&out.stdout),
String::from_utf8_lossy(&out.stderr)
);
}
#[test]
fn duckdb_path_from_yaml_swaps_extension() {
let p = duckdb_path_from_yaml(Path::new("/tmp/rags/docs.yaml"));
+372 -36
View File
@@ -3,10 +3,125 @@ use crate::rag::{DocumentId, RagData};
use anyhow::{Context, Result, bail};
use async_trait::async_trait;
use parking_lot::RwLock;
use reqwest::header::{HeaderMap, HeaderValue};
use reqwest::{Client, Response, StatusCode};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use url::{Host, Url};
/// 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
/// LangChain writes by default.
///
/// `DocumentId` packs `(file_index, document_index)` into one `usize` with the
/// file index in the high half, so this bit is only reachable at a file index of
/// 2^31. Nothing local gets near that, and an attached RAG builds no local index
/// at all — `data.files` and `data.vectors` stay empty and every
/// `DocumentId::split` caller early-returns on `data.attached`. Along the
/// attached path the id is an opaque key carried through RRF, which is what
/// makes a synthetic one safe here and nowhere else.
const SYNTHETIC_ID_TAG: usize = 1 << (usize::BITS - 1);
/// Two-way map between a raw Qdrant point id and the `DocumentId` the retrieval
/// pipeline sees.
///
/// Only ids that cannot survive the round trip are interned. A plain `u64` that
/// fits below the tag keeps mapping to itself, so integer-keyed collections
/// behave exactly as they did before this map existed.
#[derive(Default)]
struct PointIdInterner {
handles: HashMap<String, DocumentId>,
raw: HashMap<DocumentId, Value>,
next: usize,
}
impl PointIdInterner {
/// The `DocumentId` for a raw point id, minting a handle if one is needed.
///
/// `None` only for a missing id, which is a malformed response.
fn document_id(&mut self, raw: &Value) -> Option<DocumentId> {
if raw.is_null() {
return None;
}
// The pre-existing integer path, unchanged. `try_from` rather than `as`
// so a value too wide for the target's `usize` is interned instead of
// silently truncated into a different point.
if let Some(n) = raw.as_u64()
&& let Ok(n) = usize::try_from(n)
&& n & SYNTHETIC_ID_TAG == 0
{
return Some(DocumentId(n));
}
Some(self.intern(raw))
}
fn intern(&mut self, raw: &Value) -> DocumentId {
// Keyed on the JSON rendering, so the string "1" and the integer 1 are
// not conflated into one point.
let key = raw.to_string();
if let Some(handle) = self.handles.get(&key) {
return *handle;
}
let handle = DocumentId(SYNTHETIC_ID_TAG | self.next);
self.next += 1;
self.handles.insert(key, handle);
self.raw.insert(handle, raw.clone());
handle
}
/// The original id for a handle, or `None` when the id was never interned —
/// i.e. it is a plain integer that is already its own id.
fn raw_id(&self, handle: DocumentId) -> Option<&Value> {
self.raw.get(&handle)
}
/// Builds the `ids` array for an outbound `/points` fetch. Every entry is the
/// id Qdrant issued, integer or string; a synthetic handle must never leave
/// this process.
fn outbound_ids(&self, ids: &[DocumentId]) -> Vec<Value> {
ids.iter()
.map(|id| match self.raw_id(*id) {
Some(raw) => raw.clone(),
None => Value::from(id.0 as u64),
})
.collect()
}
}
fn parse_search_hits(
interner: &mut PointIdInterner,
body: &Value,
min_score: f32,
) -> Result<Vec<(DocumentId, f32)>> {
let hits = body["result"]
.as_array()
.context("Unexpected /points/search response shape")?;
Ok(hits
.iter()
.filter_map(|pt| {
let score = pt["score"].as_f64()? as f32;
Some((interner.document_id(&pt["id"])?, score))
})
.filter(|(_, score)| min_score <= 0.0 || *score > min_score)
.collect())
}
fn parse_points(interner: &mut PointIdInterner, body: &Value) -> Result<Vec<(DocumentId, String)>> {
let points = body["result"]
.as_array()
.context("Unexpected /points response shape")?;
Ok(points
.iter()
.filter_map(|pt| {
let text = pt["payload"]["page_content"].as_str()?.to_string();
Some((interner.document_id(&pt["id"])?, text))
})
.collect())
}
/// Render Qdrant's error envelope into a human-readable message.
///
@@ -65,10 +180,26 @@ pub struct QdrantProvider {
client: Client,
base_url: String,
collection: String,
point_ids: Arc<RwLock<PointIdInterner>>,
}
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();
if let Some(key) = api_key {
let mut value =
@@ -76,10 +207,11 @@ impl QdrantProvider {
value.set_sensitive(true);
headers.insert("api-key", value);
}
Client::builder()
.default_headers(headers)
.build()
.context("Failed to build reqwest client")
let mut builder = Client::builder().default_headers(headers);
if Self::skips_proxy(base_url) {
builder = builder.no_proxy();
}
builder.build().context("Failed to build reqwest client")
}
pub(crate) fn normalize_base_url(host: &str) -> String {
@@ -104,7 +236,7 @@ impl QdrantProvider {
api_key: Option<&str>,
) -> Result<Value> {
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
.get(format!("{base_url}/collections/{collection}"))
.send()
@@ -122,7 +254,7 @@ impl QdrantProvider {
pub async fn new(host: &str, collection: &str, api_key: Option<&str>) -> Result<Self> {
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
.get(format!("{base_url}/collections/{collection}"))
.send()
@@ -139,12 +271,13 @@ impl QdrantProvider {
client,
base_url,
collection: collection.to_string(),
point_ids: Arc::default(),
})
}
pub async fn list_collections(host: &str, api_key: Option<&str>) -> Result<Vec<String>> {
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
.get(format!("{base_url}/collections"))
.send()
@@ -194,7 +327,7 @@ impl QdrantProvider {
api_key: Option<&str>,
) -> Result<Option<String>> {
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 body = serde_json::json!({ "limit": 1, "with_payload": false });
@@ -237,7 +370,8 @@ impl RagProvider for QdrantProvider {
// `score_threshold` is deliberately NOT sent. It is metric-aware: on Cosine
// collections 0.0 means "no floor" as expected, but Euclid collections score
// 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!({
"vector": embedding,
"limit": top_k,
@@ -252,22 +386,11 @@ impl RagProvider for QdrantProvider {
);
}
let data: Value = resp.json().await?;
let results = data["result"]
.as_array()
.context("Unexpected /points/search response shape")?
.iter()
.filter_map(|pt| {
// String (UUID) IDs yield None here and are dropped. The attach
// wizard rejects such collections up front so this cannot silently
// become "zero results, no error".
let id = pt["id"].as_u64()? as usize;
let score = pt["score"].as_f64()? as f32;
Some((DocumentId(id), score))
})
.filter(|(_, score)| *score > min_score)
.collect();
// The interner is what lets a UUID-keyed collection work: a string id gets
// a synthetic handle here and the original is replayed by `fetch_content`.
let mut interner = self.point_ids.write();
Ok(results)
parse_search_hits(&mut interner, &data, min_score)
}
async fn fetch_content(&self, ids: &[DocumentId]) -> Result<Vec<(DocumentId, String)>> {
@@ -275,7 +398,8 @@ impl RagProvider for QdrantProvider {
return Ok(vec![]);
}
let url = format!("{}/collections/{}/points", self.base_url, self.collection);
let id_list: Vec<u64> = ids.iter().map(|d| d.0 as u64).collect();
// Qdrant is asked for the ids it issued, never for a synthetic handle.
let id_list = self.point_ids.read().outbound_ids(ids);
let body = serde_json::json!({
"ids": id_list,
"with_payload": true,
@@ -291,16 +415,10 @@ impl RagProvider for QdrantProvider {
);
}
let data: Value = resp.json().await?;
let mut rows: Vec<(DocumentId, String)> = data["result"]
.as_array()
.context("Unexpected /points response shape")?
.iter()
.filter_map(|pt| {
let id = pt["id"].as_u64()? as usize;
let text = pt["payload"]["page_content"].as_str()?.to_string();
Some((DocumentId(id), text))
})
.collect();
let mut rows = {
let mut interner = self.point_ids.write();
parse_points(&mut interner, &data)?
};
// `/points` does not guarantee response order matches request order, and the
// caller's RRF ranking is carried by that order. Restore it.
let position: HashMap<DocumentId, usize> =
@@ -328,10 +446,17 @@ impl RagProvider for QdrantProvider {
// Cloning the client shares the connection pool and the injected api-key
// header. Sharing is correct: both handles address the same remote
// collection, and neither of them writes to it.
//
// The point-id map is shared for the same reason, and because it MUST be:
// `Rag::clone()` hands the clone `DocumentId`s that the original minted,
// so a fresh map would resolve them to nothing and `fetch_content` would
// ask Qdrant for a synthetic handle — zero results, no error. Resetting it
// would also re-mint handles for ids the original still holds.
Box::new(Self {
client: self.client.clone(),
base_url: self.base_url.clone(),
collection: self.collection.clone(),
point_ids: Arc::clone(&self.point_ids),
})
}
}
@@ -428,6 +553,7 @@ mod tests {
client: Client::new(),
base_url: "http://localhost:6333".to_string(),
collection: "c".to_string(),
point_ids: Arc::default(),
};
let attached = RagData {
@@ -462,11 +588,221 @@ mod tests {
client: Client::new(),
base_url: "http://127.0.0.1:1".to_string(),
collection: "c".to_string(),
point_ids: Arc::default(),
};
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` →
/// `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
/// `filter_map`, i.e. zero results and no error.
#[test]
fn uuid_point_ids_round_trip_and_are_requested_verbatim() {
let mut interner = PointIdInterner::default();
let first_uuid = "3f1b0c2e-1111-4000-8000-000000000001";
let second_uuid = "3f1b0c2e-2222-4000-8000-000000000002";
let search = serde_json::json!({
"result": [
{"id": first_uuid, "score": 0.91},
{"id": second_uuid, "score": 0.42},
]
});
let hits = parse_search_hits(&mut interner, &search, 0.0).unwrap();
assert_eq!(hits.len(), 2, "string ids must not be silently dropped");
let ids: Vec<DocumentId> = hits.iter().map(|(id, _)| *id).collect();
assert_eq!(
interner.outbound_ids(&ids),
vec![Value::from(first_uuid), Value::from(second_uuid)],
"the fetch must send the ids Qdrant issued, not the handles"
);
// Qdrant may answer /points in any order; the handles still map back and
// the caller's RRF ranking is recoverable.
let points = serde_json::json!({
"result": [
{"id": second_uuid, "payload": {"page_content": "second"}},
{"id": first_uuid, "payload": {"page_content": "first"}},
]
});
let mut rows = parse_points(&mut interner, &points).unwrap();
let position: HashMap<DocumentId, usize> =
ids.iter().enumerate().map(|(i, id)| (*id, i)).collect();
rows.sort_by_key(|(id, _)| position.get(id).copied().unwrap_or(usize::MAX));
assert_eq!(
rows,
vec![
(ids[0], "first".to_string()),
(ids[1], "second".to_string())
]
);
}
/// Integer-keyed collections must be untouched by the interner: the id maps to
/// itself on the way in and goes back out as the same integer.
#[test]
fn integer_point_ids_are_passed_through_untouched() {
let mut interner = PointIdInterner::default();
let search = serde_json::json!({
"result": [{"id": 7, "score": 0.9}, {"id": 0, "score": 0.5}]
});
let hits = parse_search_hits(&mut interner, &search, 0.0).unwrap();
assert_eq!(
hits,
vec![(DocumentId(7), 0.9_f32), (DocumentId(0), 0.5_f32)]
);
let ids: Vec<DocumentId> = hits.iter().map(|(id, _)| *id).collect();
assert_eq!(
interner.outbound_ids(&ids),
vec![Value::from(7_u64), Value::from(0_u64)],
"integer ids must not be regressed into synthetic handles"
);
assert!(
interner.raw_id(DocumentId(7)).is_none(),
"a plain integer id is its own id and needs no map entry"
);
}
/// Synthetic handles are stable per point id and live in a range no packed
/// `DocumentId` can reach.
#[test]
fn synthetic_handles_are_stable_and_never_collide_with_packed_ids() {
let mut interner = PointIdInterner::default();
let uuid = Value::from("9d2f0a11-3333-4000-8000-00000000000a");
let handle = interner.document_id(&uuid).unwrap();
assert_eq!(
interner.document_id(&uuid).unwrap(),
handle,
"the same point id must keep the same handle across queries"
);
assert_ne!(
interner.document_id(&Value::from("other")).unwrap(),
handle,
"distinct point ids must not share a handle"
);
assert_ne!(handle.0 & SYNTHETIC_ID_TAG, 0, "a handle carries the tag");
// A packed (file_index, document_index) never sets the tag bit: it is the
// top bit of the file index, which would take 2^31 indexed files.
for (file_index, document_index) in [(0, 0), (1, 0), (0, 4242), (1_000_000, 999)] {
assert_eq!(
DocumentId::new(file_index, document_index).0 & SYNTHETIC_ID_TAG,
0,
"packed ({file_index}, {document_index}) must stay out of the handle range"
);
}
// The one integer id that WOULD land on the tag is interned instead of
// being handed back as itself, so it cannot alias a handle.
let collides = Value::from(SYNTHETIC_ID_TAG as u64);
let interned = interner.document_id(&collides).unwrap();
assert_eq!(interner.raw_id(interned), Some(&collides));
assert_eq!(
interner.outbound_ids(&[interned]),
vec![collides],
"the original integer must still be what Qdrant is asked for"
);
}
/// `duplicate()` shares the map rather than resetting it: `Rag::clone()` hands
/// the clone `DocumentId`s the original minted, and a fresh map would turn
/// those into requests for a synthetic handle — zero results, no error.
#[test]
fn duplicate_shares_the_point_id_map() {
let provider = QdrantProvider {
client: Client::new(),
base_url: "http://127.0.0.1:1".to_string(),
collection: "c".to_string(),
point_ids: Arc::default(),
};
let uuid = Value::from("c0ffee00-4444-4000-8000-000000000007");
let handle = provider.point_ids.write().document_id(&uuid).unwrap();
let dup = provider.duplicate(&RagData {
driver: "qdrant".to_string(),
attached: true,
..Default::default()
});
// Downcasting is not available through `dyn RagProvider`, so go via the
// shared Arc: the clone must observe the original's interning.
assert_eq!(Arc::strong_count(&provider.point_ids), 2);
assert_eq!(provider.point_ids.read().raw_id(handle), Some(&uuid));
drop(dup);
}
#[tokio::test]
#[ignore]
async fn qdrant_list_collections_requires_running_instance() {
+1 -1
View File
@@ -892,7 +892,7 @@ pub async fn run_repl_command(
".rag" => match split_first_arg(args) {
Some(("attach", rest)) => match rest {
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>"),
},
+482 -47
View File
@@ -1,9 +1,10 @@
use std::env;
use std::fs;
use std::fs::{read_dir, read_to_string};
use std::io;
use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use anyhow::{Context, Result, anyhow, bail};
use serde_yaml::Value;
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_SUFFIX: &str = ".sbx-mixin.yaml";
const KIT_SPEC_FILE_NAME: &str = "spec.yaml";
const MIXIN_FILES_DIR_NAME: &str = "files";
#[derive(Debug, Clone)]
pub struct DiscoveredMixin {
@@ -34,46 +36,168 @@ impl DiscoveredMixin {
pub fn wrap_mixin_as_kit(mixin_path: &Path) -> Result<PathBuf> {
let bytes = fs::read(mixin_path)
.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> {
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();
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 kit_dir = paths::sbx_mixin_kits_dir().join(&hash);
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)
&& existing == bytes
{
let spec_matches = fs::read(&spec_path).is_ok_and(|existing| existing == spec_bytes);
let files_ready = files.is_empty() || files_dst.is_dir();
if spec_matches && files_ready {
return Ok(kit_dir);
}
fs::create_dir_all(&kit_dir)
.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()))?;
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());
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>> {
let mut out = Vec::new();
push_if_exists(&mut out, paths::sbx_mixin_file())?;
push_if_exists(&mut out, paths::global_tools_sbx_mixin_file())?;
for path in collect_subdir_mixins(&paths::functions_dir()) {
for path in collect_mixins(&paths::functions_dir(), &[ScanMode::SubdirNamed]) {
out.push(read_mixin(path)?);
}
for path in collect_subdir_mixins(&paths::agents_data_dir()) {
for path in collect_mixins(
&paths::agents_data_dir(),
&[ScanMode::SubdirNamed, ScanMode::SubdirFlat],
) {
out.push(read_mixin(path)?);
}
for path in collect_flat_mixins(&paths::rags_dir()) {
for path in collect_mixins(&paths::rags_dir(), &[ScanMode::Flat]) {
out.push(read_mixin(path)?);
}
@@ -160,27 +284,54 @@ fn read_mixin(path: PathBuf) -> Result<DiscoveredMixin> {
})
}
fn collect_subdir_mixins(dir: &Path) -> Vec<PathBuf> {
/// One on-disk layout a mixin scan can look for. A scan takes a set of these,
/// and each mode contributes only the shape it names.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ScanMode {
/// `<dir>/*.sbx-mixin.yaml`
Flat,
/// `<dir>/*/sbx-mixin.yaml`
SubdirNamed,
/// `<dir>/*/*.sbx-mixin.yaml`
SubdirFlat,
}
/// Collects mixin paths under `dir` for every requested layout. Missing or
/// unreadable directories yield nothing rather than an error — these paths are
/// all optional on disk.
///
/// Order is deterministic: flat matches first (sorted by file name), then each
/// subdirectory in sorted order, contributing its named mixin before its
/// suffixed ones.
fn collect_mixins(dir: &Path, modes: &[ScanMode]) -> Vec<PathBuf> {
let mut result = Vec::new();
let Ok(rd) = read_dir(dir) else { return result };
let mut entries: Vec<_> = rd
.flatten()
.filter(|e| e.file_type().map(|t| t.is_dir()).unwrap_or(false))
.collect();
entries.sort_by_key(|e| e.file_name());
if modes.contains(&ScanMode::Flat) {
result.extend(suffixed_mixins_in(dir));
}
for entry in entries {
let candidate = entry.path().join(SBX_MIXIN_FILE_NAME);
if candidate.exists() {
result.push(candidate);
let named = modes.contains(&ScanMode::SubdirNamed);
let subdir_flat = modes.contains(&ScanMode::SubdirFlat);
if !named && !subdir_flat {
return result;
}
for subdir in subdirs_of(dir) {
if named {
let candidate = subdir.join(SBX_MIXIN_FILE_NAME);
if candidate.exists() {
result.push(candidate);
}
}
if subdir_flat {
result.extend(suffixed_mixins_in(&subdir));
}
}
result
}
fn collect_flat_mixins(dir: &Path) -> Vec<PathBuf> {
fn suffixed_mixins_in(dir: &Path) -> Vec<PathBuf> {
let mut result = Vec::new();
let Ok(rd) = read_dir(dir) else { return result };
@@ -195,10 +346,21 @@ fn collect_flat_mixins(dir: &Path) -> Vec<PathBuf> {
.collect();
entries.sort_by_key(|e| e.file_name());
for entry in entries {
result.push(entry.path());
}
result.extend(entries.into_iter().map(|e| e.path()));
result
}
fn subdirs_of(dir: &Path) -> Vec<PathBuf> {
let mut result = Vec::new();
let Ok(rd) = read_dir(dir) else { return result };
let mut entries: Vec<_> = rd
.flatten()
.filter(|e| e.file_type().map(|t| t.is_dir()).unwrap_or(false))
.collect();
entries.sort_by_key(|e| e.file_name());
result.extend(entries.into_iter().map(|e| e.path()));
result
}
@@ -218,6 +380,13 @@ mod tests {
root
}
fn file_names(paths: &[PathBuf]) -> Vec<&str> {
paths
.iter()
.map(|p| p.file_name().unwrap().to_str().unwrap())
.collect()
}
#[test]
fn summarize_counts_installs_and_domains() {
let root = unique_root("sbx-mixin-counts");
@@ -301,7 +470,7 @@ network:
}
#[test]
fn collect_subdir_mixins_sorts_and_skips_missing() {
fn subdir_named_scan_sorts_and_skips_missing() {
let root = unique_root("sbx-mixin-subdirs");
for name in ["zebra", "apple", "no-mixin", "mango"] {
let dir = root.join(name);
@@ -311,7 +480,7 @@ network:
}
}
let found = collect_subdir_mixins(&root);
let found = collect_mixins(&root, &[ScanMode::SubdirNamed]);
let names: Vec<String> = found
.iter()
.map(|p| {
@@ -329,9 +498,9 @@ network:
}
#[test]
fn collect_subdir_mixins_returns_empty_for_missing_dir() {
fn subdir_named_scan_returns_empty_for_missing_dir() {
let absent = env::temp_dir().join("coyote-definitely-not-here-xyz");
let found = collect_subdir_mixins(&absent);
let found = collect_mixins(&absent, &[ScanMode::SubdirNamed]);
assert!(found.is_empty());
}
@@ -508,10 +677,200 @@ network:
"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]
fn collect_flat_mixins_matches_rag_sidecars_by_suffix() {
fn flat_scan_matches_rag_sidecars_by_suffix() {
let root = unique_root("flat-mixins");
fs::write(root.join("company-docs.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
fs::write(root.join("alpha.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
@@ -519,39 +878,115 @@ network:
fs::write(root.join("notes.yaml"), "driver: yaml\n").unwrap();
fs::create_dir_all(root.join("decoy.sbx-mixin.yaml")).unwrap();
let found = collect_flat_mixins(&root);
let names: Vec<_> = found
.iter()
.map(|p| p.file_name().unwrap().to_str().unwrap())
.collect();
let found = collect_mixins(&root, &[ScanMode::Flat]);
assert_eq!(
names,
file_names(&found),
vec!["alpha.sbx-mixin.yaml", "company-docs.sbx-mixin.yaml"]
);
let _ = fs::remove_dir_all(&root);
}
/// Why `collect_flat_mixins` had to be written: the existing collector walks
/// SUBDIRECTORIES for a file named exactly `sbx-mixin.yaml`, so it cannot see
/// a flat sidecar. If this ever starts finding them, the new collector is
/// redundant — but until then, removing it silently drops every RAG mixin.
/// Every scan site in `discover()` picks its modes assuming each mode owns
/// exactly one layout and nothing else. `agents_data_dir()` requests two
/// modes at once, so an overlap would collect the same file twice and
/// `create_sandbox` would pass it as two `--kit` flags.
#[test]
fn collect_subdir_mixins_cannot_see_flat_rag_sidecars() {
let root = unique_root("flat-vs-subdir");
fs::write(root.join("company-docs.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
fn each_scan_mode_owns_exactly_one_layout() {
let root = unique_root("scan-mode-ownership");
let agent = root.join("researcher");
fs::create_dir_all(&agent).unwrap();
let flat = root.join("company-docs.sbx-mixin.yaml");
let subdir_named = agent.join("sbx-mixin.yaml");
let subdir_flat = agent.join("handbook.sbx-mixin.yaml");
for path in [&flat, &subdir_named, &subdir_flat] {
fs::write(path, "kind: mixin\n").unwrap();
}
assert!(collect_subdir_mixins(&root).is_empty());
assert_eq!(collect_flat_mixins(&root).len(), 1);
assert_eq!(collect_mixins(&root, &[ScanMode::Flat]), vec![flat.clone()]);
assert_eq!(
collect_mixins(&root, &[ScanMode::SubdirNamed]),
vec![subdir_named.clone()]
);
assert_eq!(
collect_mixins(&root, &[ScanMode::SubdirFlat]),
vec![subdir_flat.clone()]
);
let all = collect_mixins(
&root,
&[ScanMode::Flat, ScanMode::SubdirNamed, ScanMode::SubdirFlat],
);
assert_eq!(all, vec![flat, subdir_named, subdir_flat]);
let mut deduped = all.clone();
deduped.sort();
deduped.dedup();
assert_eq!(
deduped.len(),
all.len(),
"no mixin may be collected twice: {all:?}"
);
let _ = fs::remove_dir_all(&root);
}
#[test]
fn collect_flat_mixins_tolerates_a_missing_directory() {
fn flat_scan_tolerates_a_missing_directory() {
let root = unique_root("flat-missing");
let absent = root.join("nope");
assert!(collect_flat_mixins(&absent).is_empty());
assert!(collect_mixins(&absent, &[ScanMode::Flat]).is_empty());
let _ = fs::remove_dir_all(&root);
}
/// `generate_rag_sbx_mixin` writes an agent-scoped RAG sidecar next to the
/// rag yaml, at `<agents>/<agent>/<rag>.sbx-mixin.yaml`. Before `SubdirFlat`
/// existed, nothing scanned that shape and attaching a Qdrant RAG from
/// inside an agent produced no network allow rule and no credential.
#[test]
fn agent_scoped_rag_sidecar_is_discovered() {
let root = unique_root("agent-scoped-rag");
let agent = root.join("researcher");
fs::create_dir_all(&agent).unwrap();
fs::write(agent.join("company-docs.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
fs::write(agent.join("company-docs.yaml"), "driver: qdrant\n").unwrap();
let found = collect_mixins(&root, &[ScanMode::SubdirNamed, ScanMode::SubdirFlat]);
assert_eq!(found, vec![agent.join("company-docs.sbx-mixin.yaml")]);
let _ = fs::remove_dir_all(&root);
}
#[test]
fn agent_level_mixin_and_rag_sidecars_are_both_discovered() {
let root = unique_root("agent-both-shapes");
let agent = root.join("researcher");
fs::create_dir_all(&agent).unwrap();
fs::write(agent.join("sbx-mixin.yaml"), "kind: mixin\n").unwrap();
fs::write(agent.join("zebra.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
fs::write(agent.join("alpha.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
let found = collect_mixins(&root, &[ScanMode::SubdirNamed, ScanMode::SubdirFlat]);
assert_eq!(
file_names(&found),
vec![
"sbx-mixin.yaml",
"alpha.sbx-mixin.yaml",
"zebra.sbx-mixin.yaml"
]
);
let _ = fs::remove_dir_all(&root);
}
#[test]
fn subdir_flat_scan_ignores_a_directory_named_like_a_mixin() {
let root = unique_root("subdir-flat-decoy");
let agent = root.join("researcher");
fs::create_dir_all(agent.join("decoy.sbx-mixin.yaml")).unwrap();
assert!(collect_mixins(&root, &[ScanMode::SubdirFlat]).is_empty());
let _ = fs::remove_dir_all(&root);
}
+106 -20
View File
@@ -337,29 +337,20 @@ fn inject_rag_secrets(vault: &Vault, registered: &HashSet<String>) -> Result<()>
if !data.attached {
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;
};
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) {
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."
);
let service_id = mcp_credentials::secret_service_id(&stem);
if !service_id.is_empty() && !registered.contains(&service_id) {
bind_rag_secret(vault, &service_id, primary, &stem)?;
}
for name in extra {
let id = mcp_credentials::secret_service_id(name);
if !id.is_empty() && !registered.contains(&id) {
bind_rag_secret(vault, &id, name, &stem)?;
}
}
}
@@ -367,6 +358,43 @@ fn inject_rag_secrets(vault: &Vault, registered: &HashSet<String>) -> Result<()>
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 {
match provider_type {
"claude" => "anthropic".to_string(),
@@ -646,6 +674,64 @@ fn chown_agent_recursive(sandbox: &str, path: &str) -> Result<()> {
#[cfg(test)]
mod tests {
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};
#[test]