Compare commits

..
Author SHA1 Message Date
ZakiandClaude Opus 4.6 5b0d261398 feat(mcp): support custom HTTP headers for MCP server auth (#639)
Some MCP servers (e.g., Browser-Use) require custom headers instead of
OAuth for authentication. Add a `headers` field to McpServerConfig that
injects custom HTTP headers into every request to that server.

Changes:
- Add `headers: HashMap<String, String>` to McpServerConfig with serde
  default/skip_serializing_if for backward compatibility
- Add `with_headers()` builder and `new_with_config()` constructor
- Inject custom headers in McpClient::send_request() before auth header
- Add `--header` / `-H` CLI flag to `ironclaw mcp add` (Key:Value format)
- Update app.rs and extensions/manager.rs to use new_with_config()
- Show custom header names in verbose `mcp list` output

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-03-07 18:26:30 -08:00
7 changed files with 233 additions and 161 deletions
+1 -1
View File
@@ -542,7 +542,7 @@ impl AppBuilder {
server, mcp_sm, secrets, "default",
)
} else {
McpClient::new_with_name(&server_name, &server.url)
McpClient::new_with_config(server.clone())
};
match client.list_tools().await {
+32 -1
View File
@@ -47,6 +47,10 @@ pub enum McpCommand {
/// Server description
#[arg(long)]
description: Option<String>,
/// Custom HTTP headers (format: "Key:Value", can be repeated)
#[arg(long = "header", short = 'H')]
headers: Vec<String>,
},
/// Remove an MCP server
@@ -108,6 +112,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
token_url,
scopes,
description,
headers,
} => {
add_server(
name,
@@ -117,6 +122,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
token_url,
scopes,
description,
headers,
)
.await
}
@@ -133,6 +139,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
}
/// Add a new MCP server.
#[allow(clippy::too_many_arguments)]
async fn add_server(
name: String,
url: String,
@@ -141,6 +148,7 @@ async fn add_server(
token_url: Option<String>,
scopes: Option<String>,
description: Option<String>,
headers: Vec<String>,
) -> anyhow::Result<()> {
let mut config = McpServerConfig::new(&name, &url);
@@ -148,6 +156,18 @@ async fn add_server(
config = config.with_description(desc);
}
// Parse custom headers (format: "Key:Value")
if !headers.is_empty() {
let mut header_map = std::collections::HashMap::new();
for h in &headers {
let (key, value) = h.split_once(':').ok_or_else(|| {
anyhow::anyhow!("Invalid header format '{}'. Expected 'Key:Value'.", h)
})?;
header_map.insert(key.trim().to_string(), value.trim().to_string());
}
config = config.with_headers(header_map);
}
// Track if auth is required
let requires_auth = client_id.is_some();
@@ -242,6 +262,17 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
if let Some(ref desc) = server.description {
println!(" Description: {}", desc);
}
if !server.headers.is_empty() {
println!(
" Custom headers: {}",
server
.headers
.keys()
.cloned()
.collect::<Vec<_>>()
.join(", ")
);
}
if let Some(ref oauth) = server.oauth {
println!(" OAuth Client ID: {}", oauth.client_id);
if !oauth.scopes.is_empty() {
@@ -374,7 +405,7 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
return Ok(());
} else {
// No OAuth and no tokens - try unauthenticated
McpClient::new_with_name(&server.name, &server.url)
McpClient::new_with_config(server.clone())
};
// Test connection
-102
View File
@@ -169,18 +169,10 @@ pub(crate) fn parse_timestamp(s: &str) -> Result<DateTime<Utc>, String> {
}
// Naive with fractional seconds (legacy or SQLite datetime() output)
if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") {
tracing::warn!(
timestamp = s,
"parsing naive timestamp without timezone; assuming UTC — consider re-running migrations"
);
return Ok(ndt.and_utc());
}
// Naive without fractional seconds (legacy format)
if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
tracing::warn!(
timestamp = s,
"parsing naive timestamp without timezone; assuming UTC — consider re-running migrations"
);
return Ok(ndt.and_utc());
}
Err(format!("unparseable timestamp: {:?}", s))
@@ -518,98 +510,4 @@ mod tests {
);
}
}
#[test]
fn test_parse_timestamp_rfc3339() {
use super::parse_timestamp;
// Standard RFC 3339 with Z suffix
let dt = parse_timestamp("2024-01-15T10:30:00.123Z").unwrap();
assert_eq!(
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
"2024-01-15T10:30:00.123Z"
);
// RFC 3339 with +00:00 offset
let dt = parse_timestamp("2024-01-15T10:30:00.000+00:00").unwrap();
assert_eq!(
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
"2024-01-15T10:30:00.000Z"
);
}
#[test]
fn test_parse_timestamp_naive_fallback() {
use super::parse_timestamp;
// Naive with fractional seconds (legacy datetime('now') output)
let dt = parse_timestamp("2024-01-15 10:30:00.123").unwrap();
assert_eq!(
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
"2024-01-15T10:30:00.123Z"
);
// Naive without fractional seconds
let dt = parse_timestamp("2024-01-15 10:30:00").unwrap();
assert_eq!(
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
"2024-01-15T10:30:00.000Z"
);
}
#[test]
fn test_parse_timestamp_invalid() {
use super::parse_timestamp;
assert!(parse_timestamp("not-a-timestamp").is_err());
assert!(parse_timestamp("").is_err());
}
#[tokio::test]
async fn test_default_timestamps_are_rfc3339() {
// Verify that DEFAULT column values produce RFC 3339 timestamps
// after the migration change from datetime('now') to strftime.
// Use file-based DB because in-memory doesn't share schema across connections.
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test_ts.db");
let backend = LibSqlBackend::new_local(&db_path).await.unwrap();
backend.run_migrations().await.unwrap();
let conn = backend.connect().await.unwrap();
let id = uuid::Uuid::new_v4().to_string();
conn.execute(
"INSERT INTO conversations (id, channel, user_id) VALUES (?1, ?2, ?3)",
libsql::params![id.clone(), "test", "user1"],
)
.await
.unwrap();
let mut rows = conn
.query(
"SELECT started_at, last_activity FROM conversations WHERE id = ?1",
libsql::params![id],
)
.await
.unwrap();
let row = rows.next().await.unwrap().unwrap();
let started_at: String = row.get(0).unwrap();
let last_activity: String = row.get(1).unwrap();
// Must end with 'Z' (RFC 3339 UTC) and contain 'T' separator
assert!(
started_at.ends_with('Z') && started_at.contains('T'),
"started_at should be RFC 3339, got: {started_at}"
);
assert!(
last_activity.ends_with('Z') && last_activity.contains('T'),
"last_activity should be RFC 3339, got: {last_activity}"
);
// Must be parseable by the RFC 3339 parser directly (not just naive fallback)
use chrono::DateTime;
assert!(
DateTime::parse_from_rfc3339(&started_at).is_ok(),
"started_at not valid RFC 3339: {started_at}"
);
}
}
+55 -55
View File
@@ -26,7 +26,7 @@ pub const SCHEMA: &str = r#"
CREATE TABLE IF NOT EXISTS _migrations (
version INTEGER PRIMARY KEY,
name TEXT NOT NULL,
applied_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
applied_at TEXT NOT NULL DEFAULT (datetime('now'))
);
-- ==================== Conversations ====================
@@ -36,8 +36,8 @@ CREATE TABLE IF NOT EXISTS conversations (
channel TEXT NOT NULL,
user_id TEXT NOT NULL,
thread_id TEXT,
started_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
last_activity TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
started_at TEXT NOT NULL DEFAULT (datetime('now')),
last_activity TEXT NOT NULL DEFAULT (datetime('now')),
metadata TEXT NOT NULL DEFAULT '{}'
);
@@ -59,7 +59,7 @@ CREATE TABLE IF NOT EXISTS conversation_messages (
conversation_id TEXT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
role TEXT NOT NULL,
content TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_conversation_messages_conversation
@@ -91,7 +91,7 @@ CREATE TABLE IF NOT EXISTS agent_jobs (
failure_reason TEXT,
stuck_since TEXT,
repair_attempts INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
started_at TEXT,
completed_at TEXT
);
@@ -116,7 +116,7 @@ CREATE TABLE IF NOT EXISTS job_actions (
duration_ms INTEGER,
success INTEGER NOT NULL,
error_message TEXT,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE(job_id, sequence_num)
);
@@ -137,8 +137,8 @@ CREATE TABLE IF NOT EXISTS dynamic_tools (
failure_count INTEGER NOT NULL DEFAULT 0,
last_error TEXT,
status TEXT NOT NULL DEFAULT 'active',
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_dynamic_tools_status ON dynamic_tools(status);
@@ -156,7 +156,7 @@ CREATE TABLE IF NOT EXISTS llm_calls (
output_tokens INTEGER NOT NULL,
cost TEXT NOT NULL,
purpose TEXT,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_llm_calls_job ON llm_calls(job_id);
@@ -176,7 +176,7 @@ CREATE TABLE IF NOT EXISTS estimation_snapshots (
actual_time_secs INTEGER,
estimated_value TEXT NOT NULL,
actual_value TEXT,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_estimation_category ON estimation_snapshots(category);
@@ -192,7 +192,7 @@ CREATE TABLE IF NOT EXISTS repair_attempts (
action_taken TEXT NOT NULL,
success INTEGER NOT NULL,
error_message TEXT,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_repair_attempts_target ON repair_attempts(target_type, target_id);
@@ -206,8 +206,8 @@ CREATE TABLE IF NOT EXISTS memory_documents (
agent_id TEXT,
path TEXT NOT NULL,
content TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
metadata TEXT NOT NULL DEFAULT '{}',
UNIQUE (user_id, agent_id, path)
);
@@ -222,7 +222,7 @@ CREATE TRIGGER IF NOT EXISTS update_memory_documents_updated_at
FOR EACH ROW
WHEN NEW.updated_at = OLD.updated_at
BEGIN
UPDATE memory_documents SET updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') WHERE id = NEW.id;
UPDATE memory_documents SET updated_at = datetime('now') WHERE id = NEW.id;
END;
-- ==================== Workspace: Memory Chunks ====================
@@ -234,7 +234,7 @@ CREATE TABLE IF NOT EXISTS memory_chunks (
chunk_index INTEGER NOT NULL,
content TEXT NOT NULL,
embedding BLOB,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (document_id, chunk_index)
);
@@ -296,8 +296,8 @@ CREATE TABLE IF NOT EXISTS secrets (
expires_at TEXT,
last_used_at TEXT,
usage_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (user_id, name)
);
@@ -318,8 +318,8 @@ CREATE TABLE IF NOT EXISTS wasm_tools (
source_url TEXT,
trust_level TEXT NOT NULL DEFAULT 'user',
status TEXT NOT NULL DEFAULT 'active',
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (user_id, name, version)
);
@@ -340,8 +340,8 @@ CREATE TABLE IF NOT EXISTS wasm_channels (
binary_hash BLOB NOT NULL,
capabilities_json TEXT NOT NULL DEFAULT '{}',
status TEXT NOT NULL DEFAULT 'active',
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (user_id, name)
);
@@ -359,8 +359,8 @@ CREATE TABLE IF NOT EXISTS tool_capabilities (
max_response_body_bytes INTEGER NOT NULL DEFAULT 10485760,
workspace_read_prefixes TEXT NOT NULL DEFAULT '[]',
http_timeout_secs INTEGER NOT NULL DEFAULT 30,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (wasm_tool_id)
);
@@ -373,7 +373,7 @@ CREATE TABLE IF NOT EXISTS leak_detection_patterns (
severity TEXT NOT NULL DEFAULT 'high',
action TEXT NOT NULL DEFAULT 'block',
enabled INTEGER NOT NULL DEFAULT 1,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
-- ==================== Rate Limit State ====================
@@ -382,9 +382,9 @@ CREATE TABLE IF NOT EXISTS tool_rate_limit_state (
id TEXT PRIMARY KEY,
wasm_tool_id TEXT NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
user_id TEXT NOT NULL,
minute_window_start TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
minute_window_start TEXT NOT NULL DEFAULT (datetime('now')),
minute_count INTEGER NOT NULL DEFAULT 0,
hour_window_start TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
hour_window_start TEXT NOT NULL DEFAULT (datetime('now')),
hour_count INTEGER NOT NULL DEFAULT 0,
UNIQUE (wasm_tool_id, user_id)
);
@@ -400,7 +400,7 @@ CREATE TABLE IF NOT EXISTS secret_usage_log (
target_path TEXT,
success INTEGER NOT NULL,
error_message TEXT,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_secret_usage_user ON secret_usage_log(user_id);
@@ -415,7 +415,7 @@ CREATE TABLE IF NOT EXISTS leak_detection_events (
source TEXT NOT NULL,
action_taken TEXT NOT NULL,
context_preview TEXT,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
-- ==================== Tool Failures ====================
@@ -425,8 +425,8 @@ CREATE TABLE IF NOT EXISTS tool_failures (
tool_name TEXT NOT NULL UNIQUE,
error_message TEXT,
error_count INTEGER DEFAULT 1,
first_failure TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
last_failure TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
first_failure TEXT DEFAULT (datetime('now')),
last_failure TEXT DEFAULT (datetime('now')),
last_build_result TEXT,
repaired_at TEXT,
repair_attempts INTEGER DEFAULT 0
@@ -441,7 +441,7 @@ CREATE TABLE IF NOT EXISTS job_events (
job_id TEXT NOT NULL REFERENCES agent_jobs(id),
event_type TEXT NOT NULL,
data TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_job_events_job ON job_events(job_id, id);
@@ -471,8 +471,8 @@ CREATE TABLE IF NOT EXISTS routines (
next_fire_at TEXT,
run_count INTEGER NOT NULL DEFAULT 0,
consecutive_failures INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (user_id, name)
);
@@ -485,13 +485,13 @@ CREATE TABLE IF NOT EXISTS routine_runs (
routine_id TEXT NOT NULL REFERENCES routines(id) ON DELETE CASCADE,
trigger_type TEXT NOT NULL,
trigger_detail TEXT,
started_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
started_at TEXT NOT NULL DEFAULT (datetime('now')),
completed_at TEXT,
status TEXT NOT NULL DEFAULT 'running',
result_summary TEXT,
tokens_used INTEGER,
job_id TEXT REFERENCES agent_jobs(id),
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_routine_runs_routine ON routine_runs(routine_id);
@@ -502,7 +502,7 @@ CREATE TABLE IF NOT EXISTS settings (
user_id TEXT NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
PRIMARY KEY (user_id, key)
);
@@ -558,24 +558,24 @@ CREATE INDEX IF NOT EXISTS idx_heartbeat_next_run ON heartbeat_state(next_run);
-- Pre-populate leak detection patterns (matches PostgreSQL V2 migration).
INSERT OR IGNORE INTO leak_detection_patterns (id, name, pattern, severity, action, enabled, created_at) VALUES
('550e8400-e29b-41d4-a716-446655440001', 'openai_api_key', 'sk-(?:proj-)?[a-zA-Z0-9]{20,}(?:T3BlbkFJ[a-zA-Z0-9_-]*)?', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440002', 'anthropic_api_key', 'sk-ant-api[a-zA-Z0-9_-]{90,}', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440003', 'aws_access_key', 'AKIA[0-9A-Z]{16}', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440004', 'aws_secret_key', '(?<![A-Za-z0-9/+=])[A-Za-z0-9/+=]{40}(?![A-Za-z0-9/+=])', 'high', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440005', 'github_token', 'gh[pousr]_[A-Za-z0-9_]{36,}', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440006', 'github_fine_grained_pat', 'github_pat_[a-zA-Z0-9]{22}_[a-zA-Z0-9]{59}', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440007', 'stripe_api_key', 'sk_(?:live|test)_[a-zA-Z0-9]{24,}', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440008', 'nearai_session', 'sess_[a-zA-Z0-9]{32,}', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440009', 'bearer_token', 'Bearer\s+[a-zA-Z0-9_-]{20,}', 'high', 'redact', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-44665544000a', 'pem_private_key', '-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-44665544000b', 'ssh_private_key', '-----BEGIN\s+(?:OPENSSH|EC|DSA)\s+PRIVATE\s+KEY-----', 'critical', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-44665544000c', 'google_api_key', 'AIza[0-9A-Za-z_-]{35}', 'high', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-44665544000d', 'slack_token', 'xox[baprs]-[0-9a-zA-Z-]{10,}', 'high', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-44665544000e', 'discord_token', '[MN][A-Za-z\d]{23,}\.[\w-]{6}\.[\w-]{27}', 'high', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-44665544000f', 'twilio_api_key', 'SK[a-fA-F0-9]{32}', 'high', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440010', 'sendgrid_api_key', 'SG\.[a-zA-Z0-9_-]{22}\.[a-zA-Z0-9_-]{43}', 'high', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440011', 'mailchimp_api_key', '[a-f0-9]{32}-us[0-9]{1,2}', 'medium', 'block', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
('550e8400-e29b-41d4-a716-446655440012', 'high_entropy_hex', '(?<![a-fA-F0-9])[a-fA-F0-9]{64}(?![a-fA-F0-9])', 'medium', 'warn', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'));
('550e8400-e29b-41d4-a716-446655440001', 'openai_api_key', 'sk-(?:proj-)?[a-zA-Z0-9]{20,}(?:T3BlbkFJ[a-zA-Z0-9_-]*)?', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440002', 'anthropic_api_key', 'sk-ant-api[a-zA-Z0-9_-]{90,}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440003', 'aws_access_key', 'AKIA[0-9A-Z]{16}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440004', 'aws_secret_key', '(?<![A-Za-z0-9/+=])[A-Za-z0-9/+=]{40}(?![A-Za-z0-9/+=])', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440005', 'github_token', 'gh[pousr]_[A-Za-z0-9_]{36,}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440006', 'github_fine_grained_pat', 'github_pat_[a-zA-Z0-9]{22}_[a-zA-Z0-9]{59}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440007', 'stripe_api_key', 'sk_(?:live|test)_[a-zA-Z0-9]{24,}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440008', 'nearai_session', 'sess_[a-zA-Z0-9]{32,}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440009', 'bearer_token', 'Bearer\s+[a-zA-Z0-9_-]{20,}', 'high', 'redact', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000a', 'pem_private_key', '-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000b', 'ssh_private_key', '-----BEGIN\s+(?:OPENSSH|EC|DSA)\s+PRIVATE\s+KEY-----', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000c', 'google_api_key', 'AIza[0-9A-Za-z_-]{35}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000d', 'slack_token', 'xox[baprs]-[0-9a-zA-Z-]{10,}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000e', 'discord_token', '[MN][A-Za-z\d]{23,}\.[\w-]{6}\.[\w-]{27}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000f', 'twilio_api_key', 'SK[a-fA-F0-9]{32}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440010', 'sendgrid_api_key', 'SG\.[a-zA-Z0-9_-]{22}\.[a-zA-Z0-9_-]{43}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440011', 'mailchimp_api_key', '[a-f0-9]{32}-us[0-9]{1,2}', 'medium', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440012', 'high_entropy_hex', '(?<![a-fA-F0-9])[a-fA-F0-9]{64}(?![a-fA-F0-9])', 'medium', 'warn', 1, datetime('now'));
"#;
@@ -613,7 +613,7 @@ CREATE TABLE IF NOT EXISTS memory_chunks_new (
chunk_index INTEGER NOT NULL,
content TEXT NOT NULL,
embedding BLOB,
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (document_id, chunk_index)
);
+1 -1
View File
@@ -2508,7 +2508,7 @@ impl ExtensionManager {
&self.user_id,
)
} else {
McpClient::new_with_name(&server.name, &server.url)
McpClient::new_with_config(server.clone())
};
// Try to list and create tools
+55 -1
View File
@@ -3,6 +3,7 @@
//! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers.
//! Uses the Streamable HTTP transport with session management.
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
@@ -52,6 +53,9 @@ pub struct McpClient {
/// Server configuration (for token secret name lookup).
server_config: Option<McpServerConfig>,
/// Custom HTTP headers injected into every request.
custom_headers: HashMap<String, String>,
}
impl McpClient {
@@ -75,6 +79,7 @@ impl McpClient {
secrets: None,
user_id: "default".to_string(),
server_config: None,
custom_headers: HashMap::new(),
}
}
@@ -95,6 +100,28 @@ impl McpClient {
secrets: None,
user_id: "default".to_string(),
server_config: None,
custom_headers: HashMap::new(),
}
}
/// Create a new simple MCP client from a server configuration (no authentication).
///
/// Use this when you have an `McpServerConfig` with custom headers but no OAuth.
pub fn new_with_config(config: McpServerConfig) -> Self {
Self {
server_name: config.name.clone(),
server_url: config.url.clone(),
http_client: reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()
.expect("Failed to create HTTP client"),
next_id: AtomicU64::new(1),
tools_cache: RwLock::new(None),
session_manager: None,
secrets: None,
user_id: "default".to_string(),
custom_headers: config.headers.clone(),
server_config: Some(config),
}
}
@@ -119,6 +146,7 @@ impl McpClient {
session_manager: Some(session_manager),
secrets: Some(secrets),
user_id: user_id.into(),
custom_headers: config.headers.clone(),
server_config: Some(config),
}
}
@@ -178,7 +206,12 @@ impl McpClient {
.header("Content-Type", "application/json")
.json(&request);
// Add Authorization header if we have a token
// Add custom headers from config
for (key, value) in &self.custom_headers {
req_builder = req_builder.header(key, value);
}
// Add Authorization header if we have a token (overrides custom Authorization)
if let Some(token) = self.get_access_token().await? {
req_builder = req_builder.header("Authorization", format!("Bearer {}", token));
}
@@ -474,6 +507,7 @@ impl Clone for McpClient {
secrets: self.secrets.clone(),
user_id: self.user_id.clone(),
server_config: self.server_config.clone(),
custom_headers: self.custom_headers.clone(),
}
}
}
@@ -692,6 +726,26 @@ mod tests {
assert_eq!(id3, 3);
}
#[test]
fn test_custom_headers_from_config() {
use std::collections::HashMap;
let mut headers = HashMap::new();
headers.insert("X-API-Key".to_string(), "secret".to_string());
headers.insert("X-Custom".to_string(), "value".to_string());
let config = McpServerConfig::new("test", "http://localhost:8080").with_headers(headers);
let client = McpClient::new_with_config(config);
assert_eq!(client.custom_headers.len(), 2);
assert_eq!(client.custom_headers.get("X-API-Key").unwrap(), "secret");
}
#[test]
fn test_new_has_no_custom_headers() {
let client = McpClient::new("http://localhost:8080");
assert!(client.custom_headers.is_empty());
}
#[test]
fn test_mcp_tool_requires_approval_destructive() {
use crate::tools::mcp::protocol::{McpTool, McpToolAnnotations};
+89
View File
@@ -32,6 +32,13 @@ pub struct McpServerConfig {
/// Optional description for the server.
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
/// Custom HTTP headers to send with every request to this server.
///
/// Useful for MCP servers that require non-OAuth authentication
/// (e.g., API keys via `X-API-Key` header).
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub headers: HashMap<String, String>,
}
fn default_true() -> bool {
@@ -47,9 +54,16 @@ impl McpServerConfig {
oauth: None,
enabled: true,
description: None,
headers: HashMap::new(),
}
}
/// Set custom HTTP headers for this server.
pub fn with_headers(mut self, headers: HashMap<String, String>) -> Self {
self.headers = headers;
self
}
/// Set OAuth configuration.
pub fn with_oauth(mut self, oauth: OAuthConfig) -> Self {
self.oauth = Some(oauth);
@@ -593,4 +607,79 @@ mod tests {
let config = McpServerConfig::new("bad", "http://mcp.example.com");
assert!(!config.requires_auth());
}
#[test]
fn test_custom_headers_default_empty() {
let config = McpServerConfig::new("test", "http://localhost:8080");
assert!(config.headers.is_empty());
}
#[test]
fn test_custom_headers_with_builder() {
let mut headers = HashMap::new();
headers.insert("X-API-Key".to_string(), "secret123".to_string());
headers.insert("X-Custom".to_string(), "value".to_string());
let config = McpServerConfig::new("browser-use", "https://mcp.browser-use.com")
.with_headers(headers.clone());
assert_eq!(config.headers.len(), 2);
assert_eq!(config.headers.get("X-API-Key").unwrap(), "secret123");
}
#[test]
fn test_custom_headers_serde_roundtrip() {
let mut headers = HashMap::new();
headers.insert("Authorization".to_string(), "Bearer tok_123".to_string());
let config =
McpServerConfig::new("test-serde", "http://localhost:3000").with_headers(headers);
let json = serde_json::to_string(&config).unwrap();
let deserialized: McpServerConfig = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.headers.len(), 1);
assert_eq!(
deserialized.headers.get("Authorization").unwrap(),
"Bearer tok_123"
);
}
#[test]
fn test_custom_headers_absent_in_json_defaults_empty() {
let json = serde_json::json!({
"name": "legacy",
"url": "http://localhost:8080"
});
let config: McpServerConfig = serde_json::from_value(json).unwrap();
assert!(config.headers.is_empty());
}
#[test]
fn test_custom_headers_skipped_when_empty_in_serialization() {
let config = McpServerConfig::new("minimal", "http://localhost:8080");
let json = serde_json::to_value(&config).unwrap();
// Empty headers map should not appear in serialized output
assert!(json.get("headers").is_none());
}
#[tokio::test]
async fn test_custom_headers_persist_to_disk() {
let dir = tempdir().unwrap();
let path = dir.path().join("mcp-headers-test.json");
let mut headers = HashMap::new();
headers.insert("X-API-Key".to_string(), "key123".to_string());
let mut config = McpServersFile::default();
config.upsert(
McpServerConfig::new("headered", "http://localhost:9090").with_headers(headers),
);
save_mcp_servers_to(&config, &path).await.unwrap();
let loaded = load_mcp_servers_from(&path).await.unwrap();
let server = loaded.get("headered").unwrap();
assert_eq!(server.headers.get("X-API-Key").unwrap(), "key123");
}
}