Compare commits
27
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d31110cd67
|
||
|
|
0f35e03a85
|
||
|
|
c2b0c120d7
|
||
|
|
65c9be36b2
|
||
|
|
2658ca776e
|
||
|
|
f8682102a0
|
||
|
|
3fa0f5c428
|
||
|
|
68135b97d1
|
||
|
|
b87a3460c4 | ||
|
|
cb23da6490
|
||
|
|
2a40a5a81d | ||
|
|
ebba976a27 | ||
|
|
c84f9522e9 | ||
|
|
81ed769f8a | ||
|
|
6f7defe25f | ||
|
|
b837f82d7e | ||
|
|
54685be9a2 | ||
|
|
4dd6e794b2 | ||
|
|
af9622d31c | ||
|
|
6f586bd535 | ||
|
|
78740db170 | ||
|
|
64d594f4ee | ||
|
|
1322d73c7b | ||
|
|
de91ffa517 | ||
|
|
74bc613d94 | ||
|
|
6d0a5550fe | ||
|
|
7b1c0342b4 |
@@ -5,3 +5,6 @@
|
||||
.idea/
|
||||
/coyote.iml
|
||||
/.idea/
|
||||
.coyote/**
|
||||
.sisyphus/**
|
||||
.coyote-project.json
|
||||
|
||||
+15
@@ -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 && \
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -33,7 +33,11 @@ source "$LLM_PROMPT_UTILS_FILE"
|
||||
|
||||
# shellcheck disable=SC2154
|
||||
main() {
|
||||
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")"
|
||||
# Command substitution strips *all* trailing newlines and `jq -r` appends one
|
||||
# of its own, so read with `-j` and pin the real end of the content with a
|
||||
# sentinel that is removed afterwards.
|
||||
argc_contents="$(jq -j '.content' <<< "$LLM_TOOL_RAW_JSON"; printf x)"
|
||||
argc_contents="${argc_contents%x}"
|
||||
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
|
||||
|
||||
if [[ ! -f "$argc_path" ]]; then
|
||||
@@ -41,7 +45,11 @@ main() {
|
||||
exit 1
|
||||
fi
|
||||
|
||||
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"))"
|
||||
# Same sentinel guard on the patched result, otherwise the trailing newline
|
||||
# is stripped again on the way back out. `rc` preserves patch_file's exit
|
||||
# status so a failure still aborts under `set -e`.
|
||||
new_contents="$(patch_file "$argc_path" <(printf "%s" "$argc_contents"); rc=$?; printf x; exit "$rc")"
|
||||
new_contents="${new_contents%x}"
|
||||
printf "%s" "$new_contents" | git diff --no-index "$argc_path" - || true
|
||||
|
||||
guard_operation "Apply changes?"
|
||||
|
||||
@@ -15,7 +15,12 @@ source "$LLM_PROMPT_UTILS_FILE"
|
||||
|
||||
# shellcheck disable=SC2154
|
||||
main() {
|
||||
argc_contents="$(jq -r '.content' <<< "$LLM_TOOL_RAW_JSON")"
|
||||
# Command substitution strips *all* trailing newlines and `jq -r` appends one
|
||||
# of its own, so read with `-j` and pin the real end of the content with a
|
||||
# sentinel that is removed afterwards. Without this every written file loses
|
||||
# its final newline, which breaks formatters such as `cargo fmt --check`.
|
||||
argc_contents="$(jq -j '.content' <<< "$LLM_TOOL_RAW_JSON"; printf x)"
|
||||
argc_contents="${argc_contents%x}"
|
||||
argc_path="$(jq -r '.path' <<< "$LLM_TOOL_RAW_JSON")"
|
||||
|
||||
if [[ -f "$argc_path" ]]; then
|
||||
|
||||
@@ -841,6 +841,11 @@
|
||||
referrer: coyote
|
||||
echo_pkce_in_token_exchange: true
|
||||
models:
|
||||
- name: grok-4.6
|
||||
input_price: 2
|
||||
output_price: 6
|
||||
max_input_tokens: 500000
|
||||
supports_function_calling: true
|
||||
- name: grok-4.5
|
||||
input_price: 2
|
||||
output_price: 6
|
||||
|
||||
+161
-1
@@ -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
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
+324
-41
@@ -2,13 +2,13 @@ use super::access_token::{is_valid_access_token, set_access_token};
|
||||
use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider;
|
||||
use super::{ClientConfig, ProviderModels};
|
||||
use crate::config::paths;
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use anyhow::{Context, Error, Result, anyhow, bail};
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||
use chrono::Utc;
|
||||
use indexmap::IndexMap;
|
||||
use inquire::Text;
|
||||
use reqwest::{Client as ReqwestClient, RequestBuilder};
|
||||
use reqwest::{Client as ReqwestClient, RequestBuilder, StatusCode};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
@@ -16,6 +16,10 @@ use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::Duration;
|
||||
use tokio::sync;
|
||||
use url::Url;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -197,10 +201,20 @@ pub struct OAuthTokens {
|
||||
pub account_id: Option<String>,
|
||||
}
|
||||
|
||||
const TOKEN_ENDPOINT_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> {
|
||||
match provider.flow() {
|
||||
OAuthFlow::Pkce => run_pkce_flow(provider, client_name).await,
|
||||
OAuthFlow::ClientCredentials => run_client_credentials_flow(provider, client_name).await,
|
||||
OAuthFlow::ClientCredentials => {
|
||||
run_client_credentials_flow(provider, client_name).await?;
|
||||
println!(
|
||||
"Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.",
|
||||
client_name,
|
||||
provider.provider_name()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
OAuthFlow::DeviceCode => run_device_code_flow(provider, client_name).await,
|
||||
}
|
||||
}
|
||||
@@ -301,12 +315,20 @@ async fn run_pkce_flow(provider: &dyn OAuthProvider, client_name: &str) -> Resul
|
||||
|
||||
let access_token = response["access_token"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow!("Missing access_token in response: {response}"))?
|
||||
.ok_or_else(|| {
|
||||
anyhow!(
|
||||
"Missing access_token in response (keys: {})",
|
||||
token_response_keys(&response)
|
||||
)
|
||||
})?
|
||||
.to_string();
|
||||
let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string());
|
||||
let expires_in = response["expires_in"]
|
||||
.as_i64()
|
||||
.ok_or_else(|| anyhow!("Missing expires_in in response: {response}"))?;
|
||||
let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
|
||||
anyhow!(
|
||||
"Missing expires_in in response (keys: {})",
|
||||
token_response_keys(&response)
|
||||
)
|
||||
})?;
|
||||
|
||||
let expires_at = Utc::now().timestamp() + expires_in;
|
||||
|
||||
@@ -334,7 +356,9 @@ async fn run_client_credentials_flow(
|
||||
provider: &dyn OAuthProvider,
|
||||
client_name: &str,
|
||||
) -> Result<()> {
|
||||
let client = ReqwestClient::new();
|
||||
let client = ReqwestClient::builder()
|
||||
.timeout(TOKEN_ENDPOINT_TIMEOUT)
|
||||
.build()?;
|
||||
let scopes = provider.scopes();
|
||||
let mut params: Vec<(&str, &str)> = vec![
|
||||
("grant_type", "client_credentials"),
|
||||
@@ -349,11 +373,19 @@ async fn run_client_credentials_flow(
|
||||
|
||||
let access_token = response["access_token"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow!("Missing access_token in client_credentials response: {response}"))?
|
||||
.ok_or_else(|| {
|
||||
anyhow!(
|
||||
"Missing access_token in client_credentials response (keys: {})",
|
||||
token_response_keys(&response)
|
||||
)
|
||||
})?
|
||||
.to_string();
|
||||
let expires_in = response["expires_in"]
|
||||
.as_i64()
|
||||
.ok_or_else(|| anyhow!("Missing expires_in in client_credentials response: {response}"))?;
|
||||
let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
|
||||
anyhow!(
|
||||
"Missing expires_in in client_credentials response (keys: {})",
|
||||
token_response_keys(&response)
|
||||
)
|
||||
})?;
|
||||
let expires_at = Utc::now().timestamp() + expires_in;
|
||||
|
||||
let tokens = OAuthTokens {
|
||||
@@ -363,11 +395,6 @@ async fn run_client_credentials_flow(
|
||||
account_id: provider.extract_account_id(&response),
|
||||
};
|
||||
save_oauth_tokens(client_name, &tokens)?;
|
||||
println!(
|
||||
"Successfully authenticated client '{}' with {} via OAuth (client_credentials). Tokens saved.",
|
||||
client_name,
|
||||
provider.provider_name()
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -417,19 +444,28 @@ async fn run_device_code_flow(provider: &dyn OAuthProvider, client_name: &str) -
|
||||
let device_code = device_response["device_code"]
|
||||
.as_str()
|
||||
.ok_or_else(|| {
|
||||
anyhow!("Missing device_code in device authorization response: {device_response}")
|
||||
anyhow!(
|
||||
"Missing device_code in device authorization response (keys: {})",
|
||||
token_response_keys(&device_response)
|
||||
)
|
||||
})?
|
||||
.to_string();
|
||||
let user_code = device_response["user_code"]
|
||||
.as_str()
|
||||
.ok_or_else(|| {
|
||||
anyhow!("Missing user_code in device authorization response: {device_response}")
|
||||
anyhow!(
|
||||
"Missing user_code in device authorization response (keys: {})",
|
||||
token_response_keys(&device_response)
|
||||
)
|
||||
})?
|
||||
.to_string();
|
||||
let verification_uri = device_response["verification_uri"]
|
||||
.as_str()
|
||||
.ok_or_else(|| {
|
||||
anyhow!("Missing verification_uri in device authorization response: {device_response}")
|
||||
anyhow!(
|
||||
"Missing verification_uri in device authorization response (keys: {})",
|
||||
token_response_keys(&device_response)
|
||||
)
|
||||
})?
|
||||
.to_string();
|
||||
let verification_uri_complete = device_response["verification_uri_complete"]
|
||||
@@ -551,10 +587,76 @@ fn save_oauth_tokens(client_name: &str, tokens: &OAuthTokens) -> Result<()> {
|
||||
fs::create_dir_all(parent)?;
|
||||
}
|
||||
let json = serde_json::to_string_pretty(tokens)?;
|
||||
fs::write(path, json)?;
|
||||
// Write-then-rename so a crash mid-write never truncates the live token file.
|
||||
let mut tmp = path.clone().into_os_string();
|
||||
tmp.push(".tmp");
|
||||
let tmp = PathBuf::from(tmp);
|
||||
// Tokens are live credentials: create the file owner-only, not umask-default.
|
||||
let mut options = fs::OpenOptions::new();
|
||||
options.write(true).create(true).truncate(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
options.mode(0o600);
|
||||
}
|
||||
options.open(&tmp)?.write_all(json.as_bytes())?;
|
||||
fs::rename(&tmp, &path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn token_response_keys(response: &Value) -> String {
|
||||
match response.as_object() {
|
||||
Some(map) => {
|
||||
let keys: Vec<&str> = map.keys().map(String::as_str).collect();
|
||||
format!("[{}]", keys.join(", "))
|
||||
}
|
||||
None => "<non-object response>".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_refresh_response(
|
||||
status: StatusCode,
|
||||
response: &Value,
|
||||
previous_refresh_token: Option<&str>,
|
||||
) -> Result<(String, Option<String>, i64)> {
|
||||
if let Some(error) = response["error"].as_str() {
|
||||
let description = response["error_description"]
|
||||
.as_str()
|
||||
.unwrap_or("no description");
|
||||
if matches!(error, "invalid_grant" | "invalid_token") {
|
||||
bail!(
|
||||
"OAuth refresh token was rejected ({error}: {description}). Please re-authenticate."
|
||||
);
|
||||
}
|
||||
bail!("Token refresh failed ({error}: {description})");
|
||||
}
|
||||
if !status.is_success() {
|
||||
bail!("Token refresh failed with HTTP status {status}");
|
||||
}
|
||||
|
||||
let access_token = response["access_token"]
|
||||
.as_str()
|
||||
.ok_or_else(|| {
|
||||
anyhow!(
|
||||
"Missing access_token in refresh response (keys: {})",
|
||||
token_response_keys(response)
|
||||
)
|
||||
})?
|
||||
.to_string();
|
||||
let refresh_token = response["refresh_token"]
|
||||
.as_str()
|
||||
.or(previous_refresh_token)
|
||||
.map(str::to_string);
|
||||
let expires_in = response["expires_in"].as_i64().ok_or_else(|| {
|
||||
anyhow!(
|
||||
"Missing expires_in in refresh response (keys: {})",
|
||||
token_response_keys(response)
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok((access_token, refresh_token, expires_in))
|
||||
}
|
||||
|
||||
pub async fn refresh_oauth_token(
|
||||
client: &ReqwestClient,
|
||||
provider: &dyn OAuthProvider,
|
||||
@@ -577,19 +679,23 @@ pub async fn refresh_oauth_token(
|
||||
],
|
||||
);
|
||||
|
||||
let response: Value = request.send().await?.json().await?;
|
||||
let (status, response) = tokio::time::timeout(TOKEN_ENDPOINT_TIMEOUT, async {
|
||||
let response = request.send().await?;
|
||||
let status = response.status();
|
||||
let body: Value = response.json().await?;
|
||||
Ok::<_, Error>((status, body))
|
||||
})
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow!(
|
||||
"Token refresh for '{}' timed out after {}s",
|
||||
client_name,
|
||||
TOKEN_ENDPOINT_TIMEOUT.as_secs()
|
||||
)
|
||||
})??;
|
||||
|
||||
let access_token = response["access_token"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow!("Missing access_token in refresh response: {response}"))?
|
||||
.to_string();
|
||||
let refresh_token = response["refresh_token"]
|
||||
.as_str()
|
||||
.map(|s| s.to_string())
|
||||
.or_else(|| tokens.refresh_token.clone());
|
||||
let expires_in = response["expires_in"]
|
||||
.as_i64()
|
||||
.ok_or_else(|| anyhow!("Missing expires_in in refresh response: {response}"))?;
|
||||
let (access_token, refresh_token, expires_in) =
|
||||
parse_refresh_response(status, &response, tokens.refresh_token.as_deref())?;
|
||||
|
||||
let expires_at = Utc::now().timestamp() + expires_in;
|
||||
|
||||
@@ -609,6 +715,20 @@ pub async fn refresh_oauth_token(
|
||||
Ok(new_tokens)
|
||||
}
|
||||
|
||||
/// Per-client lock so concurrent requests perform a single refresh.
|
||||
/// Returns a clone of the Arc so the parking_lot guard is dropped before the
|
||||
/// caller awaits on the tokio mutex.
|
||||
fn refresh_guard(client_name: &str) -> Arc<sync::Mutex<()>> {
|
||||
static GUARDS: OnceLock<parking_lot::Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
|
||||
OnceLock::new();
|
||||
GUARDS
|
||||
.get_or_init(Default::default)
|
||||
.lock()
|
||||
.entry(client_name.to_string())
|
||||
.or_default()
|
||||
.clone()
|
||||
}
|
||||
|
||||
pub async fn prepare_oauth_access_token(
|
||||
client: &ReqwestClient,
|
||||
provider: &dyn OAuthProvider,
|
||||
@@ -624,15 +744,35 @@ pub async fn prepare_oauth_access_token(
|
||||
};
|
||||
|
||||
let tokens = if Utc::now().timestamp() >= tokens.expires_at {
|
||||
match provider.flow() {
|
||||
OAuthFlow::Pkce | OAuthFlow::DeviceCode => {
|
||||
refresh_oauth_token(client, provider, client_name, &tokens).await?
|
||||
}
|
||||
OAuthFlow::ClientCredentials => {
|
||||
run_client_credentials_flow(provider, client_name).await?;
|
||||
load_oauth_tokens(client_name)
|
||||
.ok_or_else(|| anyhow!("Token file missing after client_credentials refresh"))?
|
||||
let guard = refresh_guard(client_name);
|
||||
let _guard = guard.lock().await;
|
||||
|
||||
// A concurrent caller may have refreshed while we waited for the
|
||||
// lock; a valid in-memory token means the winner already populated
|
||||
// the cache.
|
||||
if is_valid_access_token(client_name) {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let tokens = match load_oauth_tokens(client_name) {
|
||||
Some(t) => t,
|
||||
None => return Ok(false),
|
||||
};
|
||||
|
||||
if Utc::now().timestamp() >= tokens.expires_at {
|
||||
match provider.flow() {
|
||||
OAuthFlow::Pkce | OAuthFlow::DeviceCode => {
|
||||
refresh_oauth_token(client, provider, client_name, &tokens).await?
|
||||
}
|
||||
OAuthFlow::ClientCredentials => {
|
||||
run_client_credentials_flow(provider, client_name).await?;
|
||||
load_oauth_tokens(client_name).ok_or_else(|| {
|
||||
anyhow!("Token file missing after client_credentials refresh")
|
||||
})?
|
||||
}
|
||||
}
|
||||
} else {
|
||||
tokens
|
||||
}
|
||||
} else {
|
||||
tokens
|
||||
@@ -886,11 +1026,54 @@ pub(crate) fn client_config_info(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::ffi::OsString;
|
||||
use std::path::PathBuf;
|
||||
use std::str;
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use super::*;
|
||||
use crate::client::openai_compatible::OpenAICompatibleConfig;
|
||||
use crate::client::{ModelData, ProviderModels};
|
||||
use crate::utils::get_env_name;
|
||||
use serial_test::serial;
|
||||
use std::{env, time::SystemTime};
|
||||
|
||||
fn with_temp_cache<F: FnOnce()>(f: F) {
|
||||
struct Restore {
|
||||
key: String,
|
||||
prev: Option<OsString>,
|
||||
root: PathBuf,
|
||||
}
|
||||
impl Drop for Restore {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
match self.prev.take() {
|
||||
Some(v) => env::set_var(&self.key, v),
|
||||
None => env::remove_var(&self.key),
|
||||
}
|
||||
}
|
||||
let _ = fs::remove_dir_all(&self.root);
|
||||
}
|
||||
}
|
||||
|
||||
let unique = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
let root = env::temp_dir().join(format!("coyote-client-oauth-test-{unique}"));
|
||||
fs::create_dir_all(&root).unwrap();
|
||||
let env_key = get_env_name("cache_dir");
|
||||
let prev = env::var_os(&env_key);
|
||||
unsafe {
|
||||
env::set_var(&env_key, &root);
|
||||
}
|
||||
let _restore = Restore {
|
||||
key: env_key,
|
||||
prev,
|
||||
root,
|
||||
};
|
||||
f();
|
||||
}
|
||||
|
||||
fn base_config() -> OAuthConfig {
|
||||
OAuthConfig {
|
||||
@@ -1468,4 +1651,104 @@ scopes:
|
||||
"body missing grant_type param: {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn save_oauth_tokens_roundtrips_and_leaves_no_tmp_file() {
|
||||
with_temp_cache(|| {
|
||||
let tokens = OAuthTokens {
|
||||
access_token: "at-123".into(),
|
||||
refresh_token: Some("rt-456".into()),
|
||||
expires_at: 1234567890,
|
||||
account_id: Some("acct-789".into()),
|
||||
};
|
||||
|
||||
save_oauth_tokens("atomic-test", &tokens).unwrap();
|
||||
|
||||
let loaded = load_oauth_tokens("atomic-test").unwrap();
|
||||
assert_eq!(loaded.access_token, "at-123");
|
||||
assert_eq!(loaded.refresh_token.as_deref(), Some("rt-456"));
|
||||
assert_eq!(loaded.expires_at, 1234567890);
|
||||
assert_eq!(loaded.account_id.as_deref(), Some("acct-789"));
|
||||
|
||||
let dir = paths::oauth_tokens_dir();
|
||||
let leftover_tmp = fs::read_dir(&dir)
|
||||
.unwrap()
|
||||
.any(|e| e.unwrap().file_name().to_string_lossy().ends_with(".tmp"));
|
||||
assert!(!leftover_tmp, "temp file left behind in {dir:?}");
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let mode = fs::metadata(paths::token_file("atomic-test"))
|
||||
.unwrap()
|
||||
.permissions()
|
||||
.mode();
|
||||
assert_eq!(mode & 0o777, 0o600, "token file mode was {mode:o}");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_refresh_response_invalid_grant_redacts_and_prompts_reauth() {
|
||||
let response = serde_json::json!({
|
||||
"error": "invalid_grant",
|
||||
"error_description": "refresh token revoked",
|
||||
"refresh_token": "planted-secret-token",
|
||||
});
|
||||
|
||||
let err = parse_refresh_response(StatusCode::BAD_REQUEST, &response, Some("old-rt"))
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
|
||||
assert!(err.contains("re-authenticate"), "unexpected error: {err}");
|
||||
assert!(
|
||||
!err.contains("planted-secret-token"),
|
||||
"error leaked token material: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_response_keys_lists_keys_without_values() {
|
||||
let response = serde_json::json!({
|
||||
"access_token": "secret-at",
|
||||
"token_type": "SecretBearer",
|
||||
});
|
||||
|
||||
let keys = token_response_keys(&response);
|
||||
|
||||
assert!(keys.contains("access_token"), "missing key name: {keys}");
|
||||
assert!(keys.contains("token_type"), "missing key name: {keys}");
|
||||
assert!(!keys.contains("secret-at"), "leaked value: {keys}");
|
||||
assert!(!keys.contains("SecretBearer"), "leaked value: {keys}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_refresh_response_rotates_refresh_token_when_present() {
|
||||
let response = serde_json::json!({
|
||||
"access_token": "new-at",
|
||||
"refresh_token": "new-rt",
|
||||
"expires_in": 3600,
|
||||
});
|
||||
|
||||
let (access_token, refresh_token, expires_in) =
|
||||
parse_refresh_response(StatusCode::OK, &response, Some("old-rt")).unwrap();
|
||||
|
||||
assert_eq!(access_token, "new-at");
|
||||
assert_eq!(refresh_token.as_deref(), Some("new-rt"));
|
||||
assert_eq!(expires_in, 3600);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_refresh_response_keeps_old_refresh_token_when_absent() {
|
||||
let response = serde_json::json!({
|
||||
"access_token": "new-at",
|
||||
"expires_in": 3600,
|
||||
});
|
||||
|
||||
let (_, refresh_token, _) =
|
||||
parse_refresh_response(StatusCode::OK, &response, Some("old-rt")).unwrap();
|
||||
|
||||
assert_eq!(refresh_token.as_deref(), Some("old-rt"));
|
||||
}
|
||||
}
|
||||
|
||||
+12
-3
@@ -4,6 +4,7 @@ use crate::{
|
||||
client::Model,
|
||||
config::memory,
|
||||
function::{Functions, run_llm_function},
|
||||
graph, rag,
|
||||
};
|
||||
|
||||
use super::rag_cache::RagKey;
|
||||
@@ -185,7 +186,7 @@ impl Agent {
|
||||
&rag_path_clone,
|
||||
&document_paths,
|
||||
abort,
|
||||
false,
|
||||
true,
|
||||
)
|
||||
.await
|
||||
})
|
||||
@@ -247,6 +248,10 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
if rag.is_some() && app.function_calling_support && graph_for_rag.is_none() {
|
||||
functions.append_rag_query_functions();
|
||||
}
|
||||
|
||||
agent_config.replace_tools_placeholder(&functions);
|
||||
|
||||
Ok(Self {
|
||||
@@ -1021,11 +1026,11 @@ async fn init_graph_rags(
|
||||
// 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) = crate::graph::validator::rag_driver_error(driver)
|
||||
&& let Some(message) = graph::validator::rag_driver_error(driver)
|
||||
{
|
||||
bail!("rag node '{node_id}': {message}");
|
||||
}
|
||||
let config = rag_init_config(rag_node);
|
||||
let mut config = rag_init_config(rag_node);
|
||||
let fully_specified = config.embedding_model.is_some()
|
||||
&& config.chunk_size.is_some()
|
||||
&& config.chunk_overlap.is_some();
|
||||
@@ -1051,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 =
|
||||
|
||||
@@ -103,7 +103,7 @@ impl McpFactory {
|
||||
}
|
||||
|
||||
let bearer_token = if spec.is_remote() {
|
||||
oauth::load_valid_mcp_token(name)
|
||||
oauth::load_or_refresh_mcp_token(name).await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
+8
-1
@@ -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"
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
use super::{FunctionDeclaration, JsonSchema};
|
||||
use crate::config::RequestContext;
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use indexmap::IndexMap;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
pub const RAG_FUNCTION_PREFIX: &str = "rag__";
|
||||
|
||||
pub fn rag_query_function_declarations() -> Vec<FunctionDeclaration> {
|
||||
vec![FunctionDeclaration {
|
||||
name: format!("{RAG_FUNCTION_PREFIX}query"),
|
||||
description: "Search the RAG knowledge base attached to this session and return \
|
||||
the most relevant text chunks with their source paths. The relevant \
|
||||
context has already been injected into the prompt up-front; use this \
|
||||
tool to pull additional context on-demand when the initial retrieval \
|
||||
does not fully answer the question. Prefer specific, keyword-rich queries."
|
||||
.to_string(),
|
||||
parameters: JsonSchema {
|
||||
type_value: Some("object".to_string()),
|
||||
properties: Some(IndexMap::from([
|
||||
(
|
||||
"query".to_string(),
|
||||
JsonSchema {
|
||||
type_value: Some("string".to_string()),
|
||||
description: Some(
|
||||
"Natural language search query used to retrieve relevant chunks."
|
||||
.into(),
|
||||
),
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
(
|
||||
"top_k".to_string(),
|
||||
JsonSchema {
|
||||
type_value: Some("integer".to_string()),
|
||||
description: Some(
|
||||
"Maximum number of chunks to return. Defaults to the RAG's \
|
||||
configured top_k when omitted."
|
||||
.into(),
|
||||
),
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
])),
|
||||
required: Some(vec!["query".to_string()]),
|
||||
..Default::default()
|
||||
},
|
||||
agent: false,
|
||||
}]
|
||||
}
|
||||
|
||||
pub async fn handle_rag_tool(
|
||||
ctx: &mut RequestContext,
|
||||
cmd_name: &str,
|
||||
args: &Value,
|
||||
) -> Result<Value> {
|
||||
let action = cmd_name
|
||||
.strip_prefix(RAG_FUNCTION_PREFIX)
|
||||
.unwrap_or(cmd_name);
|
||||
|
||||
match action {
|
||||
"query" => handle_query(ctx, args).await,
|
||||
_ => Err(anyhow!("Unknown RAG action: {action}")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_query(ctx: &RequestContext, args: &Value) -> Result<Value> {
|
||||
let rag = ctx
|
||||
.rag
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow!("No RAG is attached to this session"))?;
|
||||
|
||||
let query = args
|
||||
.get("query")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| anyhow!("'query' is required"))?;
|
||||
|
||||
let top_k = args
|
||||
.get("top_k")
|
||||
.and_then(Value::as_u64)
|
||||
.map(|v| v as usize)
|
||||
.unwrap_or_else(|| rag.configured_top_k());
|
||||
|
||||
let rerank_model = rag.configured_reranker().map(|s| s.to_string());
|
||||
|
||||
let chunks = rag
|
||||
.search_chunks(query, top_k, rerank_model.as_deref())
|
||||
.await?;
|
||||
|
||||
let chunks_json: Vec<Value> = chunks
|
||||
.into_iter()
|
||||
.map(|(text, source)| json!({ "text": text, "source": source }))
|
||||
.collect();
|
||||
|
||||
Ok(json!({
|
||||
"rag_name": rag.name(),
|
||||
"count": chunks_json.len(),
|
||||
"chunks": chunks_json,
|
||||
}))
|
||||
}
|
||||
+12
@@ -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)?;
|
||||
|
||||
@@ -0,0 +1,443 @@
|
||||
use crate::cli::{Cli, McpScopeArg, McpTransportArg};
|
||||
use crate::config::{ensure_parent_exists, paths};
|
||||
use crate::mcp::{JsonField, McpOAuthConfig, McpServer, McpServersConfig, McpTransportType};
|
||||
use crate::vault::{SECRET_RE, Vault};
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use indexmap::{IndexMap, IndexSet};
|
||||
use inquire::Confirm;
|
||||
use std::collections::HashSet;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
impl From<McpTransportArg> for McpTransportType {
|
||||
fn from(value: McpTransportArg) -> Self {
|
||||
match value {
|
||||
McpTransportArg::Stdio => McpTransportType::Stdio,
|
||||
McpTransportArg::Http => McpTransportType::Http,
|
||||
McpTransportArg::Sse => McpTransportType::Sse,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn handle(cli: &Cli, vault: &Vault) -> Result<()> {
|
||||
if cli.mcp_list {
|
||||
return handle_list(cli.scope);
|
||||
}
|
||||
if let Some(name) = &cli.mcp_get {
|
||||
return handle_get(name, cli.scope);
|
||||
}
|
||||
if let Some(name) = &cli.mcp_remove {
|
||||
return handle_remove(name, cli.scope, cli.mcp_force);
|
||||
}
|
||||
if let Some(name) = &cli.mcp_add {
|
||||
return handle_add(cli, name, vault);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_list(scope: Option<McpScopeArg>) -> Result<()> {
|
||||
let show_user = scope != Some(McpScopeArg::Workspace);
|
||||
let show_workspace = scope != Some(McpScopeArg::User);
|
||||
|
||||
if show_user {
|
||||
let user_path = paths::mcp_config_file();
|
||||
let user_cfg = load_config_raw(&user_path)?;
|
||||
println!("User ({})", user_path.display());
|
||||
print_server_list(&user_cfg);
|
||||
}
|
||||
|
||||
if show_workspace {
|
||||
match paths::workspace_mcp_config_file() {
|
||||
Some(ws_path) => {
|
||||
let ws_cfg = load_config_raw(&ws_path)?;
|
||||
if show_user {
|
||||
println!();
|
||||
}
|
||||
println!("Workspace ({})", ws_path.display());
|
||||
print_server_list(&ws_cfg);
|
||||
}
|
||||
None if scope == Some(McpScopeArg::Workspace) => {
|
||||
println!("Workspace: no mcp.json found in current directory");
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn print_server_list(cfg: &McpServersConfig) {
|
||||
if cfg.mcp_servers.is_empty() {
|
||||
println!(" (none)");
|
||||
return;
|
||||
}
|
||||
let name_width = cfg.mcp_servers.keys().map(String::len).max().unwrap_or(0);
|
||||
for (name, spec) in &cfg.mcp_servers {
|
||||
let transport = match spec.transport_type {
|
||||
McpTransportType::Stdio => "stdio",
|
||||
McpTransportType::Http => "http",
|
||||
McpTransportType::Sse => "sse",
|
||||
};
|
||||
let target = spec.url.clone().unwrap_or_else(|| {
|
||||
let cmd = spec.command.clone().unwrap_or_default();
|
||||
let args = spec.args.as_ref().map(|a| a.join(" ")).unwrap_or_default();
|
||||
if args.is_empty() {
|
||||
cmd
|
||||
} else {
|
||||
format!("{cmd} {args}")
|
||||
}
|
||||
});
|
||||
println!(
|
||||
" {name:<name_width$} {transport:<5} {target}",
|
||||
name_width = name_width
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_get(name: &str, scope: Option<McpScopeArg>) -> Result<()> {
|
||||
let (path, cfg) = load_for_scope_or_search(name, scope)?;
|
||||
let spec = cfg
|
||||
.mcp_servers
|
||||
.get(name)
|
||||
.ok_or_else(|| anyhow!("MCP server '{name}' not found"))?;
|
||||
let pretty =
|
||||
serde_json::to_string_pretty(spec).context("failed to serialize MCP server config")?;
|
||||
println!("# {}", path.display());
|
||||
println!("{pretty}");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_remove(name: &str, scope: Option<McpScopeArg>, force: bool) -> Result<()> {
|
||||
let (path, mut cfg) = load_for_scope_or_search(name, scope)?;
|
||||
if !force {
|
||||
let ok = Confirm::new(&format!(
|
||||
"Remove MCP server '{name}' from {}?",
|
||||
path.display()
|
||||
))
|
||||
.with_default(false)
|
||||
.prompt()?;
|
||||
|
||||
if !ok {
|
||||
println!("Aborted.");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
cfg.mcp_servers.shift_remove(name);
|
||||
save_config(&path, &cfg)?;
|
||||
println!("✓ Removed MCP server '{name}' from {}", path.display());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_add(cli: &Cli, name: &str, vault: &Vault) -> Result<()> {
|
||||
validate_name(name)?;
|
||||
let server = build_server(cli)?;
|
||||
server.validate(name)?;
|
||||
|
||||
let scope = cli.scope.unwrap_or_default();
|
||||
let path = write_path_for_scope(scope);
|
||||
let mut cfg = load_config_raw(&path)?;
|
||||
|
||||
if cfg.mcp_servers.contains_key(name) && !cli.mcp_force {
|
||||
let ok = Confirm::new(&format!(
|
||||
"MCP server '{name}' already exists in {}. Overwrite?",
|
||||
path.display()
|
||||
))
|
||||
.with_default(false)
|
||||
.prompt()?;
|
||||
if !ok {
|
||||
println!("Aborted. Use --mcp-force to overwrite without prompting.");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
provision_secrets(cli, vault)?;
|
||||
|
||||
cfg.mcp_servers.insert(name.to_string(), server);
|
||||
save_config(&path, &cfg)?;
|
||||
println!("✓ Added MCP server '{name}' to {}", path.display());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_name(name: &str) -> Result<()> {
|
||||
if name.is_empty() {
|
||||
bail!("MCP server name cannot be empty");
|
||||
}
|
||||
|
||||
if !name
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
|
||||
{
|
||||
bail!("Invalid MCP server name '{name}': only letters, digits, '-', and '_' are allowed");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn build_server(cli: &Cli) -> Result<McpServer> {
|
||||
let has_command = !cli.mcp_command.is_empty();
|
||||
let has_url = cli.url.is_some();
|
||||
|
||||
let transport = cli
|
||||
.transport
|
||||
.map(McpTransportType::from)
|
||||
.unwrap_or_else(|| {
|
||||
if has_command {
|
||||
McpTransportType::Stdio
|
||||
} else {
|
||||
McpTransportType::Http
|
||||
}
|
||||
});
|
||||
|
||||
match transport {
|
||||
McpTransportType::Stdio => build_stdio(cli, has_url),
|
||||
McpTransportType::Http | McpTransportType::Sse => build_remote(cli, transport, has_command),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_stdio(cli: &Cli, has_url: bool) -> Result<McpServer> {
|
||||
if cli.mcp_command.is_empty() {
|
||||
bail!(
|
||||
"stdio MCP server requires a command. Pass it after `--`, e.g. \
|
||||
`--mcp-add NAME -- npx some-server --flag`"
|
||||
);
|
||||
}
|
||||
if has_url {
|
||||
bail!("stdio MCP server does not accept --url");
|
||||
}
|
||||
if !cli.header.is_empty() {
|
||||
bail!("stdio MCP server does not accept --header");
|
||||
}
|
||||
if cli.client_id.is_some()
|
||||
|| cli.client_secret.is_some()
|
||||
|| cli.callback_port.is_some()
|
||||
|| cli.redirect_host.is_some()
|
||||
{
|
||||
bail!("stdio MCP server does not accept OAuth flags");
|
||||
}
|
||||
|
||||
let (cmd, args) = cli.mcp_command.split_first().unwrap();
|
||||
|
||||
let mut env: IndexMap<String, JsonField> = IndexMap::new();
|
||||
for kv in &cli.env {
|
||||
let (k, v) = kv
|
||||
.split_once('=')
|
||||
.ok_or_else(|| anyhow!("invalid --env value '{kv}': expected KEY=VALUE"))?;
|
||||
if k.is_empty() {
|
||||
bail!("invalid --env value '{kv}': KEY cannot be empty");
|
||||
}
|
||||
env.insert(k.to_string(), JsonField::Str(v.to_string()));
|
||||
}
|
||||
|
||||
Ok(McpServer {
|
||||
transport_type: McpTransportType::Stdio,
|
||||
command: Some(cmd.clone()),
|
||||
args: (!args.is_empty()).then(|| args.to_vec()),
|
||||
env: (!env.is_empty()).then_some(env),
|
||||
cwd: cli.cwd.clone(),
|
||||
url: None,
|
||||
headers: None,
|
||||
oauth: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn build_remote(cli: &Cli, transport: McpTransportType, has_command: bool) -> Result<McpServer> {
|
||||
if has_command {
|
||||
bail!(
|
||||
"http/sse MCP server does not accept a trailing `-- <cmd>`. Use `--url` \
|
||||
to specify the endpoint."
|
||||
);
|
||||
}
|
||||
let url = cli
|
||||
.url
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow!("http/sse MCP server requires --url <URL>"))?;
|
||||
if !cli.env.is_empty() {
|
||||
bail!("http/sse MCP server does not accept --env; use --header instead");
|
||||
}
|
||||
if cli.cwd.is_some() {
|
||||
bail!("http/sse MCP server does not accept --cwd");
|
||||
}
|
||||
|
||||
let mut headers: IndexMap<String, String> = IndexMap::new();
|
||||
for h in &cli.header {
|
||||
let (name, value) = h
|
||||
.split_once(':')
|
||||
.ok_or_else(|| anyhow!("invalid --header value '{h}': expected 'Name: Value'"))?;
|
||||
let name = name.trim();
|
||||
let value = value.trim_start_matches(' ');
|
||||
if name.is_empty() {
|
||||
bail!("invalid --header value '{h}': header name cannot be empty");
|
||||
}
|
||||
headers.insert(name.to_string(), value.to_string());
|
||||
}
|
||||
|
||||
let oauth = if cli.client_id.is_some()
|
||||
|| cli.client_secret.is_some()
|
||||
|| cli.callback_port.is_some()
|
||||
|| cli.redirect_host.is_some()
|
||||
{
|
||||
Some(McpOAuthConfig {
|
||||
client_id: cli.client_id.clone(),
|
||||
client_secret: cli.client_secret.clone(),
|
||||
callback_port: cli.callback_port,
|
||||
redirect_host: cli.redirect_host.clone(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(McpServer {
|
||||
transport_type: transport,
|
||||
command: None,
|
||||
args: None,
|
||||
env: None,
|
||||
cwd: None,
|
||||
url: Some(url),
|
||||
headers: (!headers.is_empty()).then_some(headers),
|
||||
oauth,
|
||||
})
|
||||
}
|
||||
|
||||
fn provision_secrets(cli: &Cli, vault: &Vault) -> Result<()> {
|
||||
let mut sources: Vec<&str> = Vec::new();
|
||||
if let Some(s) = cli.url.as_deref() {
|
||||
sources.push(s);
|
||||
}
|
||||
if let Some(s) = cli.client_secret.as_deref() {
|
||||
sources.push(s);
|
||||
}
|
||||
if let Some(s) = cli.client_id.as_deref() {
|
||||
sources.push(s);
|
||||
}
|
||||
if let Some(s) = cli.redirect_host.as_deref() {
|
||||
sources.push(s);
|
||||
}
|
||||
if let Some(s) = cli.cwd.as_deref() {
|
||||
sources.push(s);
|
||||
}
|
||||
sources.extend(cli.env.iter().map(String::as_str));
|
||||
sources.extend(cli.header.iter().map(String::as_str));
|
||||
|
||||
let mut needed: IndexSet<String> = IndexSet::new();
|
||||
for value in sources {
|
||||
for caps in SECRET_RE.captures_iter(value).filter_map(Result::ok) {
|
||||
if let Some(m) = caps.get(1) {
|
||||
needed.insert(m.as_str().trim().to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if needed.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let existing: HashSet<String> = vault.list_secrets(false)?.into_iter().collect();
|
||||
for name in needed {
|
||||
if existing.contains(&name) {
|
||||
continue;
|
||||
}
|
||||
eprintln!("Value references vault secret {{{{ {name} }}}} which is not stored yet.");
|
||||
let ok = Confirm::new(&format!("Add '{name}' to the vault now?"))
|
||||
.with_default(true)
|
||||
.prompt()?;
|
||||
if !ok {
|
||||
bail!(
|
||||
"Vault secret '{name}' is required by the config; aborting. \
|
||||
Add it later with `coyote --add-secret {name}`."
|
||||
);
|
||||
}
|
||||
vault.add_secret(&name)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn load_for_scope_or_search(
|
||||
name: &str,
|
||||
scope: Option<McpScopeArg>,
|
||||
) -> Result<(PathBuf, McpServersConfig)> {
|
||||
if let Some(s) = scope {
|
||||
let path = match s {
|
||||
McpScopeArg::User => paths::mcp_config_file(),
|
||||
McpScopeArg::Workspace => paths::workspace_mcp_config_file()
|
||||
.ok_or_else(|| anyhow!("no workspace mcp.json found in the current directory"))?,
|
||||
};
|
||||
let cfg = load_config_raw(&path)?;
|
||||
if !cfg.mcp_servers.contains_key(name) {
|
||||
bail!(
|
||||
"MCP server '{name}' not found in {} scope ({})",
|
||||
scope_label(s),
|
||||
path.display()
|
||||
);
|
||||
}
|
||||
|
||||
return Ok((path, cfg));
|
||||
}
|
||||
|
||||
let user_path = paths::mcp_config_file();
|
||||
let user_cfg = load_config_raw(&user_path)?;
|
||||
if user_cfg.mcp_servers.contains_key(name) {
|
||||
return Ok((user_path, user_cfg));
|
||||
}
|
||||
|
||||
if let Some(ws_path) = paths::workspace_mcp_config_file() {
|
||||
let ws_cfg = load_config_raw(&ws_path)?;
|
||||
if ws_cfg.mcp_servers.contains_key(name) {
|
||||
return Ok((ws_path, ws_cfg));
|
||||
}
|
||||
}
|
||||
|
||||
bail!("MCP server '{name}' not found in any scope");
|
||||
}
|
||||
|
||||
fn write_path_for_scope(scope: McpScopeArg) -> PathBuf {
|
||||
match scope {
|
||||
McpScopeArg::User => paths::mcp_config_file(),
|
||||
McpScopeArg::Workspace => paths::workspace_mcp_config_file()
|
||||
.unwrap_or_else(|| paths::workspace_config_dir().join("mcp.json")),
|
||||
}
|
||||
}
|
||||
|
||||
fn scope_label(scope: McpScopeArg) -> &'static str {
|
||||
match scope {
|
||||
McpScopeArg::User => "user",
|
||||
McpScopeArg::Workspace => "workspace",
|
||||
}
|
||||
}
|
||||
|
||||
fn load_config_raw(path: &Path) -> Result<McpServersConfig> {
|
||||
if !path.exists() {
|
||||
return Ok(McpServersConfig {
|
||||
mcp_servers: IndexMap::new(),
|
||||
});
|
||||
}
|
||||
|
||||
let raw = fs::read_to_string(path)
|
||||
.with_context(|| format!("failed to read MCP config at {}", path.display()))?;
|
||||
if raw.trim().is_empty() {
|
||||
return Ok(McpServersConfig {
|
||||
mcp_servers: IndexMap::new(),
|
||||
});
|
||||
}
|
||||
|
||||
serde_json::from_str(&raw)
|
||||
.with_context(|| format!("failed to parse MCP config at {}", path.display()))
|
||||
}
|
||||
|
||||
fn save_config(path: &Path, config: &McpServersConfig) -> Result<()> {
|
||||
ensure_parent_exists(path)?;
|
||||
let serialized =
|
||||
serde_json::to_string_pretty(config).context("failed to serialize MCP config")?;
|
||||
let tmp = path.with_extension("json.tmp");
|
||||
fs::write(&tmp, &serialized)
|
||||
.with_context(|| format!("failed to write temporary MCP config at {}", tmp.display()))?;
|
||||
fs::rename(&tmp, path)
|
||||
.with_context(|| format!("failed to finalize MCP config at {}", path.display()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
+4
-1
@@ -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")]
|
||||
@@ -324,7 +327,7 @@ impl McpRegistry {
|
||||
.with_context(|| format!("MCP server not found in config: {id}"))?;
|
||||
|
||||
let bearer_token = if spec.is_remote() {
|
||||
oauth::load_valid_mcp_token(&id)
|
||||
oauth::load_or_refresh_mcp_token(&id).await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
+300
-28
@@ -1,15 +1,25 @@
|
||||
use crate::client::oauth::{OAuthProvider, TokenRequestFormat, load_oauth_tokens, run_oauth_flow};
|
||||
use crate::client::oauth::{
|
||||
OAuthProvider, OAuthTokens, TokenRequestFormat, load_oauth_tokens, refresh_oauth_token,
|
||||
run_oauth_flow, token_response_keys,
|
||||
};
|
||||
use crate::config::paths;
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use chrono::Utc;
|
||||
use inquire::Text;
|
||||
use log::warn;
|
||||
use log::{debug, warn};
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::net::TcpListener;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::sync;
|
||||
use url::Url;
|
||||
|
||||
const REFRESH_HTTP_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const REFRESH_FAILURE_BACKOFF: Duration = Duration::from_secs(60);
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ProtectedResourceMetadata {
|
||||
#[serde(default)]
|
||||
@@ -34,6 +44,10 @@ struct McpRegistration {
|
||||
client_id: String,
|
||||
#[serde(default)]
|
||||
redirect_uri: Option<String>,
|
||||
#[serde(default)]
|
||||
token_url: Option<String>,
|
||||
#[serde(default)]
|
||||
resource: Option<String>,
|
||||
}
|
||||
|
||||
struct DiscoveredOAuth {
|
||||
@@ -124,8 +138,19 @@ pub async fn run_mcp_oauth_flow(
|
||||
None
|
||||
};
|
||||
|
||||
let (client_id, redirect_uri) = if let Some(reused) = cached_reuse {
|
||||
reused
|
||||
let (client_id, redirect_uri) = if let Some((client_id, redirect_uri)) = cached_reuse {
|
||||
// Re-save so registrations cached before token_url/resource were
|
||||
// persisted gain them, enabling token refresh next time.
|
||||
if let Err(e) = save_registration(
|
||||
server_name,
|
||||
&client_id,
|
||||
&redirect_uri,
|
||||
&metadata.token_endpoint,
|
||||
&resource,
|
||||
) {
|
||||
debug!("Failed to update cached MCP registration for '{server_name}': {e}");
|
||||
}
|
||||
(client_id, redirect_uri)
|
||||
} else {
|
||||
let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0));
|
||||
let listener = TcpListener::bind(&bind_addr)?;
|
||||
@@ -137,10 +162,7 @@ pub async fn run_mcp_oauth_flow(
|
||||
id.to_string()
|
||||
} else if let Some(reg_endpoint) = &metadata.registration_endpoint {
|
||||
match register_client(reg_endpoint, &redirect_uri).await {
|
||||
Ok(id) => {
|
||||
let _ = save_registration(server_name, &id, &redirect_uri);
|
||||
id
|
||||
}
|
||||
Ok(id) => id,
|
||||
Err(e) => {
|
||||
warn!("Dynamic client registration failed: {e}. Falling back to manual entry.");
|
||||
Text::new("Enter the OAuth client ID for this MCP server:")
|
||||
@@ -153,6 +175,18 @@ pub async fn run_mcp_oauth_flow(
|
||||
.prompt()
|
||||
.context("Failed to read client ID")?
|
||||
};
|
||||
// Persist regardless of how the client_id was obtained (DCR, config,
|
||||
// or manual entry) so refresh_mcp_token can run the refresh_token
|
||||
// grant later without interactive re-auth.
|
||||
if let Err(e) = save_registration(
|
||||
server_name,
|
||||
&client_id,
|
||||
&redirect_uri,
|
||||
&metadata.token_endpoint,
|
||||
&resource,
|
||||
) {
|
||||
debug!("Failed to cache MCP registration for '{server_name}': {e}");
|
||||
}
|
||||
(client_id, redirect_uri)
|
||||
};
|
||||
|
||||
@@ -168,12 +202,113 @@ pub async fn run_mcp_oauth_flow(
|
||||
run_oauth_flow(&provider, &mcp_token_key(server_name)).await
|
||||
}
|
||||
|
||||
pub fn load_valid_mcp_token(server_name: &str) -> Option<String> {
|
||||
let tokens = load_oauth_tokens(&mcp_token_key(server_name))?;
|
||||
pub async fn load_or_refresh_mcp_token(server_name: &str) -> Option<String> {
|
||||
let key = mcp_token_key(server_name);
|
||||
let tokens = load_oauth_tokens(&key)?;
|
||||
if Utc::now().timestamp() < tokens.expires_at {
|
||||
Some(tokens.access_token)
|
||||
} else {
|
||||
None
|
||||
return Some(tokens.access_token);
|
||||
}
|
||||
|
||||
if in_refresh_failure_backoff(server_name) {
|
||||
debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed");
|
||||
return None;
|
||||
}
|
||||
|
||||
let lock = refresh_lock(server_name);
|
||||
let _guard = lock.lock().await;
|
||||
|
||||
// A concurrent caller may have refreshed while we waited for the lock.
|
||||
let tokens = load_oauth_tokens(&key)?;
|
||||
if Utc::now().timestamp() < tokens.expires_at {
|
||||
return Some(tokens.access_token);
|
||||
}
|
||||
|
||||
if in_refresh_failure_backoff(server_name) {
|
||||
debug!("Skipping token refresh for MCP server '{server_name}': recent attempt failed");
|
||||
return None;
|
||||
}
|
||||
|
||||
match refresh_mcp_token(server_name, &key, &tokens).await {
|
||||
Ok(access_token) => Some(access_token),
|
||||
Err(e) => {
|
||||
note_refresh_failure(server_name);
|
||||
warn!(
|
||||
"Failed to refresh OAuth token for MCP server '{server_name}'. \
|
||||
Run `.mcp auth {server_name}` to re-authenticate."
|
||||
);
|
||||
debug!(
|
||||
"Token refresh error for MCP server '{server_name}': {}",
|
||||
redact_refresh_error(&e)
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn refresh_mcp_token(server_name: &str, key: &str, tokens: &OAuthTokens) -> Result<String> {
|
||||
if tokens.refresh_token.is_none() {
|
||||
return Err(anyhow!("no refresh token stored"));
|
||||
}
|
||||
|
||||
let reg =
|
||||
load_registration(server_name).ok_or_else(|| anyhow!("no cached client registration"))?;
|
||||
let token_url = reg.token_url.ok_or_else(|| {
|
||||
anyhow!("cached registration has no token URL (saved by an older version)")
|
||||
})?;
|
||||
let resource = reg.resource.ok_or_else(|| {
|
||||
anyhow!("cached registration has no resource (saved by an older version)")
|
||||
})?;
|
||||
|
||||
let provider = McpOAuthProvider {
|
||||
client_id: reg.client_id,
|
||||
authorize_url: String::new(),
|
||||
token_url,
|
||||
scopes: String::new(),
|
||||
fixed_redirect: String::new(),
|
||||
resource,
|
||||
};
|
||||
|
||||
let client = Client::builder().timeout(REFRESH_HTTP_TIMEOUT).build()?;
|
||||
let refreshed = refresh_oauth_token(&client, &provider, key, tokens).await?;
|
||||
Ok(refreshed.access_token)
|
||||
}
|
||||
|
||||
fn refresh_lock(server_name: &str) -> Arc<sync::Mutex<()>> {
|
||||
static LOCKS: OnceLock<parking_lot::Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
|
||||
OnceLock::new();
|
||||
LOCKS
|
||||
.get_or_init(Default::default)
|
||||
.lock()
|
||||
.entry(server_name.to_string())
|
||||
.or_default()
|
||||
.clone()
|
||||
}
|
||||
|
||||
fn refresh_failures() -> &'static parking_lot::Mutex<HashMap<String, Instant>> {
|
||||
static FAILURES: OnceLock<parking_lot::Mutex<HashMap<String, Instant>>> = OnceLock::new();
|
||||
FAILURES.get_or_init(Default::default)
|
||||
}
|
||||
|
||||
fn note_refresh_failure(server_name: &str) {
|
||||
refresh_failures()
|
||||
.lock()
|
||||
.insert(server_name.to_string(), Instant::now());
|
||||
}
|
||||
|
||||
fn in_refresh_failure_backoff(server_name: &str) -> bool {
|
||||
refresh_failures()
|
||||
.lock()
|
||||
.get(server_name)
|
||||
.is_some_and(|failed_at| failed_at.elapsed() < REFRESH_FAILURE_BACKOFF)
|
||||
}
|
||||
|
||||
/// Refresh errors may embed the token endpoint's JSON response, which can
|
||||
/// contain live tokens; strip everything from the first `{` before logging.
|
||||
fn redact_refresh_error(e: &anyhow::Error) -> String {
|
||||
let msg = e.to_string();
|
||||
match msg.find('{') {
|
||||
Some(idx) => format!("{}<response body redacted>", &msg[..idx]),
|
||||
None => msg,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -187,7 +322,13 @@ fn load_registration(server_name: &str) -> Option<McpRegistration> {
|
||||
serde_json::from_str(&content).ok()
|
||||
}
|
||||
|
||||
fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) -> Result<()> {
|
||||
fn save_registration(
|
||||
server_name: &str,
|
||||
client_id: &str,
|
||||
redirect_uri: &str,
|
||||
token_url: &str,
|
||||
resource: &str,
|
||||
) -> Result<()> {
|
||||
let dir = paths::oauth_tokens_dir();
|
||||
fs::create_dir_all(&dir)?;
|
||||
|
||||
@@ -195,6 +336,8 @@ fn save_registration(server_name: &str, client_id: &str, redirect_uri: &str) ->
|
||||
let reg = McpRegistration {
|
||||
client_id: client_id.to_string(),
|
||||
redirect_uri: Some(redirect_uri.to_string()),
|
||||
token_url: Some(token_url.to_string()),
|
||||
resource: Some(resource.to_string()),
|
||||
};
|
||||
|
||||
fs::write(path, serde_json::to_string_pretty(®)?)?;
|
||||
@@ -244,7 +387,12 @@ async fn register_client(endpoint: &str, redirect_uri: &str) -> Result<String> {
|
||||
|
||||
response["client_id"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow!("Missing client_id in registration response: {response}"))
|
||||
.ok_or_else(|| {
|
||||
anyhow!(
|
||||
"Missing client_id in registration response (keys: {})",
|
||||
token_response_keys(&response)
|
||||
)
|
||||
})
|
||||
.map(|s| s.to_string())
|
||||
}
|
||||
|
||||
@@ -426,11 +574,31 @@ mod tests {
|
||||
use crate::utils::get_env_name;
|
||||
use serial_test::serial;
|
||||
use std::{
|
||||
env, fs,
|
||||
env,
|
||||
ffi::OsString,
|
||||
fs,
|
||||
path::PathBuf,
|
||||
time::{self, SystemTime},
|
||||
};
|
||||
|
||||
fn with_temp_cache<F: FnOnce()>(f: F) {
|
||||
struct Restore {
|
||||
key: String,
|
||||
prev: Option<OsString>,
|
||||
root: PathBuf,
|
||||
}
|
||||
impl Drop for Restore {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
match self.prev.take() {
|
||||
Some(v) => env::set_var(&self.key, v),
|
||||
None => env::remove_var(&self.key),
|
||||
}
|
||||
}
|
||||
let _ = fs::remove_dir_all(&self.root);
|
||||
}
|
||||
}
|
||||
|
||||
let unique = SystemTime::now()
|
||||
.duration_since(time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
@@ -442,14 +610,12 @@ mod tests {
|
||||
unsafe {
|
||||
env::set_var(&env_key, &root);
|
||||
}
|
||||
let _restore = Restore {
|
||||
key: env_key,
|
||||
prev,
|
||||
root,
|
||||
};
|
||||
f();
|
||||
unsafe {
|
||||
match prev {
|
||||
Some(v) => env::set_var(&env_key, v),
|
||||
None => env::remove_var(&env_key),
|
||||
}
|
||||
}
|
||||
let _ = fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -685,12 +851,19 @@ mod tests {
|
||||
"notion",
|
||||
"client-xyz-123",
|
||||
"http://127.0.0.1:49152/callback",
|
||||
"https://as.example/token",
|
||||
"https://mcp.example/mcp",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let loaded = load_registration("notion");
|
||||
let loaded = load_registration("notion").unwrap();
|
||||
|
||||
assert_eq!(loaded.unwrap().client_id, "client-xyz-123");
|
||||
assert_eq!(loaded.client_id, "client-xyz-123");
|
||||
assert_eq!(
|
||||
loaded.token_url.as_deref(),
|
||||
Some("https://as.example/token")
|
||||
);
|
||||
assert_eq!(loaded.resource.as_deref(), Some("https://mcp.example/mcp"));
|
||||
});
|
||||
}
|
||||
|
||||
@@ -708,8 +881,22 @@ mod tests {
|
||||
#[serial]
|
||||
fn registration_second_save_overwrites_first() {
|
||||
with_temp_cache(|| {
|
||||
save_registration("github", "first-id", "http://127.0.0.1:49152/callback").unwrap();
|
||||
save_registration("github", "second-id", "http://127.0.0.1:49153/callback").unwrap();
|
||||
save_registration(
|
||||
"github",
|
||||
"first-id",
|
||||
"http://127.0.0.1:49152/callback",
|
||||
"https://as.example/token",
|
||||
"https://mcp.example/mcp",
|
||||
)
|
||||
.unwrap();
|
||||
save_registration(
|
||||
"github",
|
||||
"second-id",
|
||||
"http://127.0.0.1:49153/callback",
|
||||
"https://as.example/token",
|
||||
"https://mcp.example/mcp",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let loaded = load_registration("github").unwrap();
|
||||
|
||||
@@ -737,6 +924,8 @@ mod tests {
|
||||
|
||||
assert_eq!(loaded.client_id, "legacy-id");
|
||||
assert_eq!(loaded.redirect_uri, None);
|
||||
assert_eq!(loaded.token_url, None);
|
||||
assert_eq!(loaded.resource, None);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -744,7 +933,14 @@ mod tests {
|
||||
#[serial]
|
||||
fn save_registration_persists_redirect_uri() {
|
||||
with_temp_cache(|| {
|
||||
save_registration("aws", "client-abc", "http://127.0.0.1:49152/callback").unwrap();
|
||||
save_registration(
|
||||
"aws",
|
||||
"client-abc",
|
||||
"http://127.0.0.1:49152/callback",
|
||||
"https://as.example/token",
|
||||
"https://mcp.example/mcp",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let loaded = load_registration("aws").unwrap();
|
||||
|
||||
@@ -756,6 +952,82 @@ mod tests {
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_registration_deserializes_without_new_fields_and_roundtrips() {
|
||||
let old: McpRegistration = serde_json::from_str(r#"{"client_id":"legacy-id"}"#).unwrap();
|
||||
|
||||
assert_eq!(old.client_id, "legacy-id");
|
||||
assert_eq!(old.token_url, None);
|
||||
assert_eq!(old.resource, None);
|
||||
|
||||
let full = McpRegistration {
|
||||
client_id: "client-abc".into(),
|
||||
redirect_uri: Some("http://127.0.0.1:49152/callback".into()),
|
||||
token_url: Some("https://as.example/token".into()),
|
||||
resource: Some("https://mcp.example/mcp".into()),
|
||||
};
|
||||
let json = serde_json::to_string(&full).unwrap();
|
||||
let back: McpRegistration = serde_json::from_str(&json).unwrap();
|
||||
|
||||
assert_eq!(back.token_url.as_deref(), Some("https://as.example/token"));
|
||||
assert_eq!(back.resource.as_deref(), Some("https://mcp.example/mcp"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn expired_token_with_old_format_registration_returns_none() {
|
||||
with_temp_cache(|| {
|
||||
let dir = paths::oauth_tokens_dir();
|
||||
fs::create_dir_all(&dir).unwrap();
|
||||
fs::write(
|
||||
paths::token_file("mcp_legacyref"),
|
||||
r#"{"access_token":"stale","refresh_token":"refresh-abc","expires_at":0}"#,
|
||||
)
|
||||
.unwrap();
|
||||
fs::write(
|
||||
dir.join("mcp_legacyref_registration.json"),
|
||||
r#"{"client_id":"legacy-id"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
let token = rt.block_on(load_or_refresh_mcp_token("legacyref"));
|
||||
|
||||
assert_eq!(token, None);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refresh_failure_backoff_memoizes_per_server() {
|
||||
assert!(!in_refresh_failure_backoff("backoff-test-server"));
|
||||
|
||||
note_refresh_failure("backoff-test-server");
|
||||
|
||||
assert!(in_refresh_failure_backoff("backoff-test-server"));
|
||||
assert!(!in_refresh_failure_backoff("backoff-other-server"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_refresh_error_strips_response_body() {
|
||||
let with_body = anyhow!(
|
||||
"Missing access_token in refresh response: {}",
|
||||
r#"{"access_token":"live-secret"}"#
|
||||
);
|
||||
let without_body = anyhow!("no refresh token stored");
|
||||
|
||||
assert_eq!(
|
||||
redact_refresh_error(&with_body),
|
||||
"Missing access_token in refresh response: <response body redacted>"
|
||||
);
|
||||
assert_eq!(
|
||||
redact_refresh_error(&without_body),
|
||||
"no refresh token stored"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_redirect_port_matches() {
|
||||
let port = cached_redirect_port("http://127.0.0.1:49152/callback", "127.0.0.1", None);
|
||||
|
||||
+77
-38
@@ -267,31 +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.",
|
||||
];
|
||||
let sel = Select::new("RAG storage driver:", options)
|
||||
.with_starting_cursor(0)
|
||||
.prompt()?;
|
||||
if sel.starts_with("duckdb") {
|
||||
println!(
|
||||
"Note: several Coyote processes can query a duckdb RAG at the same time, \
|
||||
but while one process is ingesting or rebuilding it the others cannot \
|
||||
read it until that finishes. Changing its driver later means deleting \
|
||||
and recreating the RAG."
|
||||
);
|
||||
"duckdb"
|
||||
} else {
|
||||
"yaml"
|
||||
}
|
||||
select_rag_driver()?
|
||||
} else {
|
||||
"yaml"
|
||||
"yaml".to_string()
|
||||
};
|
||||
let reranker_model = app.rag_reranker_model.clone();
|
||||
let top_k = app.rag_top_k;
|
||||
@@ -318,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() {
|
||||
@@ -586,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 =
|
||||
@@ -856,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,
|
||||
@@ -1169,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)
|
||||
@@ -1185,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();
|
||||
@@ -1842,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>> {
|
||||
@@ -3005,7 +3044,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn reciprocal_rank_fusion_empty_lists() {
|
||||
let result = super::reciprocal_rank_fusion(vec![], vec![], 5);
|
||||
let result = reciprocal_rank_fusion(vec![], vec![], 5);
|
||||
assert!(result.is_empty(), "empty input should produce empty output");
|
||||
}
|
||||
|
||||
@@ -3013,7 +3052,7 @@ mod tests {
|
||||
fn reciprocal_rank_fusion_deduplicates_across_signals() {
|
||||
let doc_a = DocumentId::new(0, 0);
|
||||
let doc_b = DocumentId::new(0, 1);
|
||||
let result = super::reciprocal_rank_fusion(
|
||||
let result = reciprocal_rank_fusion(
|
||||
vec![vec![doc_a, doc_b], vec![doc_a, doc_b]],
|
||||
vec![1.0, 1.0],
|
||||
5,
|
||||
@@ -3030,7 +3069,7 @@ mod tests {
|
||||
#[test]
|
||||
fn reciprocal_rank_fusion_respects_top_k() {
|
||||
let docs: Vec<DocumentId> = (0..10).map(|i| DocumentId::new(0, i)).collect();
|
||||
let result = super::reciprocal_rank_fusion(vec![docs], vec![1.0], 3);
|
||||
let result = reciprocal_rank_fusion(vec![docs], vec![1.0], 3);
|
||||
assert_eq!(result.len(), 3, "result should be capped at top_k=3");
|
||||
}
|
||||
|
||||
@@ -3038,7 +3077,7 @@ mod tests {
|
||||
fn reciprocal_rank_fusion_weights_affect_ranking() {
|
||||
let doc_a = DocumentId::new(0, 0);
|
||||
let doc_b = DocumentId::new(0, 1);
|
||||
let result = super::reciprocal_rank_fusion(
|
||||
let result = reciprocal_rank_fusion(
|
||||
vec![vec![doc_a, doc_b], vec![doc_b, doc_a]],
|
||||
vec![10.0, 1.0],
|
||||
2,
|
||||
|
||||
+96
-11
@@ -9,6 +9,7 @@ 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
|
||||
@@ -104,7 +105,7 @@ fn parse_search_hits(
|
||||
let score = pt["score"].as_f64()? as f32;
|
||||
Some((interner.document_id(&pt["id"])?, score))
|
||||
})
|
||||
.filter(|(_, score)| *score > min_score)
|
||||
.filter(|(_, score)| min_score <= 0.0 || *score > min_score)
|
||||
.collect())
|
||||
}
|
||||
|
||||
@@ -183,7 +184,22 @@ pub struct QdrantProvider {
|
||||
}
|
||||
|
||||
impl QdrantProvider {
|
||||
fn make_client(api_key: Option<&str>) -> Result<Client> {
|
||||
fn skips_proxy(base_url: &str) -> bool {
|
||||
let Ok(url) = Url::parse(base_url) else {
|
||||
return false;
|
||||
};
|
||||
match url.host() {
|
||||
Some(Host::Domain(name)) => {
|
||||
name == "localhost" || name.ends_with(".localhost") || name.ends_with(".local")
|
||||
}
|
||||
Some(Host::Ipv4(ip)) => ip.is_loopback() || ip.is_private() || ip.is_link_local(),
|
||||
// No stable is_unique_local, so fc00::/7 is matched directly.
|
||||
Some(Host::Ipv6(ip)) => ip.is_loopback() || ip.segments()[0] & 0xfe00 == 0xfc00,
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_client(base_url: &str, api_key: Option<&str>) -> Result<Client> {
|
||||
let mut headers = HeaderMap::new();
|
||||
if let Some(key) = api_key {
|
||||
let mut value =
|
||||
@@ -191,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 {
|
||||
@@ -219,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()
|
||||
@@ -237,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()
|
||||
@@ -260,7 +277,7 @@ impl QdrantProvider {
|
||||
|
||||
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()
|
||||
@@ -310,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 });
|
||||
|
||||
@@ -353,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,
|
||||
@@ -576,6 +594,73 @@ mod tests {
|
||||
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
|
||||
|
||||
+34
-2
@@ -39,8 +39,10 @@ static INLINE_CODE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"`([^`\n]+
|
||||
static IMAGE_RE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"!\[([^\]]*)\]\(([^)]+)\)").unwrap());
|
||||
static LINK_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\[([^\]]+)\]\(([^)]+)\)").unwrap());
|
||||
static BOLD_AST_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\*\*([^*\n]+)\*\*").unwrap());
|
||||
static BOLD_US_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"__([^_\n]+)__").unwrap());
|
||||
static BOLD_AST_RE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"\*\*((?:[^*\n]|\*(?!\*))+?)\*\*").unwrap());
|
||||
static BOLD_US_RE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"__((?:[^_\n]|_(?!_))+?)__").unwrap());
|
||||
static ITALIC_AST_RE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"(?<![*\w])\*(?!\s)([^*\n]+?)(?<!\s)\*(?!\*)").unwrap());
|
||||
static ITALIC_US_RE: LazyLock<Regex> =
|
||||
@@ -2419,6 +2421,36 @@ std::error::Error>> {
|
||||
assert!(result.contains("\x1b[9m"), "strikethrough SGR: {result:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bold_asterisk_wraps_italic_asterisk() {
|
||||
let styles = test_styles();
|
||||
|
||||
let result = apply_inline("**loud *soft* loud**", &styles);
|
||||
|
||||
assert!(
|
||||
!result.contains("**"),
|
||||
"outer bold markers stripped: {result:?}"
|
||||
);
|
||||
assert!(result.contains("\x1b[1m"), "bold SGR present: {result:?}");
|
||||
assert!(result.contains("\x1b[3m"), "italic SGR present: {result:?}");
|
||||
assert!(result.contains("soft"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bold_underscore_wraps_italic_underscore() {
|
||||
let styles = test_styles();
|
||||
|
||||
let result = apply_inline("__loud _soft_ loud__", &styles);
|
||||
|
||||
assert!(
|
||||
!result.contains("__"),
|
||||
"outer bold markers stripped: {result:?}"
|
||||
);
|
||||
assert!(result.contains("\x1b[1m"), "bold SGR present: {result:?}");
|
||||
assert!(result.contains("\x1b[3m"), "italic SGR present: {result:?}");
|
||||
assert!(result.contains("soft"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bold_wraps_inline_code() {
|
||||
let styles = test_styles();
|
||||
|
||||
+1
-1
@@ -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>"),
|
||||
},
|
||||
|
||||
+318
-7
@@ -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,33 +36,152 @@ 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();
|
||||
|
||||
@@ -556,6 +677,196 @@ 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]
|
||||
|
||||
+106
-20
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user