mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 15:40:18 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5b0d261398 |
+2
-209
@@ -25,7 +25,6 @@ use crate::agent::routine::{
|
||||
};
|
||||
use crate::channels::{IncomingMessage, OutgoingResponse};
|
||||
use crate::config::RoutineConfig;
|
||||
use crate::context::JobState;
|
||||
use crate::db::Database;
|
||||
use crate::error::RoutineError;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
|
||||
@@ -181,130 +180,6 @@ impl RoutineEngine {
|
||||
}
|
||||
}
|
||||
|
||||
/// Sync dispatched routine runs with their linked background job status.
|
||||
///
|
||||
/// Full-job routines are fire-and-forget: the routine run is created with
|
||||
/// `Running` status when the job is dispatched, but the run record is never
|
||||
/// updated when the background job completes or fails. This method checks
|
||||
/// all `Running` routine runs that have a linked job, queries the job's
|
||||
/// current state, and updates the routine run accordingly. It also sends
|
||||
/// failure/success notifications that would otherwise be lost.
|
||||
pub async fn sync_dispatched_runs(&self) {
|
||||
let runs = match self.store.list_dispatched_routine_runs().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to list dispatched routine runs: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
for run in runs {
|
||||
let Some(job_id) = run.job_id else {
|
||||
continue;
|
||||
};
|
||||
|
||||
// Check the linked job's current state
|
||||
let job = match self.store.get_job(job_id).await {
|
||||
Ok(Some(j)) => j,
|
||||
Ok(None) => {
|
||||
// Job was deleted — mark the routine run as failed
|
||||
tracing::warn!(
|
||||
run_id = %run.id,
|
||||
job_id = %job_id,
|
||||
"Linked job not found, marking routine run as failed"
|
||||
);
|
||||
self.complete_dispatched_run(
|
||||
&run,
|
||||
RunStatus::Failed,
|
||||
"Linked job not found (may have been deleted)",
|
||||
)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
run_id = %run.id,
|
||||
job_id = %job_id,
|
||||
"Failed to query linked job: {}", e
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Extract the reason from the most recent state transition
|
||||
let last_reason = job.transitions.last().and_then(|t| t.reason.clone());
|
||||
|
||||
// Map job state to routine run status
|
||||
let (new_status, summary) = match job.state {
|
||||
JobState::Completed | JobState::Submitted | JobState::Accepted => {
|
||||
let summary =
|
||||
last_reason.unwrap_or_else(|| "Job completed successfully".to_string());
|
||||
(RunStatus::Ok, summary)
|
||||
}
|
||||
JobState::Failed => {
|
||||
let summary = last_reason
|
||||
.unwrap_or_else(|| "Job failed (no error message recorded)".to_string());
|
||||
(RunStatus::Failed, summary)
|
||||
}
|
||||
JobState::Cancelled => (RunStatus::Failed, "Job was cancelled".to_string()),
|
||||
// Still in progress — skip
|
||||
JobState::Pending | JobState::InProgress | JobState::Stuck => continue,
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
run_id = %run.id,
|
||||
job_id = %job_id,
|
||||
status = %new_status,
|
||||
"Syncing dispatched routine run with completed job"
|
||||
);
|
||||
|
||||
self.complete_dispatched_run(&run, new_status, &summary)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Complete a dispatched routine run and send the appropriate notification.
|
||||
async fn complete_dispatched_run(&self, run: &RoutineRun, status: RunStatus, summary: &str) {
|
||||
if let Err(e) = self
|
||||
.store
|
||||
.complete_routine_run(run.id, status, Some(summary), None)
|
||||
.await
|
||||
{
|
||||
tracing::error!(
|
||||
run_id = %run.id,
|
||||
"Failed to update dispatched routine run: {}", e
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Look up the routine to get its notify config and name
|
||||
match self.store.get_routine(run.routine_id).await {
|
||||
Ok(Some(routine)) => {
|
||||
send_notification(
|
||||
&self.notify_tx,
|
||||
&routine.notify,
|
||||
&routine.name,
|
||||
status,
|
||||
Some(summary),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Ok(None) => {
|
||||
tracing::debug!(
|
||||
routine_id = %run.routine_id,
|
||||
"Routine not found for notification (may have been deleted)"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
routine_id = %run.routine_id,
|
||||
"Failed to look up routine for notification: {}", e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fire a routine manually (from tool call or CLI).
|
||||
///
|
||||
/// Bypasses cooldown checks (those only apply to cron/event triggers).
|
||||
@@ -659,10 +534,9 @@ async fn execute_full_job(
|
||||
);
|
||||
|
||||
let summary = format!(
|
||||
"Dispatched job {job_id} for full execution with tool access (max_iterations: {max_iterations}). \
|
||||
Status will be updated when the job completes."
|
||||
"Dispatched job {job_id} for full execution with tool access (max_iterations: {max_iterations})"
|
||||
);
|
||||
Ok((RunStatus::Running, Some(summary), None))
|
||||
Ok((RunStatus::Ok, Some(summary), None))
|
||||
}
|
||||
|
||||
/// Execute a lightweight routine (single LLM call).
|
||||
@@ -838,7 +712,6 @@ pub fn spawn_cron_ticker(
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
engine.check_cron_triggers().await;
|
||||
engine.sync_dispatched_runs().await;
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -883,84 +756,4 @@ mod tests {
|
||||
let _ = status.to_string();
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_running_status_does_not_notify() {
|
||||
// Running status should not trigger notifications (job still in progress)
|
||||
let config = NotifyConfig {
|
||||
on_success: true,
|
||||
on_failure: true,
|
||||
on_attention: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// RunStatus::Running maps to false in send_notification's match
|
||||
let should_notify = match RunStatus::Running {
|
||||
RunStatus::Ok => config.on_success,
|
||||
RunStatus::Attention => config.on_attention,
|
||||
RunStatus::Failed => config.on_failure,
|
||||
RunStatus::Running => false,
|
||||
};
|
||||
assert!(!should_notify);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_full_job_dispatch_returns_running_status() {
|
||||
// Verify the status text for Running is "running"
|
||||
assert_eq!(RunStatus::Running.to_string(), "running");
|
||||
}
|
||||
|
||||
/// Regression test for #697: full_job routines were immediately marked Ok
|
||||
/// on dispatch, so failures/completions were never synced back. The fix
|
||||
/// changed dispatch to return Running and added sync_dispatched_runs which
|
||||
/// maps terminal job states to routine run statuses.
|
||||
#[test]
|
||||
fn test_job_state_to_run_status_mapping() {
|
||||
use crate::context::JobState;
|
||||
|
||||
// Helper that replicates the mapping logic from sync_dispatched_runs
|
||||
let map_state = |state: JobState, reason: Option<&str>| -> Option<(RunStatus, String)> {
|
||||
let last_reason = reason.map(|s| s.to_string());
|
||||
match state {
|
||||
JobState::Completed | JobState::Submitted | JobState::Accepted => {
|
||||
let summary =
|
||||
last_reason.unwrap_or_else(|| "Job completed successfully".to_string());
|
||||
Some((RunStatus::Ok, summary))
|
||||
}
|
||||
JobState::Failed => {
|
||||
let summary = last_reason
|
||||
.unwrap_or_else(|| "Job failed (no error message recorded)".to_string());
|
||||
Some((RunStatus::Failed, summary))
|
||||
}
|
||||
JobState::Cancelled => Some((RunStatus::Failed, "Job was cancelled".to_string())),
|
||||
JobState::Pending | JobState::InProgress | JobState::Stuck => None,
|
||||
}
|
||||
};
|
||||
|
||||
// Terminal states produce a status update
|
||||
let (status, _) = map_state(JobState::Completed, None).unwrap();
|
||||
assert_eq!(status, RunStatus::Ok);
|
||||
|
||||
let (status, _) = map_state(JobState::Submitted, None).unwrap();
|
||||
assert_eq!(status, RunStatus::Ok);
|
||||
|
||||
let (status, _) = map_state(JobState::Accepted, None).unwrap();
|
||||
assert_eq!(status, RunStatus::Ok);
|
||||
|
||||
let (status, summary) = map_state(JobState::Failed, Some("OOM killed")).unwrap();
|
||||
assert_eq!(status, RunStatus::Failed);
|
||||
assert_eq!(summary, "OOM killed");
|
||||
|
||||
let (status, summary) = map_state(JobState::Failed, None).unwrap();
|
||||
assert_eq!(status, RunStatus::Failed);
|
||||
assert!(summary.contains("no error message"));
|
||||
|
||||
let (status, _) = map_state(JobState::Cancelled, None).unwrap();
|
||||
assert_eq!(status, RunStatus::Failed);
|
||||
|
||||
// In-progress states should NOT produce a status update (skip)
|
||||
assert!(map_state(JobState::Pending, None).is_none());
|
||||
assert!(map_state(JobState::InProgress, None).is_none());
|
||||
assert!(map_state(JobState::Stuck, None).is_none());
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -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
@@ -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
|
||||
|
||||
@@ -423,29 +423,4 @@ impl RoutineStore for LibSqlBackend {
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||
let conn = self.connect().await?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
&format!(
|
||||
"SELECT {} FROM routine_runs \
|
||||
WHERE status = 'running' AND job_id IS NOT NULL",
|
||||
ROUTINE_RUN_COLUMNS
|
||||
),
|
||||
params![],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?;
|
||||
|
||||
let mut runs = Vec::new();
|
||||
while let Some(row) = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(e.to_string()))?
|
||||
{
|
||||
runs.push(row_to_routine_run_libsql(&row)?);
|
||||
}
|
||||
Ok(runs)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -303,10 +303,6 @@ pub trait RoutineStore: Send + Sync {
|
||||
run_id: Uuid,
|
||||
job_id: Uuid,
|
||||
) -> Result<(), DatabaseError>;
|
||||
/// List routine runs that were dispatched as full_job (status = 'running'
|
||||
/// with a linked job_id). Used by the routine engine to sync completion
|
||||
/// status from the background job.
|
||||
async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -494,10 +494,6 @@ impl RoutineStore for PgBackend {
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store.link_routine_run_to_job(run_id, job_id).await
|
||||
}
|
||||
|
||||
async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||
self.store.list_dispatched_routine_runs().await
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== ToolFailureStore ====================
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -1295,17 +1295,6 @@ impl Store {
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
let rows = conn
|
||||
.query(
|
||||
"SELECT * FROM routine_runs WHERE status = 'running' AND job_id IS NOT NULL",
|
||||
&[],
|
||||
)
|
||||
.await?;
|
||||
rows.iter().map(row_to_routine_run).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
|
||||
+55
-1
@@ -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};
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user