mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
30
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
98e6e2601c | ||
|
|
b5c9c9f5d6 | ||
|
|
5b95d22218 | ||
|
|
7ad035773a | ||
|
|
6040e56984 | ||
|
|
9946149a25 | ||
|
|
b18376f9f1 | ||
|
|
cd617500a8 | ||
|
|
3e866a9c0b | ||
|
|
da1db8f0e2 | ||
|
|
e313efc680 | ||
|
|
c91c63f810 | ||
|
|
54981c32f4 | ||
|
|
c779360730 | ||
|
|
6f687aabd2 | ||
|
|
386ed298c1 | ||
|
|
28192a8f30 | ||
|
|
6861c57638 | ||
|
|
28e95378d2 | ||
|
|
9a8f8cebc3 | ||
|
|
8bcdf1608f | ||
|
|
8a5346f417 | ||
|
|
708755f34a | ||
|
|
ce46e75dec | ||
|
|
42b66b33b2 | ||
|
|
e7ca8bb435 | ||
|
|
2af70642de | ||
|
|
ff971d0f14 | ||
|
|
5c8ee16f81 | ||
|
|
065a7498d5 |
+4
-5
@@ -17,6 +17,8 @@ target/
|
||||
# Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
*.pyd
|
||||
|
||||
# Benchmark results (local runs, not committed)
|
||||
bench-results/
|
||||
@@ -34,8 +36,5 @@ trace_*.json
|
||||
.claude/settings.local.json
|
||||
.worktrees/
|
||||
|
||||
# Python cache
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
*.pyd
|
||||
# JetBrains IDE
|
||||
.idea
|
||||
|
||||
Generated
+1
-1
@@ -269,7 +269,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "whatsapp-channel"
|
||||
version = "0.1.0"
|
||||
version = "0.2.0"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
||||
@@ -47,6 +47,15 @@ impl Agent {
|
||||
thread_id: Uuid,
|
||||
initial_messages: Vec<ChatMessage>,
|
||||
) -> Result<AgenticLoopResult, Error> {
|
||||
if let Some(ext_mgr) = self.deps.extension_manager.as_ref()
|
||||
&& let Err(e) = ext_mgr.ensure_nearai_companion_active_if_ready().await
|
||||
{
|
||||
tracing::debug!(
|
||||
"Failed to auto-activate NEAR AI companion MCP before turn: {}",
|
||||
e
|
||||
);
|
||||
}
|
||||
|
||||
// Detect group chat from channel metadata (needed before loading system prompt)
|
||||
let is_group_chat = message
|
||||
.metadata
|
||||
|
||||
+20
-1
@@ -449,6 +449,8 @@ impl AppBuilder {
|
||||
|
||||
let mcp_session_manager = Arc::new(McpSessionManager::new());
|
||||
let mcp_process_manager = Arc::new(McpProcessManager::new());
|
||||
let companion_mcp_server =
|
||||
crate::tools::mcp::config::derive_nearai_companion_mcp_server(&self.config);
|
||||
|
||||
// Create WASM tool runtime eagerly so extensions installed after startup
|
||||
// (e.g. via the web UI) can still be activated. The tools directory is only
|
||||
@@ -526,6 +528,7 @@ impl AppBuilder {
|
||||
let mcp_sm = Arc::clone(&mcp_session_manager);
|
||||
let pm = Arc::clone(&mcp_process_manager);
|
||||
let owner_id = self.config.owner_id.clone();
|
||||
let companion_mcp_server = companion_mcp_server.clone();
|
||||
async move {
|
||||
let servers_result = if let Some(ref d) = db {
|
||||
load_mcp_servers_from_db(d.as_ref(), &owner_id).await
|
||||
@@ -533,7 +536,16 @@ impl AppBuilder {
|
||||
crate::tools::mcp::config::load_mcp_servers().await
|
||||
};
|
||||
match servers_result {
|
||||
Ok(servers) => {
|
||||
Ok(mut servers) => {
|
||||
if let Some(companion) = companion_mcp_server {
|
||||
let companion_name = companion.name.clone();
|
||||
if !servers.insert_if_absent(companion) {
|
||||
tracing::debug!(
|
||||
"Skipping derived MCP companion '{}': an existing config with that name is already present",
|
||||
companion_name
|
||||
);
|
||||
}
|
||||
}
|
||||
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
|
||||
if !enabled.is_empty() {
|
||||
tracing::debug!(
|
||||
@@ -545,6 +557,8 @@ impl AppBuilder {
|
||||
let mut join_set = tokio::task::JoinSet::new();
|
||||
for server in enabled {
|
||||
let mcp_sm = Arc::clone(&mcp_sm);
|
||||
let nearai_session = Arc::clone(&self.session);
|
||||
let nearai_api_key = self.config.llm.nearai.api_key.clone();
|
||||
let secrets = secrets_store.clone();
|
||||
let tools = Arc::clone(&tools);
|
||||
let pm = Arc::clone(&pm);
|
||||
@@ -556,6 +570,8 @@ impl AppBuilder {
|
||||
let client = match crate::tools::mcp::create_client_from_config(
|
||||
server,
|
||||
&mcp_sm,
|
||||
Some(nearai_session),
|
||||
nearai_api_key,
|
||||
&pm,
|
||||
secrets,
|
||||
&owner_id,
|
||||
@@ -712,6 +728,8 @@ impl AppBuilder {
|
||||
let manager = Arc::new(ExtensionManager::new(
|
||||
Arc::clone(&mcp_session_manager),
|
||||
Arc::clone(&mcp_process_manager),
|
||||
Some(Arc::clone(&self.session)),
|
||||
self.config.llm.nearai.api_key.clone(),
|
||||
ext_secrets,
|
||||
Arc::clone(tools),
|
||||
Some(Arc::clone(hooks)),
|
||||
@@ -721,6 +739,7 @@ impl AppBuilder {
|
||||
self.config.tunnel.public_url.clone(),
|
||||
self.config.owner_id.clone(),
|
||||
self.db.clone(),
|
||||
companion_mcp_server,
|
||||
catalog_entries.clone(),
|
||||
));
|
||||
tools.register_extension_tools(Arc::clone(&manager));
|
||||
|
||||
@@ -70,6 +70,7 @@ pub async fn extensions_list_handler(
|
||||
tools: ext.tools,
|
||||
needs_setup: ext.needs_setup,
|
||||
has_auth: ext.has_auth,
|
||||
derived: ext.derived,
|
||||
activation_status,
|
||||
activation_error: ext.activation_error,
|
||||
version: ext.version,
|
||||
|
||||
+440
-2
@@ -836,10 +836,10 @@ async fn oauth_callback_handler(
|
||||
|
||||
let result: Result<(), String> = async {
|
||||
let token_response = if let Some(proxy_url) = &exchange_proxy_url {
|
||||
let gateway_token = flow.gateway_token.as_deref().unwrap_or_default();
|
||||
let oauth_proxy_auth_token = flow.oauth_proxy_auth_token().unwrap_or_default();
|
||||
oauth_defaults::exchange_via_proxy(oauth_defaults::ProxyTokenExchangeRequest {
|
||||
proxy_url,
|
||||
gateway_token,
|
||||
gateway_token: oauth_proxy_auth_token,
|
||||
token_url: &flow.token_url,
|
||||
client_id: &flow.client_id,
|
||||
client_secret: flow.client_secret.as_deref(),
|
||||
@@ -2092,6 +2092,7 @@ async fn extensions_list_handler(
|
||||
tools: ext.tools,
|
||||
needs_setup: ext.needs_setup,
|
||||
has_auth: ext.has_auth,
|
||||
derived: ext.derived,
|
||||
activation_status,
|
||||
activation_error: ext.activation_error,
|
||||
version: ext.version,
|
||||
@@ -2948,6 +2949,7 @@ mod tests {
|
||||
tools: Vec::new(),
|
||||
needs_setup: true,
|
||||
has_auth: false,
|
||||
derived: false,
|
||||
installed: true,
|
||||
activation_error: None,
|
||||
version: None,
|
||||
@@ -2985,6 +2987,7 @@ mod tests {
|
||||
tools: Vec::new(),
|
||||
needs_setup: true,
|
||||
has_auth: false,
|
||||
derived: false,
|
||||
installed: true,
|
||||
activation_error: None,
|
||||
version: None,
|
||||
@@ -3057,6 +3060,160 @@ mod tests {
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct RecordedOauthProxyRequest {
|
||||
authorization: Option<String>,
|
||||
form: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockOauthProxyState {
|
||||
requests: Arc<tokio::sync::Mutex<Vec<RecordedOauthProxyRequest>>>,
|
||||
}
|
||||
|
||||
struct MockOauthProxyServer {
|
||||
addr: std::net::SocketAddr,
|
||||
requests: Arc<tokio::sync::Mutex<Vec<RecordedOauthProxyRequest>>>,
|
||||
shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
|
||||
server_task: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl MockOauthProxyServer {
|
||||
async fn start() -> Self {
|
||||
async fn exchange_handler(
|
||||
State(state): State<MockOauthProxyState>,
|
||||
headers: axum::http::HeaderMap,
|
||||
axum::Form(form): axum::Form<std::collections::HashMap<String, String>>,
|
||||
) -> Json<serde_json::Value> {
|
||||
state.requests.lock().await.push(RecordedOauthProxyRequest {
|
||||
authorization: headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string),
|
||||
form,
|
||||
});
|
||||
Json(serde_json::json!({
|
||||
"access_token": "proxy-access-token",
|
||||
"refresh_token": "proxy-refresh-token",
|
||||
"expires_in": 7200
|
||||
}))
|
||||
}
|
||||
|
||||
let requests = Arc::new(tokio::sync::Mutex::new(Vec::new()));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind mock oauth proxy");
|
||||
let addr = listener.local_addr().expect("mock oauth proxy addr");
|
||||
let app = Router::new()
|
||||
.route("/oauth/exchange", post(exchange_handler))
|
||||
.with_state(MockOauthProxyState {
|
||||
requests: Arc::clone(&requests),
|
||||
});
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
|
||||
let server_task = tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await;
|
||||
});
|
||||
|
||||
Self {
|
||||
addr,
|
||||
requests,
|
||||
shutdown_tx: Some(shutdown_tx),
|
||||
server_task: Some(server_task),
|
||||
}
|
||||
}
|
||||
|
||||
fn base_url(&self) -> String {
|
||||
format!("http://{}", self.addr)
|
||||
}
|
||||
|
||||
async fn requests(&self) -> Vec<RecordedOauthProxyRequest> {
|
||||
self.requests.lock().await.clone()
|
||||
}
|
||||
|
||||
async fn shutdown(mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(task) = self.server_task.take() {
|
||||
let _ = task.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MockOauthProxyServer {
|
||||
fn drop(&mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(task) = self.server_task.take() {
|
||||
task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct EnvVarGuard {
|
||||
key: &'static str,
|
||||
original: Option<String>,
|
||||
}
|
||||
|
||||
impl Drop for EnvVarGuard {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: Tests use lock_env() to serialize environment access.
|
||||
unsafe {
|
||||
if let Some(ref value) = self.original {
|
||||
std::env::set_var(self.key, value);
|
||||
} else {
|
||||
std::env::remove_var(self.key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_env_var(key: &'static str, value: Option<&str>) -> EnvVarGuard {
|
||||
let original = std::env::var(key).ok();
|
||||
// SAFETY: Tests use lock_env() to serialize environment access.
|
||||
unsafe {
|
||||
if let Some(value) = value {
|
||||
std::env::set_var(key, value);
|
||||
} else {
|
||||
std::env::remove_var(key);
|
||||
}
|
||||
}
|
||||
EnvVarGuard { key, original }
|
||||
}
|
||||
|
||||
fn fresh_pending_oauth_flow(
|
||||
secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync>,
|
||||
sse_manager: Option<Arc<SseManager>>,
|
||||
oauth_proxy_auth_token: Option<String>,
|
||||
) -> crate::cli::oauth_defaults::PendingOAuthFlow {
|
||||
crate::cli::oauth_defaults::PendingOAuthFlow {
|
||||
extension_name: "test_tool".to_string(),
|
||||
display_name: "Test Tool".to_string(),
|
||||
token_url: "https://example.com/token".to_string(),
|
||||
client_id: "client123".to_string(),
|
||||
client_secret: None,
|
||||
redirect_uri: "https://example.com/oauth/callback".to_string(),
|
||||
code_verifier: Some("test-code-verifier".to_string()),
|
||||
access_token_field: "access_token".to_string(),
|
||||
secret_name: "test_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
validation_endpoint: None,
|
||||
scopes: vec!["email".to_string()],
|
||||
user_id: "test".to_string(),
|
||||
secrets,
|
||||
sse_manager,
|
||||
gateway_token: oauth_proxy_auth_token,
|
||||
token_exchange_extra_params: std::collections::HashMap::new(),
|
||||
client_id_secret_name: None,
|
||||
created_at: std::time::Instant::now(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_extensions_setup_submit_returns_failure_when_not_activated() {
|
||||
use axum::body::Body;
|
||||
@@ -3714,6 +3871,284 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_callback_accepts_versioned_hosted_state_without_instance_name() {
|
||||
use axum::body::Body;
|
||||
use tower::ServiceExt;
|
||||
|
||||
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
|
||||
Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new(
|
||||
crate::secrets::SecretsCrypto::new(secrecy::SecretString::from(
|
||||
TEST_GATEWAY_CRYPTO_KEY.to_string(),
|
||||
))
|
||||
.expect("crypto"),
|
||||
)));
|
||||
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone());
|
||||
|
||||
let Some(created_at) = expired_flow_created_at() else {
|
||||
eprintln!(
|
||||
"Skipping versioned OAuth state without instance test: monotonic uptime below expiry window"
|
||||
);
|
||||
return;
|
||||
};
|
||||
let flow = crate::cli::oauth_defaults::PendingOAuthFlow {
|
||||
extension_name: "test_tool".to_string(),
|
||||
display_name: "Test Tool".to_string(),
|
||||
token_url: "https://example.com/token".to_string(),
|
||||
client_id: "client123".to_string(),
|
||||
client_secret: None,
|
||||
redirect_uri: "https://example.com/oauth/callback".to_string(),
|
||||
code_verifier: None,
|
||||
access_token_field: "access_token".to_string(),
|
||||
secret_name: "test_token".to_string(),
|
||||
provider: None,
|
||||
validation_endpoint: None,
|
||||
scopes: vec![],
|
||||
user_id: "test".to_string(),
|
||||
secrets,
|
||||
sse_manager: None,
|
||||
gateway_token: None,
|
||||
token_exchange_extra_params: std::collections::HashMap::new(),
|
||||
client_id_secret_name: None,
|
||||
created_at,
|
||||
};
|
||||
|
||||
ext_mgr
|
||||
.pending_oauth_flows()
|
||||
.write()
|
||||
.await
|
||||
.insert("test_nonce".to_string(), flow);
|
||||
|
||||
let state = test_gateway_state(Some(ext_mgr.clone()));
|
||||
let app = test_oauth_router(state);
|
||||
let versioned_state =
|
||||
crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", None);
|
||||
|
||||
let req = axum::http::Request::builder()
|
||||
.uri(format!(
|
||||
"/oauth/callback?code=fake_code&state={}",
|
||||
urlencoding::encode(&versioned_state)
|
||||
))
|
||||
.body(Body::empty())
|
||||
.expect("request");
|
||||
|
||||
let resp = ServiceExt::<axum::http::Request<Body>>::oneshot(app, req)
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
let body = axum::body::to_bytes(resp.into_body(), 1024 * 64)
|
||||
.await
|
||||
.expect("body");
|
||||
let html = String::from_utf8_lossy(&body);
|
||||
assert!(html.contains("Authorization Failed"));
|
||||
assert!(
|
||||
ext_mgr
|
||||
.pending_oauth_flows()
|
||||
.read()
|
||||
.await
|
||||
.get("test_nonce")
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn test_oauth_callback_happy_path_with_gateway_token_fallback() {
|
||||
use axum::body::Body;
|
||||
use tower::ServiceExt;
|
||||
|
||||
let proxy = MockOauthProxyServer::start().await;
|
||||
// Keep the process-wide env locked for the full callback so the handler
|
||||
// sees a stable proxy URL/token configuration throughout the test.
|
||||
let _env_guard = crate::config::helpers::lock_env();
|
||||
let _exchange_url_guard =
|
||||
set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", Some(&proxy.base_url()));
|
||||
let _proxy_auth_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
|
||||
let secrets = test_secrets_store();
|
||||
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(Arc::clone(&secrets));
|
||||
let sse_mgr = Arc::new(SseManager::new());
|
||||
let mut receiver = sse_mgr.sender().subscribe();
|
||||
let flow = fresh_pending_oauth_flow(
|
||||
Arc::clone(&secrets),
|
||||
Some(Arc::clone(&sse_mgr)),
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token(),
|
||||
);
|
||||
|
||||
ext_mgr
|
||||
.pending_oauth_flows()
|
||||
.write()
|
||||
.await
|
||||
.insert("test_nonce".to_string(), flow);
|
||||
|
||||
let state = test_gateway_state(Some(ext_mgr.clone()));
|
||||
let app = test_oauth_router(state);
|
||||
let versioned_state =
|
||||
crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", Some("myinstance"));
|
||||
|
||||
let req = axum::http::Request::builder()
|
||||
.uri(format!(
|
||||
"/oauth/callback?code=fake_code&state={}",
|
||||
urlencoding::encode(&versioned_state)
|
||||
))
|
||||
.body(Body::empty())
|
||||
.expect("request");
|
||||
|
||||
let resp = ServiceExt::<axum::http::Request<Body>>::oneshot(app, req)
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
let body = axum::body::to_bytes(resp.into_body(), 1024 * 64)
|
||||
.await
|
||||
.expect("body");
|
||||
let html = String::from_utf8_lossy(&body);
|
||||
assert!(html.contains("Test Tool Connected"));
|
||||
|
||||
let requests = proxy.requests().await;
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].authorization.as_deref(),
|
||||
Some("Bearer gateway-test-token")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code").map(String::as_str),
|
||||
Some("fake_code")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code_verifier").map(String::as_str),
|
||||
Some("test-code-verifier")
|
||||
);
|
||||
|
||||
let access_token = secrets
|
||||
.get_decrypted("test", "test_token")
|
||||
.await
|
||||
.expect("access token stored");
|
||||
assert_eq!(access_token.expose(), "proxy-access-token");
|
||||
|
||||
let refresh_token = secrets
|
||||
.get_decrypted("test", "test_token_refresh_token")
|
||||
.await
|
||||
.expect("refresh token stored");
|
||||
assert_eq!(refresh_token.expose(), "proxy-refresh-token");
|
||||
|
||||
match receiver.recv().await.expect("auth_completed event").event {
|
||||
crate::channels::web::types::AppEvent::AuthCompleted {
|
||||
extension_name,
|
||||
success,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(extension_name, "test_tool");
|
||||
assert!(success, "OAuth callback should broadcast success");
|
||||
}
|
||||
event => panic!("expected AuthCompleted event, got {event:?}"),
|
||||
}
|
||||
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn test_oauth_callback_happy_path_with_dedicated_proxy_auth_token() {
|
||||
use axum::body::Body;
|
||||
use tower::ServiceExt;
|
||||
|
||||
let proxy = MockOauthProxyServer::start().await;
|
||||
// Keep the process-wide env locked for the full callback so the handler
|
||||
// sees a stable proxy URL/token configuration throughout the test.
|
||||
let _env_guard = crate::config::helpers::lock_env();
|
||||
let _exchange_url_guard =
|
||||
set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", Some(&proxy.base_url()));
|
||||
let _proxy_auth_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_PROXY_AUTH_TOKEN",
|
||||
Some("shared-oauth-proxy-secret"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
|
||||
|
||||
let secrets = test_secrets_store();
|
||||
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(Arc::clone(&secrets));
|
||||
let sse_mgr = Arc::new(SseManager::new());
|
||||
let mut receiver = sse_mgr.sender().subscribe();
|
||||
let flow = fresh_pending_oauth_flow(
|
||||
Arc::clone(&secrets),
|
||||
Some(Arc::clone(&sse_mgr)),
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token(),
|
||||
);
|
||||
|
||||
ext_mgr
|
||||
.pending_oauth_flows()
|
||||
.write()
|
||||
.await
|
||||
.insert("test_nonce".to_string(), flow);
|
||||
|
||||
let state = test_gateway_state(Some(ext_mgr.clone()));
|
||||
let app = test_oauth_router(state);
|
||||
let versioned_state =
|
||||
crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", None);
|
||||
|
||||
let req = axum::http::Request::builder()
|
||||
.uri(format!(
|
||||
"/oauth/callback?code=fake_code&state={}",
|
||||
urlencoding::encode(&versioned_state)
|
||||
))
|
||||
.body(Body::empty())
|
||||
.expect("request");
|
||||
|
||||
let resp = ServiceExt::<axum::http::Request<Body>>::oneshot(app, req)
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
let body = axum::body::to_bytes(resp.into_body(), 1024 * 64)
|
||||
.await
|
||||
.expect("body");
|
||||
let html = String::from_utf8_lossy(&body);
|
||||
assert!(html.contains("Test Tool Connected"));
|
||||
|
||||
let requests = proxy.requests().await;
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].authorization.as_deref(),
|
||||
Some("Bearer shared-oauth-proxy-secret")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code").map(String::as_str),
|
||||
Some("fake_code")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code_verifier").map(String::as_str),
|
||||
Some("test-code-verifier")
|
||||
);
|
||||
|
||||
let access_token = secrets
|
||||
.get_decrypted("test", "test_token")
|
||||
.await
|
||||
.expect("access token stored");
|
||||
assert_eq!(access_token.expose(), "proxy-access-token");
|
||||
|
||||
let refresh_token = secrets
|
||||
.get_decrypted("test", "test_token_refresh_token")
|
||||
.await
|
||||
.expect("refresh token stored");
|
||||
assert_eq!(refresh_token.expose(), "proxy-refresh-token");
|
||||
|
||||
match receiver.recv().await.expect("auth_completed event").event {
|
||||
crate::channels::web::types::AppEvent::AuthCompleted {
|
||||
extension_name,
|
||||
success,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(extension_name, "test_tool");
|
||||
assert!(success, "OAuth callback should broadcast success");
|
||||
}
|
||||
event => panic!("expected AuthCompleted event, got {event:?}"),
|
||||
}
|
||||
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
// --- Slack relay OAuth CSRF tests ---
|
||||
|
||||
fn test_relay_oauth_router(state: Arc<GatewayState>) -> Router {
|
||||
@@ -3745,6 +4180,8 @@ mod tests {
|
||||
let ext_mgr = Arc::new(ExtensionManager::new(
|
||||
mcp_sm,
|
||||
mcp_pm,
|
||||
None,
|
||||
None,
|
||||
secrets,
|
||||
tool_registry,
|
||||
None,
|
||||
@@ -3754,6 +4191,7 @@ mod tests {
|
||||
None,
|
||||
"test".to_string(),
|
||||
None,
|
||||
None,
|
||||
vec![],
|
||||
));
|
||||
(ext_mgr, wasm_tools_dir, wasm_channels_dir)
|
||||
|
||||
@@ -1160,7 +1160,7 @@ function addToolCard(name) {
|
||||
|
||||
const toolName = document.createElement('span');
|
||||
toolName.className = 'activity-tool-name';
|
||||
toolName.textContent = name;
|
||||
toolName.textContent = humanizeToolName(name);
|
||||
|
||||
const duration = document.createElement('span');
|
||||
duration.className = 'activity-tool-duration';
|
||||
@@ -1344,7 +1344,7 @@ function finalizeActivityGroup() {
|
||||
|
||||
function humanizeToolName(rawName) {
|
||||
if (!rawName) return '';
|
||||
return String(rawName)
|
||||
return stripDerivedCompanionToolPrefix(String(rawName))
|
||||
.replace(/[_-]+/g, ' ')
|
||||
.replace(/([a-z0-9])([A-Z])/g, '$1 $2')
|
||||
.replace(/^tool([a-zA-Z])/, 'tool $1')
|
||||
@@ -1352,6 +1352,12 @@ function humanizeToolName(rawName) {
|
||||
.trim();
|
||||
}
|
||||
|
||||
function stripDerivedCompanionToolPrefix(rawName) {
|
||||
if (!rawName) return '';
|
||||
const prefix = '_nearai_companion_mcp_';
|
||||
return rawName.startsWith(prefix) ? rawName.slice(prefix.length) : rawName;
|
||||
}
|
||||
|
||||
function shouldShowChannelConnectedMessage(extensionName, success) {
|
||||
if (!success || !extensionName) return false;
|
||||
return String(extensionName).toLowerCase().includes('telegram');
|
||||
@@ -1871,7 +1877,7 @@ function createToolCallsSummaryElement(toolCalls) {
|
||||
const icon = tc.has_error ? '\u2717' : '\u2713';
|
||||
const nameSpan = document.createElement('span');
|
||||
nameSpan.className = 'tool-call-name';
|
||||
nameSpan.textContent = icon + ' ' + tc.name;
|
||||
nameSpan.textContent = icon + ' ' + humanizeToolName(tc.name);
|
||||
item.appendChild(nameSpan);
|
||||
|
||||
if (tc.result_preview) {
|
||||
@@ -2906,7 +2912,12 @@ function renderExtensionCard(ext) {
|
||||
if (ext.tools && ext.tools.length > 0) {
|
||||
const tools = document.createElement('div');
|
||||
tools.className = 'ext-tools';
|
||||
tools.textContent = 'Tools: ' + ext.tools.join(', ');
|
||||
const toolNames = ext.tools.map((toolName) => (
|
||||
ext.derived && ext.kind === 'mcp_server'
|
||||
? stripDerivedCompanionToolPrefix(toolName)
|
||||
: toolName
|
||||
));
|
||||
tools.textContent = 'Tools: ' + toolNames.join(', ');
|
||||
card.appendChild(tools);
|
||||
}
|
||||
|
||||
@@ -2967,7 +2978,7 @@ function renderExtensionCard(ext) {
|
||||
// Skip when has_auth is true but needs_setup is false and not yet authenticated —
|
||||
// this means OAuth credentials resolve automatically (builtin/env) and the user
|
||||
// just needs to complete the OAuth flow, not fill in a config form.
|
||||
if (ext.needs_setup || (ext.has_auth && ext.authenticated)) {
|
||||
if (!ext.derived && (ext.needs_setup || (ext.has_auth && ext.authenticated))) {
|
||||
const configBtn = document.createElement('button');
|
||||
configBtn.className = 'btn-ext configure';
|
||||
configBtn.textContent = ext.authenticated ? I18n.t('ext.reconfigure') : I18n.t('ext.configure');
|
||||
@@ -2976,11 +2987,13 @@ function renderExtensionCard(ext) {
|
||||
}
|
||||
}
|
||||
|
||||
const removeBtn = document.createElement('button');
|
||||
removeBtn.className = 'btn-ext remove';
|
||||
removeBtn.textContent = I18n.t('ext.remove');
|
||||
removeBtn.addEventListener('click', () => removeExtension(ext.name));
|
||||
actions.appendChild(removeBtn);
|
||||
if (!ext.derived) {
|
||||
const removeBtn = document.createElement('button');
|
||||
removeBtn.className = 'btn-ext remove';
|
||||
removeBtn.textContent = I18n.t('ext.remove');
|
||||
removeBtn.addEventListener('click', () => removeExtension(ext.name));
|
||||
actions.appendChild(removeBtn);
|
||||
}
|
||||
|
||||
card.appendChild(actions);
|
||||
|
||||
|
||||
@@ -344,6 +344,9 @@ pub struct ExtensionInfo {
|
||||
/// Whether this extension has an auth configuration (OAuth or manual token).
|
||||
#[serde(default)]
|
||||
pub has_auth: bool,
|
||||
/// Whether this extension is derived from runtime/provider state.
|
||||
#[serde(default)]
|
||||
pub derived: bool,
|
||||
/// WASM channel activation status.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub activation_status: Option<ExtensionActivationStatus>,
|
||||
|
||||
+1
-2
@@ -127,8 +127,7 @@ async fn list_settings(
|
||||
}
|
||||
|
||||
let display_value = if value.len() > 60 {
|
||||
let end = crate::util::floor_char_boundary(&value, 57);
|
||||
format!("{}...", &value[..end])
|
||||
format!("{}...", &value[..57])
|
||||
} else {
|
||||
value
|
||||
};
|
||||
|
||||
+325
-41
@@ -8,7 +8,7 @@ use std::sync::Arc;
|
||||
|
||||
use clap::{Args, Subcommand};
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::config::{Config, LlmConfig};
|
||||
use crate::db::Database;
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::mcp::{
|
||||
@@ -173,6 +173,13 @@ async fn add_server(args: McpAddArgs) -> anyhow::Result<()> {
|
||||
description,
|
||||
} = args;
|
||||
|
||||
if config::is_nearai_companion_server_name(&name) {
|
||||
anyhow::bail!(
|
||||
"Server name '{}' is reserved for the NEAR AI companion MCP server",
|
||||
name
|
||||
);
|
||||
}
|
||||
|
||||
let transport_lower = transport.to_lowercase();
|
||||
|
||||
let mut config = match transport_lower.as_str() {
|
||||
@@ -244,7 +251,7 @@ async fn add_server(args: McpAddArgs) -> anyhow::Result<()> {
|
||||
|
||||
// Save (DB if available, else disk)
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
let mut servers = load_persisted_servers(db.as_deref()).await?;
|
||||
servers.upsert(config);
|
||||
save_servers(db.as_deref(), &servers).await?;
|
||||
|
||||
@@ -281,8 +288,15 @@ async fn add_server(args: McpAddArgs) -> anyhow::Result<()> {
|
||||
|
||||
/// Remove an MCP server.
|
||||
async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
if config::is_nearai_companion_server_name(&name) {
|
||||
anyhow::bail!(
|
||||
"Server '{}' is derived from the active NEAR AI provider and cannot be removed directly",
|
||||
name
|
||||
);
|
||||
}
|
||||
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
let mut servers = load_persisted_servers(db.as_deref()).await?;
|
||||
if !servers.remove(&name) {
|
||||
anyhow::bail!("Server '{}' not found", name);
|
||||
}
|
||||
@@ -298,7 +312,7 @@ async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
/// List configured MCP servers.
|
||||
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
let servers = load_servers_with_derived(db.as_deref()).await?;
|
||||
|
||||
if servers.servers.is_empty() {
|
||||
println!();
|
||||
@@ -404,12 +418,23 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
let servers = load_servers_with_derived(db.as_deref()).await?;
|
||||
let server = servers
|
||||
.get(&name)
|
||||
.cloned()
|
||||
.ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?;
|
||||
|
||||
if server.uses_runtime_auth_source() {
|
||||
println!();
|
||||
println!(
|
||||
" Server '{}' reuses your active NEAR AI authentication and does not support separate MCP OAuth.",
|
||||
name
|
||||
);
|
||||
println!(" Configure NEAR AI auth (API key or session login) instead.");
|
||||
println!();
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Initialize secrets store
|
||||
let secrets = get_secrets_store().await?;
|
||||
|
||||
@@ -477,7 +502,7 @@ async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
let servers = load_servers_with_derived(db.as_deref()).await?;
|
||||
let server = servers
|
||||
.get(&name)
|
||||
.cloned()
|
||||
@@ -488,35 +513,66 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
|
||||
// Create client
|
||||
let session_manager = Arc::new(McpSessionManager::new());
|
||||
|
||||
// Always check for stored tokens (from either pre-configured OAuth or DCR)
|
||||
let secrets = get_secrets_store().await?;
|
||||
let has_tokens = is_authenticated(&server, &secrets, &user_id).await;
|
||||
|
||||
let client = if has_tokens {
|
||||
// We have stored tokens, use authenticated client
|
||||
McpClient::new_authenticated(server.clone(), session_manager.clone(), secrets, user_id)
|
||||
} else if server.requires_auth() {
|
||||
// OAuth configured but no tokens - need to authenticate
|
||||
println!();
|
||||
println!(
|
||||
" ✗ Not authenticated. Run 'ironclaw mcp auth {}' first.",
|
||||
name
|
||||
);
|
||||
println!();
|
||||
return Ok(());
|
||||
} else {
|
||||
// Use the factory to dispatch on transport type (HTTP, stdio, unix)
|
||||
let (client, has_tokens) = if server.uses_runtime_auth_source() {
|
||||
let process_manager = Arc::new(McpProcessManager::new());
|
||||
create_client_from_config(
|
||||
server.clone(),
|
||||
&session_manager,
|
||||
&process_manager,
|
||||
None,
|
||||
"default",
|
||||
let llm = resolve_llm_for_cli(as_settings_store(db.as_deref())).await?;
|
||||
let nearai_session = crate::llm::create_session_manager(llm.session.clone()).await;
|
||||
(
|
||||
create_client_from_config(
|
||||
server.clone(),
|
||||
&session_manager,
|
||||
Some(nearai_session),
|
||||
llm.nearai.api_key.clone(),
|
||||
&process_manager,
|
||||
None,
|
||||
"default",
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?,
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
} else {
|
||||
// Only initialize the secrets store for non-runtime-auth servers that
|
||||
// can actually use persisted OAuth/DCR tokens.
|
||||
let secrets = get_secrets_store().await?;
|
||||
let has_tokens = is_authenticated(&server, &secrets, &user_id).await;
|
||||
|
||||
if has_tokens {
|
||||
(
|
||||
McpClient::new_authenticated(
|
||||
server.clone(),
|
||||
session_manager.clone(),
|
||||
secrets,
|
||||
user_id,
|
||||
),
|
||||
true,
|
||||
)
|
||||
} else if server.requires_auth() {
|
||||
println!();
|
||||
println!(
|
||||
" ✗ Not authenticated. Run 'ironclaw mcp auth {}' first.",
|
||||
name
|
||||
);
|
||||
println!();
|
||||
return Ok(());
|
||||
} else {
|
||||
// Use the factory to dispatch on transport type (HTTP, stdio, unix)
|
||||
let process_manager = Arc::new(McpProcessManager::new());
|
||||
(
|
||||
create_client_from_config(
|
||||
server.clone(),
|
||||
&session_manager,
|
||||
None,
|
||||
None,
|
||||
&process_manager,
|
||||
None,
|
||||
"default",
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?,
|
||||
false,
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
// Test connection
|
||||
@@ -581,8 +637,15 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
|
||||
/// Toggle server enabled/disabled state.
|
||||
async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> {
|
||||
if config::is_nearai_companion_server_name(&name) {
|
||||
anyhow::bail!(
|
||||
"Server '{}' is derived from the active NEAR AI provider and cannot be toggled directly",
|
||||
name
|
||||
);
|
||||
}
|
||||
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
let mut servers = load_persisted_servers(db.as_deref()).await?;
|
||||
|
||||
let server = servers
|
||||
.get_mut(&name)
|
||||
@@ -615,13 +678,30 @@ async fn connect_db() -> Option<Arc<dyn Database>> {
|
||||
crate::db::connect_from_config(&config.database).await.ok()
|
||||
}
|
||||
|
||||
/// Load MCP servers (DB if available, else disk).
|
||||
async fn load_servers(db: Option<&dyn Database>) -> Result<McpServersFile, config::ConfigError> {
|
||||
if let Some(db) = db {
|
||||
config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await
|
||||
/// Load only persisted MCP servers (DB if available, else disk).
|
||||
async fn load_persisted_servers(
|
||||
db: Option<&dyn Database>,
|
||||
) -> Result<McpServersFile, config::ConfigError> {
|
||||
Ok(if let Some(db) = db {
|
||||
config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await?
|
||||
} else {
|
||||
config::load_mcp_servers().await
|
||||
config::load_mcp_servers().await?
|
||||
})
|
||||
}
|
||||
|
||||
/// Load MCP servers plus any derived runtime companions.
|
||||
async fn load_servers_with_derived(
|
||||
db: Option<&dyn Database>,
|
||||
) -> Result<McpServersFile, config::ConfigError> {
|
||||
let mut servers = load_persisted_servers(db).await?;
|
||||
|
||||
if let Ok(llm) = resolve_llm_for_cli(as_settings_store(db)).await
|
||||
&& let Some(companion) = config::derive_nearai_companion_mcp_server_from_llm(&llm)
|
||||
{
|
||||
servers.insert_if_absent(companion);
|
||||
}
|
||||
|
||||
Ok(servers)
|
||||
}
|
||||
|
||||
/// Save MCP servers (DB if available, else disk).
|
||||
@@ -629,10 +709,15 @@ async fn save_servers(
|
||||
db: Option<&dyn Database>,
|
||||
servers: &McpServersFile,
|
||||
) -> Result<(), config::ConfigError> {
|
||||
let mut persisted = servers.clone();
|
||||
persisted
|
||||
.servers
|
||||
.retain(|server| !config::is_nearai_companion_server_name(&server.name));
|
||||
|
||||
if let Some(db) = db {
|
||||
config::save_mcp_servers_to_db(db, DEFAULT_USER_ID, servers).await
|
||||
config::save_mcp_servers_to_db(db, DEFAULT_USER_ID, &persisted).await
|
||||
} else {
|
||||
config::save_mcp_servers(servers).await
|
||||
config::save_mcp_servers(&persisted).await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -641,10 +726,84 @@ async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Syn
|
||||
crate::cli::init_secrets_store().await
|
||||
}
|
||||
|
||||
fn as_settings_store(db: Option<&dyn Database>) -> Option<&(dyn crate::db::SettingsStore + Sync)> {
|
||||
db.map(|db| db as &(dyn crate::db::SettingsStore + Sync))
|
||||
}
|
||||
|
||||
async fn resolve_llm_for_cli(
|
||||
store: Option<&(dyn crate::db::SettingsStore + Sync)>,
|
||||
) -> Result<LlmConfig, crate::error::ConfigError> {
|
||||
resolve_llm_for_cli_with_toml(store, None).await
|
||||
}
|
||||
|
||||
async fn resolve_llm_for_cli_with_toml(
|
||||
store: Option<&(dyn crate::db::SettingsStore + Sync)>,
|
||||
toml_path: Option<&std::path::Path>,
|
||||
) -> Result<LlmConfig, crate::error::ConfigError> {
|
||||
if let Some(store) = store {
|
||||
let _ = dotenvy::dotenv();
|
||||
crate::bootstrap::load_ironclaw_env();
|
||||
|
||||
let mut settings = match store.get_all_settings(DEFAULT_USER_ID).await {
|
||||
Ok(map) => crate::settings::Settings::from_db_map(&map),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to load CLI settings from DB, falling back to defaults before env/TOML resolution: {}",
|
||||
e
|
||||
);
|
||||
crate::settings::Settings::default()
|
||||
}
|
||||
};
|
||||
|
||||
apply_cli_toml_overlay(&mut settings, toml_path)?;
|
||||
return LlmConfig::resolve(&settings);
|
||||
}
|
||||
|
||||
let settings = crate::config::load_bootstrap_settings(toml_path)?;
|
||||
LlmConfig::resolve(&settings)
|
||||
}
|
||||
|
||||
fn apply_cli_toml_overlay(
|
||||
settings: &mut crate::settings::Settings,
|
||||
explicit_path: Option<&std::path::Path>,
|
||||
) -> Result<(), crate::error::ConfigError> {
|
||||
let path = explicit_path
|
||||
.map(std::path::PathBuf::from)
|
||||
.unwrap_or_else(crate::settings::Settings::default_toml_path);
|
||||
|
||||
match crate::settings::Settings::load_toml(&path) {
|
||||
Ok(Some(toml_settings)) => {
|
||||
settings.merge_from(&toml_settings);
|
||||
}
|
||||
Ok(None) => {
|
||||
if explicit_path.is_some() {
|
||||
return Err(crate::error::ConfigError::ParseError(format!(
|
||||
"Config file not found: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(crate::error::ConfigError::ParseError(e));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::error::DatabaseError;
|
||||
use crate::history::SettingRow;
|
||||
#[cfg(feature = "libsql")]
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
#[test]
|
||||
fn test_mcp_command_parsing() {
|
||||
// Just verify the command structure is valid
|
||||
@@ -701,4 +860,129 @@ mod tests {
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("invalid env var format"));
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn test_resolve_llm_for_cli_uses_db_backed_selected_model() {
|
||||
struct MockSettingsStore {
|
||||
settings: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl crate::db::SettingsStore for MockSettingsStore {
|
||||
async fn get_setting(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||
Ok(self.settings.get(key).cloned())
|
||||
}
|
||||
|
||||
async fn get_setting_full(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_key: &str,
|
||||
) -> Result<Option<SettingRow>, DatabaseError> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn set_setting(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_key: &str,
|
||||
_value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
Err(DatabaseError::Query("unused in test".to_string()))
|
||||
}
|
||||
|
||||
async fn delete_setting(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_key: &str,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
Err(DatabaseError::Query("unused in test".to_string()))
|
||||
}
|
||||
|
||||
async fn list_settings(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
) -> Result<Vec<SettingRow>, DatabaseError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn get_all_settings(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
) -> Result<HashMap<String, serde_json::Value>, DatabaseError> {
|
||||
Ok(self.settings.clone())
|
||||
}
|
||||
|
||||
async fn set_all_settings(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_settings: &HashMap<String, serde_json::Value>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
Err(DatabaseError::Query("unused in test".to_string()))
|
||||
}
|
||||
|
||||
async fn has_settings(&self, _user_id: &str) -> Result<bool, DatabaseError> {
|
||||
Ok(!self.settings.is_empty())
|
||||
}
|
||||
}
|
||||
|
||||
struct EnvGuard(&'static str, Option<String>);
|
||||
|
||||
impl Drop for EnvGuard {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: Protected by ENV_MUTEX for the duration of the test.
|
||||
unsafe {
|
||||
match &self.1 {
|
||||
Some(value) => std::env::set_var(self.0, value),
|
||||
None => std::env::remove_var(self.0),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _mutex = crate::config::helpers::ENV_MUTEX.lock().expect("env mutex");
|
||||
let prev_backend = std::env::var("LLM_BACKEND").ok();
|
||||
let prev_base_url = std::env::var("NEARAI_BASE_URL").ok();
|
||||
let prev_auth_url = std::env::var("NEARAI_AUTH_URL").ok();
|
||||
let prev_model = std::env::var("NEARAI_MODEL").ok();
|
||||
|
||||
// SAFETY: Protected by ENV_MUTEX for the duration of the test.
|
||||
unsafe {
|
||||
std::env::set_var("LLM_BACKEND", "");
|
||||
std::env::set_var("NEARAI_BASE_URL", "http://127.0.0.1:11434/v1");
|
||||
std::env::set_var("NEARAI_AUTH_URL", "http://127.0.0.1:11435");
|
||||
std::env::set_var("NEARAI_MODEL", "");
|
||||
}
|
||||
|
||||
let _backend_guard = EnvGuard("LLM_BACKEND", prev_backend);
|
||||
let _base_url_guard = EnvGuard("NEARAI_BASE_URL", prev_base_url);
|
||||
let _auth_url_guard = EnvGuard("NEARAI_AUTH_URL", prev_auth_url);
|
||||
let _model_guard = EnvGuard("NEARAI_MODEL", prev_model);
|
||||
|
||||
let empty_toml = NamedTempFile::new().expect("temp toml");
|
||||
let store = MockSettingsStore {
|
||||
settings: HashMap::from([
|
||||
("llm_backend".to_string(), serde_json::json!("nearai")),
|
||||
(
|
||||
"selected_model".to_string(),
|
||||
serde_json::json!("db-backed-nearai-model"),
|
||||
),
|
||||
]),
|
||||
};
|
||||
|
||||
let llm = resolve_llm_for_cli_with_toml(Some(&store), Some(empty_toml.path()))
|
||||
.await
|
||||
.expect("resolve llm");
|
||||
assert_eq!(llm.backend, "nearai");
|
||||
assert_eq!(llm.nearai.model, "db-backed-nearai-model");
|
||||
|
||||
let companion =
|
||||
config::derive_nearai_companion_mcp_server_from_llm(&llm).expect("derived companion");
|
||||
assert_eq!(companion.url, "http://127.0.0.1:11434/mcp");
|
||||
}
|
||||
}
|
||||
|
||||
+1
-15
@@ -256,8 +256,7 @@ fn truncate_content(s: &str, max_len: usize) -> String {
|
||||
if s.len() <= max_len {
|
||||
s.to_string()
|
||||
} else {
|
||||
let end = crate::util::floor_char_boundary(s, max_len);
|
||||
format!("{}...", &s[..end])
|
||||
format!("{}...", &s[..max_len])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -293,17 +292,4 @@ mod tests {
|
||||
assert_eq!(truncate_content("hello", 10), "hello");
|
||||
assert_eq!(truncate_content("hello world", 5), "hello...");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_content_multibyte_does_not_panic() {
|
||||
// \u{00e9} is precomposed 'é' (2 bytes in UTF-8)
|
||||
let s = "caf\u{00e9} au lait"; // "café au lait", é starts at byte 3
|
||||
let result = truncate_content(s, 4); // byte 4 is inside 2-byte é
|
||||
assert_eq!(result, "caf...");
|
||||
|
||||
// 4-byte emoji: slicing mid-emoji must not panic
|
||||
let emoji = "Hi \u{1F600} there"; // 😀 is 4 bytes, starts at byte 3
|
||||
let result = truncate_content(emoji, 4); // byte 4 is inside 😀
|
||||
assert_eq!(result, "Hi ...");
|
||||
}
|
||||
}
|
||||
|
||||
+184
-5
@@ -473,7 +473,8 @@ pub struct PendingOAuthFlow {
|
||||
pub secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||
/// SSE broadcast manager for notifying the web UI.
|
||||
pub sse_manager: Option<Arc<crate::channels::web::sse::SseManager>>,
|
||||
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
||||
/// OAuth proxy auth token for authenticating with the hosted token exchange proxy.
|
||||
/// Kept as `gateway_token` for public API compatibility.
|
||||
pub gateway_token: Option<String>,
|
||||
/// Additional form params for the token exchange request.
|
||||
/// Used for provider-specific requirements such as RFC 8707 `resource`.
|
||||
@@ -496,6 +497,12 @@ impl std::fmt::Debug for PendingOAuthFlow {
|
||||
}
|
||||
}
|
||||
|
||||
impl PendingOAuthFlow {
|
||||
pub fn oauth_proxy_auth_token(&self) -> Option<&str> {
|
||||
self.gateway_token.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// Thread-safe registry of pending OAuth flows, keyed by CSRF `state` parameter.
|
||||
pub type PendingOAuthRegistry = Arc<RwLock<HashMap<String, PendingOAuthFlow>>>;
|
||||
|
||||
@@ -529,6 +536,22 @@ pub fn exchange_proxy_url() -> Option<String> {
|
||||
.filter(|url| !url.is_empty())
|
||||
}
|
||||
|
||||
/// Returns the configured OAuth proxy auth token, if any.
|
||||
///
|
||||
/// New hosted infra can inject a dedicated shared proxy secret via
|
||||
/// `IRONCLAW_OAUTH_PROXY_AUTH_TOKEN`. Existing hosted instances continue to
|
||||
/// work by falling back to `GATEWAY_AUTH_TOKEN`.
|
||||
pub fn oauth_proxy_auth_token() -> Option<String> {
|
||||
fn normalized_env_value(key: &str) -> Option<String> {
|
||||
crate::config::helpers::env_or_override(key)
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
normalized_env_value("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN")
|
||||
.or_else(|| normalized_env_value("GATEWAY_AUTH_TOKEN"))
|
||||
}
|
||||
|
||||
/// Maximum age for pending OAuth flows (5 minutes, matching TCP listener timeout).
|
||||
pub const OAUTH_FLOW_EXPIRY: Duration = Duration::from_secs(300);
|
||||
|
||||
@@ -674,6 +697,8 @@ pub fn strip_instance_prefix(state: &str) -> &str {
|
||||
|
||||
pub struct ProxyTokenExchangeRequest<'a> {
|
||||
pub proxy_url: &'a str,
|
||||
/// OAuth proxy auth token.
|
||||
/// Kept as `gateway_token` for public API compatibility.
|
||||
pub gateway_token: &'a str,
|
||||
pub token_url: &'a str,
|
||||
pub client_id: &'a str,
|
||||
@@ -687,6 +712,8 @@ pub struct ProxyTokenExchangeRequest<'a> {
|
||||
|
||||
pub struct ProxyRefreshTokenRequest<'a> {
|
||||
pub proxy_url: &'a str,
|
||||
/// OAuth proxy auth token.
|
||||
/// Kept as `gateway_token` for public API compatibility.
|
||||
pub gateway_token: &'a str,
|
||||
pub token_url: &'a str,
|
||||
pub client_id: &'a str,
|
||||
@@ -729,7 +756,7 @@ fn oauth_token_response_from_json(
|
||||
|
||||
/// Exchange an OAuth authorization code via the platform's token exchange proxy.
|
||||
///
|
||||
/// Authenticated via the gateway auth token (Bearer header). The caller may
|
||||
/// Authenticated via an OAuth proxy auth token (Bearer header). The caller may
|
||||
/// either rely on proxy-side secret lookup or forward a `client_secret` when
|
||||
/// the provider requires it.
|
||||
///
|
||||
@@ -741,7 +768,7 @@ pub async fn exchange_via_proxy(
|
||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||
if request.gateway_token.is_empty() {
|
||||
return Err(OAuthCallbackError::Io(
|
||||
"Gateway auth token is required for proxy token exchange".to_string(),
|
||||
"OAuth proxy auth token is required for proxy token exchange".to_string(),
|
||||
));
|
||||
}
|
||||
let exchange_url = format!("{}/oauth/exchange", request.proxy_url.trim_end_matches('/'));
|
||||
@@ -796,7 +823,7 @@ pub async fn exchange_via_proxy(
|
||||
|
||||
/// Refresh an OAuth access token via the platform's token refresh proxy.
|
||||
///
|
||||
/// Authenticated via the gateway auth token (Bearer header). The caller may
|
||||
/// Authenticated via an OAuth proxy auth token (Bearer header). The caller may
|
||||
/// either rely on proxy-side secret lookup or forward a `client_secret` when
|
||||
/// the provider requires it.
|
||||
pub async fn refresh_token_via_proxy(
|
||||
@@ -804,7 +831,7 @@ pub async fn refresh_token_via_proxy(
|
||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||
if request.gateway_token.is_empty() {
|
||||
return Err(OAuthCallbackError::Io(
|
||||
"Gateway auth token is required for proxy token refresh".to_string(),
|
||||
"OAuth proxy auth token is required for proxy token refresh".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
@@ -1010,6 +1037,37 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct EnvVarGuard {
|
||||
key: &'static str,
|
||||
original: Option<String>,
|
||||
}
|
||||
|
||||
impl Drop for EnvVarGuard {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: Under ENV_MUTEX, no concurrent env access.
|
||||
unsafe {
|
||||
if let Some(ref value) = self.original {
|
||||
std::env::set_var(self.key, value);
|
||||
} else {
|
||||
std::env::remove_var(self.key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_env_var(key: &'static str, value: Option<&str>) -> EnvVarGuard {
|
||||
let original = std::env::var(key).ok();
|
||||
// SAFETY: Under ENV_MUTEX, no concurrent env access.
|
||||
unsafe {
|
||||
if let Some(value) = value {
|
||||
std::env::set_var(key, value);
|
||||
} else {
|
||||
std::env::remove_var(key);
|
||||
}
|
||||
}
|
||||
EnvVarGuard { key, original }
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hosted_proxy_client_secret_suppresses_builtin_secret() {
|
||||
let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds");
|
||||
@@ -1030,6 +1088,79 @@ mod tests {
|
||||
assert_eq!(result, client_secret);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_exchange_via_proxy_sends_auth_and_form() {
|
||||
let server = MockProxyServer::start().await;
|
||||
let mut extra_token_params = HashMap::new();
|
||||
extra_token_params.insert("resource".to_string(), "https://mcp.notion.com".to_string());
|
||||
|
||||
let response = super::exchange_via_proxy(super::ProxyTokenExchangeRequest {
|
||||
proxy_url: &server.base_url(),
|
||||
gateway_token: "shared-oauth-proxy-secret",
|
||||
code: "auth-code-123",
|
||||
redirect_uri: "https://oauth.example.com/oauth/callback",
|
||||
token_url: "https://oauth2.googleapis.com/token",
|
||||
client_id: TEST_OAUTH_CLIENT_ID,
|
||||
client_secret: Some(TEST_OAUTH_CLIENT_SECRET),
|
||||
access_token_field: "access_token",
|
||||
code_verifier: Some("code-verifier-123"),
|
||||
extra_token_params: &extra_token_params,
|
||||
})
|
||||
.await
|
||||
.expect("proxy exchange succeeds");
|
||||
|
||||
assert_eq!(response.access_token, "proxy-access-token");
|
||||
assert_eq!(
|
||||
response.refresh_token.as_deref(),
|
||||
Some("proxy-refresh-token")
|
||||
);
|
||||
assert_eq!(response.expires_in, Some(7200));
|
||||
|
||||
let requests = server.requests().await;
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].authorization.as_deref(),
|
||||
Some("Bearer shared-oauth-proxy-secret")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code").map(String::as_str),
|
||||
Some("auth-code-123")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("redirect_uri").map(String::as_str),
|
||||
Some("https://oauth.example.com/oauth/callback")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("token_url").map(String::as_str),
|
||||
Some("https://oauth2.googleapis.com/token")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("client_id").map(String::as_str),
|
||||
Some(TEST_OAUTH_CLIENT_ID)
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("client_secret").map(String::as_str),
|
||||
Some(TEST_OAUTH_CLIENT_SECRET)
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0]
|
||||
.form
|
||||
.get("access_token_field")
|
||||
.map(String::as_str),
|
||||
Some("access_token")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("code_verifier").map(String::as_str),
|
||||
Some("code-verifier-123")
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0].form.get("resource").map(String::as_str),
|
||||
Some("https://mcp.notion.com")
|
||||
);
|
||||
|
||||
server.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_token_via_proxy_sends_auth_and_form() {
|
||||
let server = MockProxyServer::start().await;
|
||||
@@ -1535,6 +1666,54 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_proxy_auth_token_prefers_dedicated_env() {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_PROXY_AUTH_TOKEN",
|
||||
Some("shared-proxy-secret"),
|
||||
);
|
||||
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||
|
||||
assert_eq!(
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||
Some("shared-proxy-secret")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_proxy_auth_token_falls_back_to_gateway_token() {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||
|
||||
assert_eq!(
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||
Some("gateway-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_proxy_auth_token_whitespace_dedicated_env_falls_back_to_gateway_token() {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", Some(" "));
|
||||
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token"));
|
||||
|
||||
assert_eq!(
|
||||
crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(),
|
||||
Some("gateway-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_proxy_auth_token_returns_none_when_unset() {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
|
||||
|
||||
assert_eq!(crate::cli::oauth_defaults::oauth_proxy_auth_token(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_instance_prefix_with_colon() {
|
||||
use crate::cli::oauth_defaults::strip_instance_prefix;
|
||||
|
||||
+479
-24
@@ -366,6 +366,8 @@ pub struct ExtensionManager {
|
||||
// MCP infrastructure
|
||||
mcp_session_manager: Arc<McpSessionManager>,
|
||||
mcp_process_manager: Arc<crate::tools::mcp::process::McpProcessManager>,
|
||||
nearai_session_manager: Option<Arc<crate::llm::SessionManager>>,
|
||||
nearai_api_key: Option<secrecy::SecretString>,
|
||||
/// Active MCP clients keyed by server name.
|
||||
mcp_clients: RwLock<HashMap<String, Arc<McpClient>>>,
|
||||
|
||||
@@ -389,6 +391,8 @@ pub struct ExtensionManager {
|
||||
user_id: String,
|
||||
/// Optional database store for DB-backed MCP config.
|
||||
store: Option<Arc<dyn crate::db::Database>>,
|
||||
/// Companion MCP server derived from the active provider config.
|
||||
companion_mcp_server: Option<McpServerConfig>,
|
||||
/// Names of WASM channels that were successfully loaded at startup.
|
||||
active_channel_names: RwLock<HashSet<String>>,
|
||||
/// Installed channel-relay extensions (no on-disk artifact, tracked in memory).
|
||||
@@ -403,9 +407,10 @@ pub struct ExtensionManager {
|
||||
/// when running in gateway mode, consumed by the web gateway's
|
||||
/// `/oauth/callback` handler.
|
||||
pending_oauth_flows: crate::cli::oauth_defaults::PendingOAuthRegistry,
|
||||
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
||||
/// Read once at construction from `GATEWAY_AUTH_TOKEN` env var.
|
||||
gateway_token: Option<String>,
|
||||
/// OAuth proxy auth token for authenticating with the hosted token exchange proxy.
|
||||
/// Resolved once at construction from `IRONCLAW_OAUTH_PROXY_AUTH_TOKEN`,
|
||||
/// then `GATEWAY_AUTH_TOKEN` as a backward-compatible fallback.
|
||||
oauth_proxy_auth_token: Option<String>,
|
||||
/// Relay config captured at startup. Used by `auth_channel_relay` and
|
||||
/// `activate_channel_relay` instead of re-reading env vars.
|
||||
relay_config: Option<crate::config::RelayConfig>,
|
||||
@@ -496,6 +501,8 @@ impl ExtensionManager {
|
||||
pub fn new(
|
||||
mcp_session_manager: Arc<McpSessionManager>,
|
||||
mcp_process_manager: Arc<crate::tools::mcp::process::McpProcessManager>,
|
||||
nearai_session_manager: Option<Arc<crate::llm::SessionManager>>,
|
||||
nearai_api_key: Option<secrecy::SecretString>,
|
||||
secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
hooks: Option<Arc<HookRegistry>>,
|
||||
@@ -505,6 +512,7 @@ impl ExtensionManager {
|
||||
tunnel_url: Option<String>,
|
||||
user_id: String,
|
||||
store: Option<Arc<dyn crate::db::Database>>,
|
||||
companion_mcp_server: Option<McpServerConfig>,
|
||||
catalog_entries: Vec<RegistryEntry>,
|
||||
) -> Self {
|
||||
let registry = if catalog_entries.is_empty() {
|
||||
@@ -517,6 +525,8 @@ impl ExtensionManager {
|
||||
discovery: OnlineDiscovery::new(),
|
||||
mcp_session_manager,
|
||||
mcp_process_manager,
|
||||
nearai_session_manager,
|
||||
nearai_api_key,
|
||||
mcp_clients: RwLock::new(HashMap::new()),
|
||||
wasm_tool_runtime,
|
||||
wasm_tools_dir,
|
||||
@@ -530,12 +540,13 @@ impl ExtensionManager {
|
||||
tunnel_url,
|
||||
user_id,
|
||||
store,
|
||||
companion_mcp_server,
|
||||
active_channel_names: RwLock::new(HashSet::new()),
|
||||
installed_relay_extensions: RwLock::new(HashSet::new()),
|
||||
activation_errors: RwLock::new(HashMap::new()),
|
||||
sse_manager: RwLock::new(None),
|
||||
pending_oauth_flows: crate::cli::oauth_defaults::new_pending_oauth_registry(),
|
||||
gateway_token: std::env::var("GATEWAY_AUTH_TOKEN").ok(),
|
||||
oauth_proxy_auth_token: crate::cli::oauth_defaults::oauth_proxy_auth_token(),
|
||||
relay_config: crate::config::RelayConfig::from_env(),
|
||||
relay_event_tx: Arc::new(tokio::sync::Mutex::new(None)),
|
||||
relay_signing_secret_cache: Arc::new(std::sync::Mutex::new(None)),
|
||||
@@ -1268,6 +1279,12 @@ impl ExtensionManager {
|
||||
tracing::info!(extension = %name, url = ?sanitized_url, kind = ?kind_hint, "Installing extension");
|
||||
Self::validate_extension_name(name)?;
|
||||
|
||||
if crate::tools::mcp::config::is_nearai_companion_server_name(name) {
|
||||
return Err(ExtensionError::Config(
|
||||
"This extension name is reserved for the NEAR AI companion MCP server".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// If we have a registry entry, use it (prefer kind_hint to resolve collisions)
|
||||
if let Some(entry) = self.registry.get_with_kind(name, kind_hint).await {
|
||||
return self.install_from_entry(&entry, user_id).await.map_err(|e| {
|
||||
@@ -1343,6 +1360,32 @@ impl ExtensionManager {
|
||||
}
|
||||
}
|
||||
|
||||
/// Activate the derived NEAR AI companion MCP server if auth is already
|
||||
/// available and the companion is not active yet.
|
||||
///
|
||||
/// Returns `Ok(true)` only when this call performed an activation.
|
||||
pub async fn ensure_nearai_companion_active_if_ready(&self) -> Result<bool, ExtensionError> {
|
||||
let Some(companion) = self.companion_mcp_server.as_ref() else {
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
let companion_name = companion.name.clone();
|
||||
|
||||
{
|
||||
let clients = self.mcp_clients.read().await;
|
||||
if clients.contains_key(&companion_name) {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
|
||||
if !self.is_runtime_authenticated(companion).await {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
self.activate(&companion_name, &self.user_id).await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// List extensions with their status.
|
||||
///
|
||||
/// When `include_available` is `true`, registry entries that are not yet
|
||||
@@ -1360,7 +1403,11 @@ impl ExtensionManager {
|
||||
match self.load_mcp_servers(user_id).await {
|
||||
Ok(servers) => {
|
||||
for server in &servers.servers {
|
||||
let authenticated = is_authenticated(server, &self.secrets, user_id).await;
|
||||
let authenticated = if server.uses_runtime_auth_source() {
|
||||
self.is_runtime_authenticated(server).await
|
||||
} else {
|
||||
is_authenticated(server, &self.secrets, user_id).await
|
||||
};
|
||||
let clients = self.mcp_clients.read().await;
|
||||
let active = clients.contains_key(&server.name);
|
||||
|
||||
@@ -1376,11 +1423,17 @@ impl ExtensionManager {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let display_name = self
|
||||
.registry
|
||||
.get_with_kind(&server.name, Some(ExtensionKind::McpServer))
|
||||
.await
|
||||
.map(|e| e.display_name);
|
||||
let display_name =
|
||||
if crate::tools::mcp::config::is_nearai_companion_server_name(
|
||||
&server.name,
|
||||
) {
|
||||
Some("NEAR AI Companion".to_string())
|
||||
} else {
|
||||
self.registry
|
||||
.get_with_kind(&server.name, Some(ExtensionKind::McpServer))
|
||||
.await
|
||||
.map(|e| e.display_name)
|
||||
};
|
||||
extensions.push(InstalledExtension {
|
||||
name: server.name.clone(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
@@ -1391,7 +1444,10 @@ impl ExtensionManager {
|
||||
active,
|
||||
tools,
|
||||
needs_setup: false,
|
||||
has_auth: false,
|
||||
has_auth: server.requires_auth(),
|
||||
derived: crate::tools::mcp::config::is_nearai_companion_server_name(
|
||||
&server.name,
|
||||
),
|
||||
installed: true,
|
||||
activation_error: None,
|
||||
version: None,
|
||||
@@ -1443,6 +1499,7 @@ impl ExtensionManager {
|
||||
tools: if active { vec![name] } else { Vec::new() },
|
||||
needs_setup: auth_state == ToolAuthState::NeedsSetup,
|
||||
has_auth: auth_state != ToolAuthState::NoAuth,
|
||||
derived: false,
|
||||
installed: true,
|
||||
activation_error: None,
|
||||
version,
|
||||
@@ -1499,6 +1556,7 @@ impl ExtensionManager {
|
||||
tools: Vec::new(),
|
||||
needs_setup: auth_state == ToolAuthState::NeedsSetup,
|
||||
has_auth: auth_state != ToolAuthState::NoAuth,
|
||||
derived: false,
|
||||
installed: true,
|
||||
activation_error,
|
||||
version,
|
||||
@@ -1537,6 +1595,7 @@ impl ExtensionManager {
|
||||
tools: Vec::new(),
|
||||
needs_setup: false,
|
||||
has_auth: true,
|
||||
derived: false,
|
||||
installed: true,
|
||||
activation_error,
|
||||
version: None,
|
||||
@@ -1571,6 +1630,7 @@ impl ExtensionManager {
|
||||
tools: Vec::new(),
|
||||
needs_setup: false,
|
||||
has_auth: false,
|
||||
derived: false,
|
||||
installed: false,
|
||||
activation_error: None,
|
||||
version: entry.version,
|
||||
@@ -1601,6 +1661,12 @@ impl ExtensionManager {
|
||||
|
||||
match kind {
|
||||
ExtensionKind::McpServer => {
|
||||
if crate::tools::mcp::config::is_nearai_companion_server_name(name) {
|
||||
return Err(ExtensionError::Config(
|
||||
"This MCP server is derived from the active NEAR AI provider and cannot be removed directly".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Unregister tools with this server's prefix
|
||||
let tool_names: Vec<String> = self
|
||||
.tool_registry
|
||||
@@ -2058,10 +2124,39 @@ impl ExtensionManager {
|
||||
user_id: &str,
|
||||
) -> Result<crate::tools::mcp::config::McpServersFile, crate::tools::mcp::config::ConfigError>
|
||||
{
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::load_mcp_servers_from_db(store.as_ref(), user_id).await
|
||||
let mut servers = if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::load_mcp_servers_from_db(store.as_ref(), user_id).await?
|
||||
} else {
|
||||
crate::tools::mcp::config::load_mcp_servers().await
|
||||
crate::tools::mcp::config::load_mcp_servers().await?
|
||||
};
|
||||
|
||||
if let Some(ref companion) = self.companion_mcp_server {
|
||||
servers.insert_if_absent(companion.clone());
|
||||
}
|
||||
|
||||
Ok(servers)
|
||||
}
|
||||
|
||||
async fn is_runtime_authenticated(&self, server: &McpServerConfig) -> bool {
|
||||
match server.auth_source {
|
||||
Some(crate::tools::mcp::config::McpAuthSource::NearAi) => {
|
||||
if self.nearai_api_key.is_some() {
|
||||
return true;
|
||||
}
|
||||
|
||||
if let Some(key) = crate::config::helpers::env_or_override("NEARAI_API_KEY")
|
||||
&& !key.trim().is_empty()
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
if let Some(ref session) = self.nearai_session_manager {
|
||||
return session.has_token().await;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2623,6 +2718,20 @@ impl ExtensionManager {
|
||||
.await
|
||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
||||
|
||||
if server.uses_runtime_auth_source() {
|
||||
if self.is_runtime_authenticated(&server).await {
|
||||
return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer));
|
||||
}
|
||||
|
||||
return Ok(AuthResult::needs_setup(
|
||||
name,
|
||||
ExtensionKind::McpServer,
|
||||
"This MCP server reuses your active NEAR AI authentication. Configure a NEAR AI API key or sign in to NEAR AI first, then try again."
|
||||
.to_string(),
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
||||
// Check if already authenticated
|
||||
if is_authenticated(&server, &self.secrets, user_id).await {
|
||||
return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer));
|
||||
@@ -2788,7 +2897,7 @@ impl ExtensionManager {
|
||||
user_id: user_id.to_string(),
|
||||
secrets: Arc::clone(&self.secrets),
|
||||
sse_manager: self.sse_manager.read().await.clone(),
|
||||
gateway_token: self.gateway_token.clone(),
|
||||
gateway_token: self.oauth_proxy_auth_token.clone(),
|
||||
token_exchange_extra_params,
|
||||
client_id_secret_name: if server.oauth.is_none() {
|
||||
Some(server.client_id_secret_name())
|
||||
@@ -3305,7 +3414,7 @@ impl ExtensionManager {
|
||||
user_id: user_id.to_string(),
|
||||
secrets: Arc::clone(&self.secrets),
|
||||
sse_manager: self.sse_manager.read().await.clone(),
|
||||
gateway_token: self.gateway_token.clone(),
|
||||
gateway_token: self.oauth_proxy_auth_token.clone(),
|
||||
token_exchange_extra_params: std::collections::HashMap::new(),
|
||||
client_id_secret_name: None,
|
||||
created_at: std::time::Instant::now(),
|
||||
@@ -3674,6 +3783,8 @@ impl ExtensionManager {
|
||||
let client = crate::tools::mcp::create_client_from_config(
|
||||
server.clone(),
|
||||
&self.mcp_session_manager,
|
||||
self.nearai_session_manager.clone(),
|
||||
self.nearai_api_key.clone(),
|
||||
&self.mcp_process_manager,
|
||||
Some(Arc::clone(&self.secrets)),
|
||||
user_id,
|
||||
@@ -5294,6 +5405,12 @@ impl ExtensionManager {
|
||||
.get_mcp_server(name, user_id)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
||||
if server.uses_runtime_auth_source() {
|
||||
return Err(ExtensionError::Other(format!(
|
||||
"Server '{}' reuses your active NEAR AI authentication and does not accept manually configured MCP tokens",
|
||||
name
|
||||
)));
|
||||
}
|
||||
let mut names = std::collections::HashSet::new();
|
||||
names.insert(server.token_secret_name());
|
||||
(names, Vec::new())
|
||||
@@ -6298,6 +6415,8 @@ mod tests {
|
||||
tools_dir: std::path::PathBuf,
|
||||
channels_dir: std::path::PathBuf,
|
||||
store: Option<Arc<dyn crate::db::Database>>,
|
||||
companion_mcp_server: Option<crate::tools::mcp::config::McpServerConfig>,
|
||||
nearai_session_manager: Option<Arc<crate::llm::SessionManager>>,
|
||||
) -> crate::extensions::manager::ExtensionManager {
|
||||
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
|
||||
use crate::tools::mcp::process::McpProcessManager;
|
||||
@@ -6316,15 +6435,18 @@ mod tests {
|
||||
crate::extensions::manager::ExtensionManager::new(
|
||||
mcp,
|
||||
Arc::new(McpProcessManager::new()),
|
||||
nearai_session_manager,
|
||||
None,
|
||||
secrets,
|
||||
tools,
|
||||
None, // hooks
|
||||
wasm_runtime,
|
||||
tools_dir,
|
||||
channels_dir,
|
||||
None, // tunnel_url
|
||||
"test".to_string(), // user_id
|
||||
None, // tunnel_url
|
||||
"test".to_string(),
|
||||
store,
|
||||
companion_mcp_server,
|
||||
vec![],
|
||||
)
|
||||
}
|
||||
@@ -6333,7 +6455,7 @@ mod tests {
|
||||
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
|
||||
tools_dir: std::path::PathBuf,
|
||||
) -> crate::extensions::manager::ExtensionManager {
|
||||
make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir, None)
|
||||
make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir, None, None, None)
|
||||
}
|
||||
|
||||
fn write_test_tool(
|
||||
@@ -6394,6 +6516,8 @@ mod tests {
|
||||
dir.path().join("tools"),
|
||||
dir.path().join("channels"),
|
||||
Some(Arc::clone(&store)),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
let field = crate::tools::wasm::ToolFieldSetupSchema {
|
||||
name: "provider".to_string(),
|
||||
@@ -6435,8 +6559,14 @@ mod tests {
|
||||
);
|
||||
let channels_dir = dir.path().join("channels");
|
||||
|
||||
let mgr =
|
||||
make_test_manager_with_dirs(None, tools_dir, channels_dir, Some(Arc::clone(&store)));
|
||||
let mgr = make_test_manager_with_dirs(
|
||||
None,
|
||||
tools_dir,
|
||||
channels_dir,
|
||||
Some(Arc::clone(&store)),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
let mut fields = std::collections::HashMap::new();
|
||||
fields.insert("llm_backend".to_string(), "openai".to_string());
|
||||
|
||||
@@ -6488,8 +6618,14 @@ mod tests {
|
||||
);
|
||||
let channels_dir = dir.path().join("channels");
|
||||
|
||||
let mgr =
|
||||
make_test_manager_with_dirs(None, tools_dir, channels_dir, Some(Arc::clone(&store)));
|
||||
let mgr = make_test_manager_with_dirs(
|
||||
None,
|
||||
tools_dir,
|
||||
channels_dir,
|
||||
Some(Arc::clone(&store)),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
let mut fields = std::collections::HashMap::new();
|
||||
fields.insert("session".to_string(), "overwrite".to_string());
|
||||
|
||||
@@ -6562,6 +6698,272 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_install_rejects_reserved_nearai_companion_name() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let manager = make_test_manager(None, dir.path().to_path_buf());
|
||||
|
||||
let err = manager
|
||||
.install(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
Some("https://mcp.example.com"),
|
||||
Some(ExtensionKind::McpServer),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.expect_err("reserved companion name should be rejected");
|
||||
|
||||
assert!(
|
||||
matches!(err, ExtensionError::Config(_)),
|
||||
"Expected config error, got: {err:?}"
|
||||
);
|
||||
assert!(
|
||||
err.to_string().contains("reserved"),
|
||||
"Expected reserved-name message, got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_ensure_nearai_companion_active_if_ready_skips_without_auth() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let companion = crate::tools::mcp::config::McpServerConfig::new(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
"https://private.near.ai/mcp",
|
||||
)
|
||||
.with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
|
||||
let manager = make_test_manager_with_dirs(
|
||||
None,
|
||||
dir.path().join("tools"),
|
||||
dir.path().join("channels"),
|
||||
None,
|
||||
Some(companion),
|
||||
None,
|
||||
);
|
||||
|
||||
let activated = manager
|
||||
.ensure_nearai_companion_active_if_ready()
|
||||
.await
|
||||
.expect("helper should not fail when auth is missing");
|
||||
|
||||
assert!(!activated, "companion should not activate without auth");
|
||||
assert!(
|
||||
!manager
|
||||
.mcp_clients
|
||||
.read()
|
||||
.await
|
||||
.contains_key(crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME)
|
||||
);
|
||||
}
|
||||
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn test_runtime_auth_detects_runtime_nearai_api_key_override() {
|
||||
struct EnvGuard(&'static str, Option<String>);
|
||||
|
||||
impl Drop for EnvGuard {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: Protected by ENV_MUTEX for the duration of the test.
|
||||
unsafe {
|
||||
match &self.1 {
|
||||
Some(value) => std::env::set_var(self.0, value),
|
||||
None => std::env::remove_var(self.0),
|
||||
}
|
||||
}
|
||||
crate::config::helpers::set_runtime_env(self.0, "");
|
||||
}
|
||||
}
|
||||
|
||||
let _mutex = crate::config::helpers::ENV_MUTEX.lock().expect("env mutex");
|
||||
let prev = std::env::var("NEARAI_API_KEY").ok();
|
||||
// SAFETY: Protected by ENV_MUTEX for the duration of the test.
|
||||
unsafe { std::env::remove_var("NEARAI_API_KEY") };
|
||||
let _env_guard = EnvGuard("NEARAI_API_KEY", prev);
|
||||
|
||||
crate::config::helpers::set_runtime_env("NEARAI_API_KEY", "runtime-overlay-key");
|
||||
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let companion = crate::tools::mcp::config::McpServerConfig::new(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
"https://private.near.ai/mcp",
|
||||
)
|
||||
.with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
|
||||
let manager = make_test_manager_with_dirs(
|
||||
None,
|
||||
dir.path().join("tools"),
|
||||
dir.path().join("channels"),
|
||||
None,
|
||||
Some(companion.clone()),
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(
|
||||
manager.is_runtime_authenticated(&companion).await,
|
||||
"runtime NEARAI_API_KEY override should count as authenticated"
|
||||
);
|
||||
}
|
||||
|
||||
async fn start_runtime_auth_mock_mcp_server() -> (String, tokio::task::JoinHandle<()>) {
|
||||
use axum::extract::State;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockState {
|
||||
auth_token: &'static str,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct JsonRpcRequest {
|
||||
id: Option<serde_json::Value>,
|
||||
method: String,
|
||||
}
|
||||
|
||||
async fn handle_mcp(
|
||||
State(state): State<Arc<MockState>>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<JsonRpcRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let auth = headers
|
||||
.get("authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("");
|
||||
|
||||
if auth != format!("Bearer {}", state.auth_token) {
|
||||
return (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
Json(serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"error": {"code": -32000, "message": "Unauthorized"}
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
if req.id.is_none() {
|
||||
return StatusCode::OK.into_response();
|
||||
}
|
||||
|
||||
let body = match req.method.as_str() {
|
||||
"initialize" => serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"serverInfo": {
|
||||
"name": "mock-mcp-server",
|
||||
"version": "1.0.0"
|
||||
},
|
||||
"capabilities": {
|
||||
"tools": {}
|
||||
}
|
||||
}
|
||||
}),
|
||||
"tools/list" => serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"result": {
|
||||
"tools": [{
|
||||
"name": "echo",
|
||||
"description": "Mock companion tool",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}
|
||||
}]
|
||||
}
|
||||
}),
|
||||
_ => serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"error": {
|
||||
"code": -32601,
|
||||
"message": format!("Method not found: {}", req.method)
|
||||
}
|
||||
}),
|
||||
};
|
||||
|
||||
Json(body).into_response()
|
||||
}
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind mock MCP server");
|
||||
let addr = listener.local_addr().expect("local addr");
|
||||
let base_url = format!("http://127.0.0.1:{}", addr.port());
|
||||
let app = Router::new()
|
||||
.route("/mcp", post(handle_mcp))
|
||||
.with_state(Arc::new(MockState {
|
||||
auth_token: "mock-access-token",
|
||||
}));
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.expect("serve mock MCP");
|
||||
});
|
||||
|
||||
(format!("{base_url}/mcp"), handle)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_ensure_nearai_companion_active_if_ready_activates_after_auth_becomes_available() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let (mcp_url, server_handle) = start_runtime_auth_mock_mcp_server().await;
|
||||
let session = Arc::new(crate::llm::SessionManager::new(
|
||||
crate::llm::SessionConfig::default(),
|
||||
));
|
||||
let companion = crate::tools::mcp::config::McpServerConfig::new(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
mcp_url,
|
||||
)
|
||||
.with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
|
||||
let manager = make_test_manager_with_dirs(
|
||||
None,
|
||||
dir.path().join("tools"),
|
||||
dir.path().join("channels"),
|
||||
None,
|
||||
Some(companion),
|
||||
Some(session.clone()),
|
||||
);
|
||||
|
||||
let first = manager
|
||||
.ensure_nearai_companion_active_if_ready()
|
||||
.await
|
||||
.expect("helper should skip cleanly before auth exists");
|
||||
assert!(!first, "companion should not activate before auth exists");
|
||||
|
||||
session
|
||||
.set_token(secrecy::SecretString::from("mock-access-token"))
|
||||
.await;
|
||||
|
||||
let second = manager
|
||||
.ensure_nearai_companion_active_if_ready()
|
||||
.await
|
||||
.expect("helper should activate once auth becomes available");
|
||||
|
||||
assert!(second, "companion should activate after auth appears");
|
||||
assert!(
|
||||
manager
|
||||
.mcp_clients
|
||||
.read()
|
||||
.await
|
||||
.contains_key(crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME)
|
||||
);
|
||||
assert!(
|
||||
manager.tool_registry.list().await.into_iter().any(|name| {
|
||||
name == format!(
|
||||
"{}_echo",
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME
|
||||
)
|
||||
}),
|
||||
"expected companion tool to be registered after delayed activation"
|
||||
);
|
||||
|
||||
server_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_capabilities_files_also_separate() {
|
||||
// capabilities.json files for tools and channels should also be separate.
|
||||
@@ -6684,6 +7086,8 @@ mod tests {
|
||||
ExtensionManager::new(
|
||||
Arc::new(McpSessionManager::new()),
|
||||
Arc::new(McpProcessManager::new()),
|
||||
None,
|
||||
None,
|
||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||
Arc::new(ToolRegistry::new()),
|
||||
None,
|
||||
@@ -6693,6 +7097,7 @@ mod tests {
|
||||
None,
|
||||
"test".to_string(),
|
||||
None,
|
||||
None,
|
||||
Vec::new(),
|
||||
)
|
||||
}
|
||||
@@ -6834,6 +7239,8 @@ mod tests {
|
||||
ExtensionManager::new(
|
||||
Arc::new(McpSessionManager::new()),
|
||||
Arc::new(McpProcessManager::new()),
|
||||
None,
|
||||
None,
|
||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||
Arc::new(ToolRegistry::new()),
|
||||
None,
|
||||
@@ -6843,6 +7250,7 @@ mod tests {
|
||||
None,
|
||||
"test".to_string(),
|
||||
Some(db),
|
||||
None,
|
||||
Vec::new(),
|
||||
)
|
||||
};
|
||||
@@ -7096,6 +7504,8 @@ mod tests {
|
||||
let manager = ExtensionManager::new(
|
||||
Arc::new(McpSessionManager::new()),
|
||||
Arc::new(McpProcessManager::new()),
|
||||
None,
|
||||
None,
|
||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||
Arc::new(ToolRegistry::new()),
|
||||
None,
|
||||
@@ -7105,6 +7515,7 @@ mod tests {
|
||||
None,
|
||||
"test".to_string(),
|
||||
Some(db.clone() as Arc<dyn crate::db::Database>),
|
||||
None,
|
||||
Vec::new(),
|
||||
);
|
||||
|
||||
@@ -7589,7 +8000,8 @@ mod tests {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let tools_dir = dir.path().join("tools");
|
||||
let channels_dir = dir.path().join("channels");
|
||||
let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone(), None);
|
||||
let mgr =
|
||||
make_test_manager_with_dirs(None, tools_dir, channels_dir.clone(), None, None, None);
|
||||
|
||||
let wasm_path = channels_dir.join("telegram.wasm");
|
||||
let cap_path = channels_dir.join("telegram.capabilities.json");
|
||||
@@ -7733,6 +8145,8 @@ mod tests {
|
||||
ExtensionManager::new(
|
||||
mcp,
|
||||
Arc::new(McpProcessManager::new()),
|
||||
None,
|
||||
None,
|
||||
secrets,
|
||||
tools,
|
||||
None,
|
||||
@@ -7742,6 +8156,7 @@ mod tests {
|
||||
tunnel_url,
|
||||
"test".to_string(),
|
||||
None,
|
||||
None,
|
||||
vec![],
|
||||
)
|
||||
}
|
||||
@@ -8024,6 +8439,46 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_configure_token_rejects_runtime_auth_companion() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let companion = crate::tools::mcp::config::McpServerConfig::new(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
"https://private.near.ai/mcp",
|
||||
)
|
||||
.with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
|
||||
let token_secret_name = companion.token_secret_name();
|
||||
let mgr = make_test_manager_with_dirs(
|
||||
None,
|
||||
dir.path().join("tools"),
|
||||
dir.path().join("channels"),
|
||||
None,
|
||||
Some(companion),
|
||||
None,
|
||||
);
|
||||
|
||||
let err = mgr
|
||||
.configure_token(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
"manual-token",
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.expect_err("runtime-auth companion should reject manual token configuration");
|
||||
|
||||
assert!(
|
||||
err.to_string().contains("active NEAR AI authentication"),
|
||||
"expected runtime-auth rejection message, got: {err}"
|
||||
);
|
||||
assert!(
|
||||
!mgr.secrets
|
||||
.exists("test", &token_secret_name)
|
||||
.await
|
||||
.unwrap_or(false),
|
||||
"configure_token must not persist a manual MCP token for the runtime-auth companion"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_auth_is_read_only_for_wasm_channel() {
|
||||
// Regression: auth() must be a pure status check — it must not store
|
||||
|
||||
@@ -506,6 +506,10 @@ pub struct InstalledExtension {
|
||||
/// Whether this extension has an auth configuration (OAuth or manual token).
|
||||
#[serde(default)]
|
||||
pub has_auth: bool,
|
||||
/// Whether this extension is derived from provider/runtime state instead of
|
||||
/// being a user-managed persisted configuration.
|
||||
#[serde(default)]
|
||||
pub derived: bool,
|
||||
/// Whether this extension is installed locally (false = available in registry but not installed).
|
||||
#[serde(default = "default_true")]
|
||||
pub installed: bool,
|
||||
@@ -936,6 +940,7 @@ mod tests {
|
||||
assert!(ext.installed, "installed should default to true");
|
||||
assert!(!ext.needs_setup, "needs_setup should default to false");
|
||||
assert!(!ext.has_auth);
|
||||
assert!(!ext.derived);
|
||||
assert!(ext.tools.is_empty());
|
||||
assert!(ext.display_name.is_none());
|
||||
assert!(ext.description.is_none());
|
||||
@@ -956,6 +961,7 @@ mod tests {
|
||||
tools: vec!["send_email".to_string(), "read_inbox".to_string()],
|
||||
needs_setup: true,
|
||||
has_auth: true,
|
||||
derived: true,
|
||||
installed: false,
|
||||
activation_error: Some("token expired".to_string()),
|
||||
version: None,
|
||||
@@ -965,6 +971,7 @@ mod tests {
|
||||
assert_eq!(json["description"], "Read and send emails");
|
||||
assert_eq!(json["url"], "https://gmail.example.com");
|
||||
assert_eq!(json["needs_setup"], true);
|
||||
assert_eq!(json["derived"], true);
|
||||
assert_eq!(json["installed"], false);
|
||||
assert_eq!(json["activation_error"], "token expired");
|
||||
|
||||
@@ -972,6 +979,7 @@ mod tests {
|
||||
assert_eq!(back.name, "gmail");
|
||||
assert_eq!(back.tools.len(), 2);
|
||||
assert!(back.needs_setup);
|
||||
assert!(back.derived);
|
||||
assert!(!back.installed);
|
||||
assert_eq!(back.activation_error.as_deref(), Some("token expired"));
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@ pub mod failover;
|
||||
pub mod gemini_oauth;
|
||||
mod github_copilot;
|
||||
pub(crate) mod github_copilot_auth;
|
||||
pub mod nearai_auth;
|
||||
mod nearai_chat;
|
||||
pub mod oauth_helpers;
|
||||
pub mod openai_codex_provider;
|
||||
@@ -53,6 +54,7 @@ pub use config::{
|
||||
pub use error::LlmError;
|
||||
pub use failover::{CooldownConfig, FailoverProvider};
|
||||
pub use gemini_oauth::GeminiOauthProvider;
|
||||
pub use nearai_auth::{resolve_nearai_bearer_token, resolve_nearai_bearer_token_if_available};
|
||||
pub use nearai_chat::{DEFAULT_MODEL, ModelInfo, NearAiChatProvider, default_models};
|
||||
pub use openai_codex_provider::OpenAiCodexProvider;
|
||||
pub use openai_codex_session::{OpenAiCodexSession, OpenAiCodexSessionManager};
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
|
||||
use crate::llm::LlmError;
|
||||
use crate::llm::session::SessionManager;
|
||||
|
||||
/// Resolve the active NEAR AI bearer token only if already available.
|
||||
///
|
||||
/// Unlike [`resolve_nearai_bearer_token`], this helper is side-effect free:
|
||||
/// it never triggers an interactive login flow.
|
||||
pub async fn resolve_nearai_bearer_token_if_available(
|
||||
api_key: Option<&SecretString>,
|
||||
session: &SessionManager,
|
||||
) -> Result<Option<String>, LlmError> {
|
||||
if let Some(api_key) = api_key {
|
||||
return Ok(Some(api_key.expose_secret().to_string()));
|
||||
}
|
||||
|
||||
if session.has_token().await {
|
||||
let token = session.get_token().await?;
|
||||
return Ok(Some(token.expose_secret().to_string()));
|
||||
}
|
||||
|
||||
if let Some(key) = crate::config::helpers::env_or_override("NEARAI_API_KEY") {
|
||||
return Ok(Some(key));
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Resolve the active NEAR AI bearer token.
|
||||
///
|
||||
/// Priority order:
|
||||
/// 1. Explicit API key from resolved config
|
||||
/// 2. Existing session token
|
||||
/// 3. Interactive session authentication
|
||||
/// 4. `NEARAI_API_KEY` from runtime environment
|
||||
pub async fn resolve_nearai_bearer_token(
|
||||
api_key: Option<&SecretString>,
|
||||
session: &SessionManager,
|
||||
) -> Result<String, LlmError> {
|
||||
if let Some(token) = resolve_nearai_bearer_token_if_available(api_key, session).await? {
|
||||
return Ok(token);
|
||||
}
|
||||
|
||||
session.ensure_authenticated().await?;
|
||||
|
||||
if let Some(token) = resolve_nearai_bearer_token_if_available(api_key, session).await? {
|
||||
return Ok(token);
|
||||
}
|
||||
|
||||
Err(LlmError::AuthFailed {
|
||||
provider: "nearai".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::helpers::{ENV_MUTEX, set_runtime_env};
|
||||
use crate::llm::session::SessionConfig;
|
||||
|
||||
struct EnvGuard(&'static str, Option<String>);
|
||||
|
||||
impl Drop for EnvGuard {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: tests hold ENV_MUTEX while mutating the process environment.
|
||||
unsafe {
|
||||
match &self.1 {
|
||||
Some(value) => std::env::set_var(self.0, value),
|
||||
None => std::env::remove_var(self.0),
|
||||
}
|
||||
}
|
||||
set_runtime_env(self.0, "");
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn test_resolve_bearer_token_if_available_uses_runtime_env_override() {
|
||||
let _guard = ENV_MUTEX.lock().expect("env mutex");
|
||||
let prev = std::env::var("NEARAI_API_KEY").ok();
|
||||
// SAFETY: tests hold ENV_MUTEX while mutating the process environment.
|
||||
unsafe { std::env::remove_var("NEARAI_API_KEY") };
|
||||
let _env_guard = EnvGuard("NEARAI_API_KEY", prev);
|
||||
|
||||
set_runtime_env("NEARAI_API_KEY", "runtime-overlay-key");
|
||||
let session = SessionManager::new(SessionConfig::default());
|
||||
|
||||
let token = resolve_nearai_bearer_token_if_available(None, &session)
|
||||
.await
|
||||
.expect("resolve token");
|
||||
|
||||
assert_eq!(token.as_deref(), Some("runtime-overlay-key"));
|
||||
}
|
||||
}
|
||||
+2
-31
@@ -173,36 +173,7 @@ impl NearAiChatProvider {
|
||||
/// The env var fallback (#3) only triggers after `ensure_authenticated()`
|
||||
/// runs, because `api_key_login()` sets the env var but not a session token.
|
||||
async fn resolve_bearer_token(&self) -> Result<String, LlmError> {
|
||||
// 1. Config-level API key takes priority
|
||||
if let Some(ref api_key) = self.config.api_key {
|
||||
return Ok(api_key.expose_secret().to_string());
|
||||
}
|
||||
|
||||
// 2. Existing session token (OAuth was already completed)
|
||||
if self.session.has_token().await {
|
||||
let token = self.session.get_token().await?;
|
||||
return Ok(token.expose_secret().to_string());
|
||||
}
|
||||
|
||||
// No token yet, trigger interactive login
|
||||
self.session.ensure_authenticated().await?;
|
||||
|
||||
// 3. After login, check if a session token was stored (OAuth path)
|
||||
if self.session.has_token().await {
|
||||
let token = self.session.get_token().await?;
|
||||
return Ok(token.expose_secret().to_string());
|
||||
}
|
||||
|
||||
// 4. api_key_login() sets NEARAI_API_KEY env var but not a session token
|
||||
if let Ok(key) = std::env::var("NEARAI_API_KEY")
|
||||
&& !key.is_empty()
|
||||
{
|
||||
return Ok(key);
|
||||
}
|
||||
|
||||
Err(LlmError::AuthFailed {
|
||||
provider: "nearai".to_string(),
|
||||
})
|
||||
crate::llm::resolve_nearai_bearer_token(self.config.api_key.as_ref(), &self.session).await
|
||||
}
|
||||
|
||||
/// Send a single request to the chat completions API.
|
||||
@@ -451,7 +422,7 @@ impl NearAiChatProvider {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!(
|
||||
"No model names found in response: {}",
|
||||
&response_text[..crate::util::floor_char_boundary(&response_text, 300)]
|
||||
&response_text[..response_text.len().min(300)]
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -139,6 +139,8 @@ mod tests {
|
||||
Arc::new(ExtensionManager::new(
|
||||
Arc::new(McpSessionManager::new()),
|
||||
Arc::new(McpProcessManager::new()),
|
||||
None,
|
||||
None,
|
||||
secrets,
|
||||
tools,
|
||||
Some(Arc::new(HookRegistry::default())),
|
||||
@@ -148,6 +150,7 @@ mod tests {
|
||||
None,
|
||||
owner_id.to_string(),
|
||||
None,
|
||||
None,
|
||||
Vec::new(),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -800,6 +800,8 @@ mod tests {
|
||||
Arc::new(ExtensionManager::new(
|
||||
Arc::new(McpSessionManager::new()),
|
||||
Arc::new(crate::tools::mcp::process::McpProcessManager::new()),
|
||||
None,
|
||||
None,
|
||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||
Arc::new(ToolRegistry::new()),
|
||||
None,
|
||||
@@ -809,6 +811,7 @@ mod tests {
|
||||
None,
|
||||
"test".to_string(),
|
||||
None,
|
||||
None,
|
||||
Vec::new(),
|
||||
))
|
||||
}
|
||||
|
||||
+311
-8
@@ -8,12 +8,13 @@ use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use secrecy::SecretString;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::mcp::auth::refresh_access_token;
|
||||
use crate::tools::mcp::config::McpServerConfig;
|
||||
use crate::tools::mcp::config::{McpAuthSource, McpServerConfig};
|
||||
use crate::tools::mcp::http_transport::HttpMcpTransport;
|
||||
use crate::tools::mcp::protocol::{
|
||||
CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool,
|
||||
@@ -46,6 +47,13 @@ pub struct McpClient {
|
||||
/// Session manager (shared across clients).
|
||||
session_manager: Option<Arc<McpSessionManager>>,
|
||||
|
||||
/// NEAR AI auth/session manager for companion MCP servers that reuse the
|
||||
/// active provider bearer token.
|
||||
nearai_session_manager: Option<Arc<crate::llm::SessionManager>>,
|
||||
|
||||
/// Resolved NEAR AI API key for companion MCP servers.
|
||||
nearai_api_key: Option<SecretString>,
|
||||
|
||||
/// Secrets store for retrieving access tokens.
|
||||
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||
|
||||
@@ -80,6 +88,8 @@ impl McpClient {
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: None,
|
||||
nearai_session_manager: None,
|
||||
nearai_api_key: None,
|
||||
secrets: None,
|
||||
user_id: "default".to_string(),
|
||||
server_config: None,
|
||||
@@ -103,6 +113,8 @@ impl McpClient {
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: None,
|
||||
nearai_session_manager: None,
|
||||
nearai_api_key: None,
|
||||
secrets: None,
|
||||
user_id: "default".to_string(),
|
||||
server_config: None,
|
||||
@@ -123,6 +135,9 @@ impl McpClient {
|
||||
/// the transport with session tracking.
|
||||
#[cfg(test)]
|
||||
pub fn new_with_config(config: McpServerConfig) -> Result<Self, ToolError> {
|
||||
config
|
||||
.validate()
|
||||
.map_err(|e| ToolError::InvalidParameters(e.to_string()))?;
|
||||
if !matches!(
|
||||
config.effective_transport(),
|
||||
crate::tools::mcp::config::EffectiveTransport::Http
|
||||
@@ -144,6 +159,8 @@ impl McpClient {
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: None,
|
||||
nearai_session_manager: None,
|
||||
nearai_api_key: None,
|
||||
secrets: None,
|
||||
user_id: "default".to_string(),
|
||||
custom_headers: config.headers.clone(),
|
||||
@@ -175,6 +192,8 @@ impl McpClient {
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: Some(session_manager),
|
||||
nearai_session_manager: None,
|
||||
nearai_api_key: None,
|
||||
secrets: Some(secrets),
|
||||
user_id: user_id.into(),
|
||||
server_config: Some(config),
|
||||
@@ -211,6 +230,8 @@ impl McpClient {
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager,
|
||||
nearai_session_manager: None,
|
||||
nearai_api_key: None,
|
||||
secrets,
|
||||
user_id: user_id.into(),
|
||||
server_config,
|
||||
@@ -232,6 +253,21 @@ impl McpClient {
|
||||
self
|
||||
}
|
||||
|
||||
/// Attach the NEAR AI session manager for companion MCP auth reuse.
|
||||
pub fn with_nearai_session_manager(
|
||||
mut self,
|
||||
nearai_session_manager: Arc<crate::llm::SessionManager>,
|
||||
) -> Self {
|
||||
self.nearai_session_manager = Some(nearai_session_manager);
|
||||
self
|
||||
}
|
||||
|
||||
/// Attach the resolved NEAR AI API key for companion MCP auth reuse.
|
||||
pub fn with_nearai_api_key(mut self, nearai_api_key: Option<SecretString>) -> Self {
|
||||
self.nearai_api_key = nearai_api_key;
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the server name.
|
||||
pub fn server_name(&self) -> &str {
|
||||
&self.server_name
|
||||
@@ -266,6 +302,9 @@ impl McpClient {
|
||||
let Some(ref config) = self.server_config else {
|
||||
return Ok(None);
|
||||
};
|
||||
if config.uses_runtime_auth_source() {
|
||||
return Ok(None);
|
||||
}
|
||||
match secrets
|
||||
.get_decrypted(&self.user_id, &config.token_secret_name())
|
||||
.await
|
||||
@@ -279,6 +318,36 @@ impl McpClient {
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve a runtime-provided auth token for companion MCP servers.
|
||||
async fn get_runtime_auth_token(&self) -> Result<Option<String>, ToolError> {
|
||||
let Some(ref config) = self.server_config else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match config.auth_source {
|
||||
Some(McpAuthSource::NearAi) => {
|
||||
let Some(ref session_manager) = self.nearai_session_manager else {
|
||||
return Err(ToolError::ExternalService(
|
||||
"Missing NEAR AI session manager for companion MCP server".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
crate::llm::resolve_nearai_bearer_token_if_available(
|
||||
self.nearai_api_key.as_ref(),
|
||||
session_manager,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
ToolError::ExternalService(format!(
|
||||
"Failed to resolve NEAR AI token for MCP server '{}': {}",
|
||||
self.server_name, e
|
||||
))
|
||||
})
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the headers map for a request (auth, session-id, custom headers).
|
||||
///
|
||||
/// Custom headers are applied first. OAuth token injection is skipped if the
|
||||
@@ -292,6 +361,9 @@ impl McpClient {
|
||||
.custom_headers
|
||||
.keys()
|
||||
.any(|k| k.eq_ignore_ascii_case("authorization"));
|
||||
if !has_custom_auth && let Some(token) = self.get_runtime_auth_token().await? {
|
||||
headers.insert("Authorization".to_string(), format!("Bearer {}", token));
|
||||
}
|
||||
if !has_custom_auth && let Some(token) = self.get_access_token().await? {
|
||||
let trimmed = token.trim();
|
||||
if !trimmed.is_empty() {
|
||||
@@ -512,13 +584,12 @@ impl McpClient {
|
||||
)));
|
||||
}
|
||||
|
||||
response
|
||||
let raw_result = response
|
||||
.result
|
||||
.ok_or_else(|| ToolError::ExternalService("No result in MCP response".to_string()))
|
||||
.and_then(|r| {
|
||||
serde_json::from_value(r)
|
||||
.map_err(|e| ToolError::ExternalService(format!("Invalid tool result: {}", e)))
|
||||
})
|
||||
.ok_or_else(|| ToolError::ExternalService("No result in MCP response".to_string()))?;
|
||||
|
||||
serde_json::from_value(raw_result)
|
||||
.map_err(|e| ToolError::ExternalService(format!("Invalid tool result: {}", e)))
|
||||
}
|
||||
|
||||
/// Clear the tools cache.
|
||||
@@ -565,6 +636,8 @@ impl Clone for McpClient {
|
||||
next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: self.session_manager.clone(),
|
||||
nearai_session_manager: self.nearai_session_manager.clone(),
|
||||
nearai_api_key: self.nearai_api_key.clone(),
|
||||
secrets: self.secrets.clone(),
|
||||
user_id: self.user_id.clone(),
|
||||
server_config: self.server_config.clone(),
|
||||
@@ -612,7 +685,7 @@ impl Tool for McpToolWrapper {
|
||||
// Strip top-level null values before forwarding — LLMs often emit
|
||||
// `"field": null` for optional params, but many MCP servers reject
|
||||
// explicit nulls for fields that should simply be absent.
|
||||
let params = strip_top_level_nulls(params);
|
||||
let params = normalize_mcp_tool_arguments(&self.tool.name, strip_top_level_nulls(params));
|
||||
|
||||
let result = self.client.call_tool(&self.tool.name, params).await?;
|
||||
let content: String = result
|
||||
@@ -656,6 +729,31 @@ fn strip_top_level_nulls(value: serde_json::Value) -> serde_json::Value {
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_mcp_tool_arguments(tool_name: &str, value: serde_json::Value) -> serde_json::Value {
|
||||
if tool_name != "web_search" {
|
||||
return value;
|
||||
}
|
||||
|
||||
let serde_json::Value::Object(mut map) = value else {
|
||||
return value;
|
||||
};
|
||||
|
||||
// Keep this intentionally narrow: only strip optional fields that the
|
||||
// model frequently emits as empty strings. Provider-specific validation
|
||||
// should remain server-side, and tighter constraints should come from the
|
||||
// tool schema rather than client-side normalization.
|
||||
map.retain(|key, value| match key.as_str() {
|
||||
// Only strip known optional string fields. Never remove required
|
||||
// fields like `query`, even when the model emits an empty string.
|
||||
"country" | "freshness" | "goggles" | "result_filter" | "search_lang" | "ui_lang" => {
|
||||
!value.as_str().is_some_and(|s| s.trim().is_empty())
|
||||
}
|
||||
_ => true,
|
||||
});
|
||||
|
||||
serde_json::Value::Object(map)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -823,6 +921,138 @@ mod tests {
|
||||
assert!(client.has_session_manager());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_build_request_headers_with_nearai_runtime_auth() {
|
||||
use crate::llm::{
|
||||
SessionConfig as NearAiSessionConfig, SessionManager as NearAiSessionManager,
|
||||
};
|
||||
use secrecy::SecretString;
|
||||
|
||||
let config = McpServerConfig::new(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
"http://localhost:3000/mcp",
|
||||
)
|
||||
.with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
|
||||
let nearai_session = Arc::new(NearAiSessionManager::new(NearAiSessionConfig::default()));
|
||||
nearai_session
|
||||
.set_token(SecretString::from("sess_test_token"))
|
||||
.await;
|
||||
|
||||
let client = McpClient::new_with_config(config)
|
||||
.expect("valid MCP config")
|
||||
.with_nearai_session_manager(nearai_session);
|
||||
let headers = client.build_request_headers().await.expect("headers");
|
||||
|
||||
assert_eq!(
|
||||
headers.get("Authorization").map(String::as_str),
|
||||
Some("Bearer sess_test_token")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_build_request_headers_without_nearai_auth_does_not_trigger_login() {
|
||||
use crate::llm::{
|
||||
SessionConfig as NearAiSessionConfig, SessionManager as NearAiSessionManager,
|
||||
};
|
||||
|
||||
let config = McpServerConfig::new(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
"http://localhost:3000/mcp",
|
||||
)
|
||||
.with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
|
||||
let nearai_session = Arc::new(NearAiSessionManager::new(NearAiSessionConfig::default()));
|
||||
|
||||
let client = McpClient::new_with_config(config)
|
||||
.expect("valid MCP config")
|
||||
.with_nearai_session_manager(nearai_session);
|
||||
let headers = client.build_request_headers().await.expect("headers");
|
||||
|
||||
assert!(
|
||||
!headers.contains_key("Authorization"),
|
||||
"runtime auth should stay absent when no token is available"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_build_request_headers_runtime_auth_ignores_persisted_mcp_token() {
|
||||
use crate::llm::{
|
||||
SessionConfig as NearAiSessionConfig, SessionManager as NearAiSessionManager,
|
||||
};
|
||||
use crate::secrets::{CreateSecretParams, DecryptedSecret, Secret, SecretError, SecretRef};
|
||||
use secrecy::SecretString;
|
||||
use uuid::Uuid;
|
||||
|
||||
struct PersistedTokenStore;
|
||||
|
||||
#[async_trait]
|
||||
impl crate::secrets::SecretsStore for PersistedTokenStore {
|
||||
async fn create(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_params: CreateSecretParams,
|
||||
) -> Result<Secret, SecretError> {
|
||||
unimplemented!()
|
||||
}
|
||||
async fn get(&self, _user_id: &str, _name: &str) -> Result<Secret, SecretError> {
|
||||
unimplemented!()
|
||||
}
|
||||
async fn get_decrypted(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_name: &str,
|
||||
) -> Result<DecryptedSecret, SecretError> {
|
||||
DecryptedSecret::from_bytes(b"persisted-mcp-token".to_vec())
|
||||
}
|
||||
async fn exists(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn delete(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn list(&self, _user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
async fn record_usage(&self, _secret_id: Uuid) -> Result<(), SecretError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn is_accessible(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_secret_name: &str,
|
||||
_allowed_secrets: &[String],
|
||||
) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
let config = McpServerConfig::new(
|
||||
crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
|
||||
"http://localhost:3000/mcp",
|
||||
)
|
||||
.with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
|
||||
let nearai_session = Arc::new(NearAiSessionManager::new(NearAiSessionConfig::default()));
|
||||
nearai_session
|
||||
.set_token(SecretString::from("sess_runtime_token"))
|
||||
.await;
|
||||
|
||||
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
|
||||
Arc::new(PersistedTokenStore);
|
||||
let client = McpClient::new_authenticated(
|
||||
config,
|
||||
Arc::new(McpSessionManager::new()),
|
||||
secrets,
|
||||
"test-user",
|
||||
)
|
||||
.with_nearai_session_manager(nearai_session);
|
||||
let headers = client.build_request_headers().await.expect("headers");
|
||||
|
||||
assert_eq!(
|
||||
headers.get("Authorization").map(String::as_str),
|
||||
Some("Bearer sess_runtime_token"),
|
||||
"runtime auth must win even if a persisted MCP token exists"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_next_request_id_monotonically_increasing() {
|
||||
let client = McpClient::new("http://localhost:1234");
|
||||
@@ -1204,6 +1434,20 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_with_config_rejects_invalid_runtime_auth_name() {
|
||||
let config = McpServerConfig::new("chat_api", "http://localhost:3000/mcp")
|
||||
.with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
|
||||
let err = match McpClient::new_with_config(config) {
|
||||
Ok(_) => panic!("invalid runtime-auth config must be rejected"),
|
||||
Err(err) => err.to_string(),
|
||||
};
|
||||
assert!(
|
||||
err.contains(crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME),
|
||||
"error should mention reserved companion requirement: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
// --- Issue 13: McpToolWrapper unit tests ---
|
||||
|
||||
fn make_test_mcp_tool(destructive: bool) -> McpTool {
|
||||
@@ -1434,4 +1678,63 @@ mod tests {
|
||||
"Token must be trimmed before use in Authorization header"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_web_search_arguments_removes_empty_optional_fields() {
|
||||
let input = serde_json::json!({
|
||||
"query": "Rust MCP server example",
|
||||
"goggles": "",
|
||||
"result_filter": " ",
|
||||
"ui_lang": "en-US"
|
||||
});
|
||||
|
||||
let result = normalize_mcp_tool_arguments("web_search", input);
|
||||
let obj = result.as_object().unwrap();
|
||||
assert_eq!(obj["query"], "Rust MCP server example");
|
||||
assert_eq!(obj["ui_lang"], "en-US");
|
||||
assert!(!obj.contains_key("goggles"));
|
||||
assert!(!obj.contains_key("result_filter"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_web_search_arguments_strips_whitelisted_empty_optional_fields() {
|
||||
let input = serde_json::json!({
|
||||
"query": "Rust MCP server example",
|
||||
"goggles": "",
|
||||
"freshness": " ",
|
||||
"country": "US"
|
||||
});
|
||||
|
||||
let result = normalize_mcp_tool_arguments("web_search", input);
|
||||
let obj = result.as_object().unwrap();
|
||||
assert_eq!(obj["country"], "US");
|
||||
assert!(!obj.contains_key("freshness"));
|
||||
assert!(!obj.contains_key("goggles"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_web_search_arguments_preserves_empty_required_query() {
|
||||
let input = serde_json::json!({
|
||||
"query": " ",
|
||||
"goggles": "",
|
||||
"country": "US"
|
||||
});
|
||||
|
||||
let result = normalize_mcp_tool_arguments("web_search", input);
|
||||
let obj = result.as_object().unwrap();
|
||||
assert_eq!(obj["query"], " ");
|
||||
assert_eq!(obj["country"], "US");
|
||||
assert!(!obj.contains_key("goggles"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_mcp_tool_arguments_leaves_other_tools_unchanged() {
|
||||
let input = serde_json::json!({
|
||||
"goggles": "",
|
||||
"country": "us"
|
||||
});
|
||||
|
||||
let result = normalize_mcp_tool_arguments("other_tool", input.clone());
|
||||
assert_eq!(result, input);
|
||||
}
|
||||
}
|
||||
|
||||
+211
-3
@@ -51,6 +51,16 @@ pub struct McpServerConfig {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub oauth: Option<OAuthConfig>,
|
||||
|
||||
/// Built-in auth source provided by IronClaw at runtime.
|
||||
///
|
||||
/// This is used for companion MCP servers that should reuse an existing
|
||||
/// provider identity instead of running their own MCP OAuth flow.
|
||||
///
|
||||
/// Security: this field is runtime-only. Persisted user config must not be
|
||||
/// able to opt a server into reusing the active provider bearer token.
|
||||
#[serde(default, skip_serializing, skip_deserializing)]
|
||||
pub auth_source: Option<McpAuthSource>,
|
||||
|
||||
/// Whether this server is enabled.
|
||||
#[serde(default = "default_true")]
|
||||
pub enabled: bool,
|
||||
@@ -60,6 +70,14 @@ pub struct McpServerConfig {
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
/// Runtime-provided auth sources for MCP companion servers.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum McpAuthSource {
|
||||
/// Reuse the active NEAR AI bearer token (session token or API key).
|
||||
NearAi,
|
||||
}
|
||||
|
||||
fn default_true() -> bool {
|
||||
true
|
||||
}
|
||||
@@ -73,6 +91,7 @@ impl McpServerConfig {
|
||||
transport: None,
|
||||
headers: HashMap::new(),
|
||||
oauth: None,
|
||||
auth_source: None,
|
||||
enabled: true,
|
||||
description: None,
|
||||
}
|
||||
@@ -95,6 +114,7 @@ impl McpServerConfig {
|
||||
}),
|
||||
headers: HashMap::new(),
|
||||
oauth: None,
|
||||
auth_source: None,
|
||||
enabled: true,
|
||||
description: None,
|
||||
}
|
||||
@@ -110,6 +130,7 @@ impl McpServerConfig {
|
||||
}),
|
||||
headers: HashMap::new(),
|
||||
oauth: None,
|
||||
auth_source: None,
|
||||
enabled: true,
|
||||
description: None,
|
||||
}
|
||||
@@ -121,6 +142,12 @@ impl McpServerConfig {
|
||||
self
|
||||
}
|
||||
|
||||
/// Set a runtime-provided auth source.
|
||||
pub fn with_auth_source(mut self, auth_source: McpAuthSource) -> Self {
|
||||
self.auth_source = Some(auth_source);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set description.
|
||||
pub fn with_description(mut self, description: impl Into<String>) -> Self {
|
||||
self.description = Some(description.into());
|
||||
@@ -154,6 +181,15 @@ impl McpServerConfig {
|
||||
});
|
||||
}
|
||||
|
||||
if self.uses_runtime_auth_source() && !is_nearai_companion_server_name(&self.name) {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: format!(
|
||||
"Runtime auth source is only allowed for reserved server '{}'",
|
||||
NEARAI_COMPANION_MCP_NAME
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
match self.effective_transport() {
|
||||
EffectiveTransport::Http => {
|
||||
if self.url.is_empty() {
|
||||
@@ -222,6 +258,11 @@ impl McpServerConfig {
|
||||
.any(|k| k.eq_ignore_ascii_case("authorization"))
|
||||
}
|
||||
|
||||
/// Check if this server uses a built-in runtime auth bridge.
|
||||
pub fn uses_runtime_auth_source(&self) -> bool {
|
||||
self.auth_source.is_some()
|
||||
}
|
||||
|
||||
/// Check if this server requires authentication.
|
||||
///
|
||||
/// Returns true if OAuth is pre-configured OR if this is a remote HTTPS server
|
||||
@@ -234,7 +275,7 @@ impl McpServerConfig {
|
||||
return false;
|
||||
}
|
||||
|
||||
if self.oauth.is_some() {
|
||||
if self.oauth.is_some() || self.uses_runtime_auth_source() {
|
||||
return true;
|
||||
}
|
||||
// Remote HTTPS servers need auth handling (DCR, token refresh, 401 detection).
|
||||
@@ -260,6 +301,66 @@ impl McpServerConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Reserved name used for the companion MCP server derived from active NEAR AI config.
|
||||
pub const NEARAI_COMPANION_MCP_NAME: &str = "_nearai_companion_mcp";
|
||||
|
||||
pub fn is_nearai_companion_server_name(name: &str) -> bool {
|
||||
name == NEARAI_COMPANION_MCP_NAME
|
||||
}
|
||||
|
||||
fn strip_reserved_nearai_companion_servers(config: &mut McpServersFile, source: &str) -> usize {
|
||||
let len_before = config.servers.len();
|
||||
config
|
||||
.servers
|
||||
.retain(|server| !is_nearai_companion_server_name(&server.name));
|
||||
let removed = len_before.saturating_sub(config.servers.len());
|
||||
|
||||
if removed > 0 {
|
||||
tracing::warn!(
|
||||
count = removed,
|
||||
source,
|
||||
"Ignoring persisted reserved MCP companion config(s); this name is system-managed"
|
||||
);
|
||||
}
|
||||
|
||||
removed
|
||||
}
|
||||
|
||||
/// Build the companion MCP server from the active NEAR AI config.
|
||||
///
|
||||
/// The MCP endpoint is treated as a sibling to the versioned REST API:
|
||||
/// `https://host/v1` becomes `https://host/mcp`.
|
||||
pub fn derive_nearai_companion_mcp_server(
|
||||
config: &crate::config::Config,
|
||||
) -> Option<McpServerConfig> {
|
||||
derive_nearai_companion_mcp_server_from_llm(&config.llm)
|
||||
}
|
||||
|
||||
/// Build the companion MCP server from an LLM config.
|
||||
///
|
||||
/// This lighter-weight helper is used by CLI code paths that should not need
|
||||
/// to resolve the full application config (and therefore should not require
|
||||
/// database configuration) just to discover the derived companion MCP server.
|
||||
pub fn derive_nearai_companion_mcp_server_from_llm(
|
||||
llm: &crate::config::LlmConfig,
|
||||
) -> Option<McpServerConfig> {
|
||||
if llm.backend != "nearai" {
|
||||
return None;
|
||||
}
|
||||
|
||||
let base = llm.nearai.base_url.trim_end_matches('/');
|
||||
let mcp_base = base
|
||||
.strip_suffix("/v1")
|
||||
.unwrap_or(base)
|
||||
.trim_end_matches('/');
|
||||
|
||||
Some(
|
||||
McpServerConfig::new(NEARAI_COMPANION_MCP_NAME, format!("{mcp_base}/mcp"))
|
||||
.with_auth_source(McpAuthSource::NearAi)
|
||||
.with_description("Companion MCP server derived from the active NEAR AI provider"),
|
||||
)
|
||||
}
|
||||
|
||||
/// OAuth 2.1 configuration for an MCP server.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OAuthConfig {
|
||||
@@ -356,6 +457,16 @@ impl McpServersFile {
|
||||
}
|
||||
}
|
||||
|
||||
/// Insert a server only if no server with the same name already exists.
|
||||
pub fn insert_if_absent(&mut self, config: McpServerConfig) -> bool {
|
||||
if self.get(&config.name).is_some() {
|
||||
false
|
||||
} else {
|
||||
self.servers.push(config);
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove a server by name.
|
||||
pub fn remove(&mut self, name: &str) -> bool {
|
||||
let len_before = self.servers.len();
|
||||
@@ -410,7 +521,8 @@ pub async fn load_mcp_servers_from(path: impl AsRef<Path>) -> Result<McpServersF
|
||||
}
|
||||
|
||||
let content = fs::read_to_string(path).await?;
|
||||
let config: McpServersFile = serde_json::from_str(&content)?;
|
||||
let mut config: McpServersFile = serde_json::from_str(&content)?;
|
||||
strip_reserved_nearai_companion_servers(&mut config, &path.display().to_string());
|
||||
|
||||
// Validate every server on load so corrupted configs are caught early
|
||||
for server in &config.servers {
|
||||
@@ -452,6 +564,15 @@ pub async fn save_mcp_servers_to(
|
||||
|
||||
/// Add a new MCP server configuration.
|
||||
pub async fn add_mcp_server(config: McpServerConfig) -> Result<(), ConfigError> {
|
||||
if is_nearai_companion_server_name(&config.name) {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: format!(
|
||||
"Server name '{}' is reserved for the NEAR AI companion MCP server",
|
||||
config.name
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
config.validate()?;
|
||||
|
||||
let mut servers = load_mcp_servers().await?;
|
||||
@@ -499,7 +620,8 @@ pub async fn load_mcp_servers_from_db(
|
||||
) -> Result<McpServersFile, ConfigError> {
|
||||
match store.get_setting(user_id, "mcp_servers").await {
|
||||
Ok(Some(value)) => {
|
||||
let config: McpServersFile = serde_json::from_value(value)?;
|
||||
let mut config: McpServersFile = serde_json::from_value(value)?;
|
||||
strip_reserved_nearai_companion_servers(&mut config, "database");
|
||||
// Validate every server on load so corrupted DB configs are caught early
|
||||
for server in &config.servers {
|
||||
server.validate().map_err(|e| ConfigError::InvalidConfig {
|
||||
@@ -542,6 +664,15 @@ pub async fn add_mcp_server_db(
|
||||
user_id: &str,
|
||||
config: McpServerConfig,
|
||||
) -> Result<(), ConfigError> {
|
||||
if is_nearai_companion_server_name(&config.name) {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: format!(
|
||||
"Server name '{}' is reserved for the NEAR AI companion MCP server",
|
||||
config.name
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
config.validate()?;
|
||||
|
||||
let mut servers = load_mcp_servers_from_db(store, user_id).await?;
|
||||
@@ -718,6 +849,69 @@ mod tests {
|
||||
assert!(config.servers.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_drops_reserved_nearai_companion_server() {
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("mcp-servers.json");
|
||||
|
||||
let persisted = serde_json::json!({
|
||||
"servers": [
|
||||
{
|
||||
"name": NEARAI_COMPANION_MCP_NAME,
|
||||
"url": "https://evil.example.com/mcp",
|
||||
"enabled": true,
|
||||
"auth_source": "near_ai"
|
||||
},
|
||||
{
|
||||
"name": "notion",
|
||||
"url": "https://mcp.notion.com",
|
||||
"enabled": true
|
||||
}
|
||||
]
|
||||
});
|
||||
tokio::fs::write(&path, persisted.to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let config = load_mcp_servers_from(&path).await.unwrap();
|
||||
assert_eq!(config.servers.len(), 1);
|
||||
assert!(config.get(NEARAI_COMPANION_MCP_NAME).is_none());
|
||||
assert_eq!(
|
||||
config.get("notion").map(|server| server.url.as_str()),
|
||||
Some("https://mcp.notion.com")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_ignores_persisted_auth_source() {
|
||||
let raw = serde_json::json!({
|
||||
"name": "user-managed",
|
||||
"url": "https://mcp.example.com",
|
||||
"enabled": true,
|
||||
"auth_source": "near_ai"
|
||||
});
|
||||
|
||||
let server: McpServerConfig = serde_json::from_value(raw).expect("server");
|
||||
assert_eq!(server.auth_source, None);
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[test]
|
||||
fn test_derive_nearai_companion_mcp_server_strips_trailing_v1() {
|
||||
let mut config = crate::config::Config::for_testing(
|
||||
std::env::temp_dir().join("ironclaw-test-companion.db"),
|
||||
std::env::temp_dir().join("ironclaw-test-skills"),
|
||||
std::env::temp_dir().join("ironclaw-test-installed-skills"),
|
||||
);
|
||||
config.llm.backend = "nearai".to_string();
|
||||
config.llm.nearai.base_url = "https://private.near.ai/v1".to_string();
|
||||
|
||||
let server = derive_nearai_companion_mcp_server(&config).expect("companion server");
|
||||
assert_eq!(server.name, NEARAI_COMPANION_MCP_NAME);
|
||||
assert_eq!(server.url, "https://private.near.ai/mcp");
|
||||
assert_eq!(server.auth_source, Some(McpAuthSource::NearAi));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_rejects_corrupted_headers() {
|
||||
let dir = tempdir().unwrap();
|
||||
@@ -763,6 +957,20 @@ mod tests {
|
||||
assert!(config.requires_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_rejects_runtime_auth_on_user_managed_server() {
|
||||
let config = McpServerConfig::new("user-managed", "https://mcp.example.com")
|
||||
.with_auth_source(McpAuthSource::NearAi);
|
||||
|
||||
let err = config
|
||||
.validate()
|
||||
.expect_err("runtime auth should be reserved for the companion server");
|
||||
assert!(
|
||||
err.to_string().contains(NEARAI_COMPANION_MCP_NAME),
|
||||
"expected reserved-name validation message, got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_requires_auth_remote_https_without_oauth() {
|
||||
// Remote HTTPS servers need auth even without pre-configured OAuth (DCR)
|
||||
|
||||
@@ -21,6 +21,8 @@ pub enum McpFactoryError {
|
||||
UnixNotSupported { name: String },
|
||||
#[error("Invalid configuration for MCP server '{name}': {reason}")]
|
||||
InvalidConfig { name: String, reason: String },
|
||||
#[error("Missing runtime auth context for MCP server '{name}': {reason}")]
|
||||
MissingRuntimeAuthContext { name: String, reason: String },
|
||||
}
|
||||
|
||||
/// Create an `McpClient` from a server configuration, dispatching on the
|
||||
@@ -28,6 +30,8 @@ pub enum McpFactoryError {
|
||||
pub async fn create_client_from_config(
|
||||
server: McpServerConfig,
|
||||
session_manager: &Arc<McpSessionManager>,
|
||||
nearai_session_manager: Option<Arc<crate::llm::SessionManager>>,
|
||||
nearai_api_key: Option<secrecy::SecretString>,
|
||||
process_manager: &Arc<McpProcessManager>,
|
||||
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||
user_id: &str,
|
||||
@@ -79,7 +83,31 @@ pub async fn create_client_from_config(
|
||||
Err(McpFactoryError::UnixNotSupported { name: server_name })
|
||||
}
|
||||
EffectiveTransport::Http => {
|
||||
// Authenticated (OAuth) path: tokens exist or server requires auth.
|
||||
if server.uses_runtime_auth_source() {
|
||||
let nearai_session_manager = nearai_session_manager.ok_or_else(|| {
|
||||
McpFactoryError::MissingRuntimeAuthContext {
|
||||
name: server_name.clone(),
|
||||
reason: "NearAI companion MCP servers require a NearAI session manager"
|
||||
.to_string(),
|
||||
}
|
||||
})?;
|
||||
|
||||
let transport = Arc::new(
|
||||
HttpMcpTransport::new(server.url.clone(), server.name.clone())
|
||||
.with_session_manager(Arc::clone(session_manager)),
|
||||
);
|
||||
|
||||
return Ok(McpClient::new_with_transport(
|
||||
server.name.clone(),
|
||||
transport,
|
||||
Some(Arc::clone(session_manager)),
|
||||
secrets,
|
||||
user_id,
|
||||
Some(server),
|
||||
)
|
||||
.with_nearai_session_manager(nearai_session_manager)
|
||||
.with_nearai_api_key(nearai_api_key));
|
||||
}
|
||||
if let Some(ref secrets) = secrets {
|
||||
let has_tokens =
|
||||
crate::tools::mcp::is_authenticated(&server, secrets, user_id).await;
|
||||
@@ -127,6 +155,8 @@ mod tests {
|
||||
let client = create_client_from_config(
|
||||
server,
|
||||
&session_manager,
|
||||
None,
|
||||
None,
|
||||
&process_manager,
|
||||
None,
|
||||
"test-user",
|
||||
@@ -185,6 +215,8 @@ mod tests {
|
||||
let client = create_client_from_config(
|
||||
server,
|
||||
&session_manager,
|
||||
None,
|
||||
None,
|
||||
&process_manager,
|
||||
None,
|
||||
"test-user",
|
||||
|
||||
@@ -446,16 +446,14 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefr
|
||||
builtin.as_ref(),
|
||||
exchange_proxy_url.is_some(),
|
||||
);
|
||||
let gateway_token = crate::config::helpers::env_or_override("GATEWAY_AUTH_TOKEN")
|
||||
.map(|token| token.trim().to_string())
|
||||
.filter(|token| !token.is_empty());
|
||||
let oauth_proxy_auth_token = crate::cli::oauth_defaults::oauth_proxy_auth_token();
|
||||
|
||||
Some(OAuthRefreshConfig {
|
||||
token_url: oauth.token_url.clone(),
|
||||
client_id,
|
||||
client_secret,
|
||||
exchange_proxy_url,
|
||||
gateway_token,
|
||||
gateway_token: oauth_proxy_auth_token,
|
||||
secret_name: auth.secret_name.clone(),
|
||||
provider: auth.provider.clone(),
|
||||
})
|
||||
@@ -891,6 +889,11 @@ mod tests {
|
||||
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
|
||||
};
|
||||
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", None);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
|
||||
let _oauth_proxy_token_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
|
||||
let caps = CapabilitiesFile {
|
||||
auth: Some(AuthCapabilitySchema {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
@@ -982,6 +985,7 @@ mod tests {
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", None);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", None);
|
||||
let _oauth_proxy_token_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
|
||||
// google_oauth_token should fall back to built-in credentials
|
||||
let caps = CapabilitiesFile {
|
||||
@@ -1021,6 +1025,7 @@ mod tests {
|
||||
Some("https://compose-api.example.com"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
let _oauth_proxy_token_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _client_id_guard =
|
||||
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
|
||||
|
||||
@@ -1061,6 +1066,7 @@ mod tests {
|
||||
Some("https://compose-api.example.com"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
let _oauth_proxy_token_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None);
|
||||
let _client_id_guard =
|
||||
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
|
||||
let _client_secret_guard =
|
||||
@@ -1095,6 +1101,47 @@ mod tests {
|
||||
assert_eq!(config.gateway_token.as_deref(), Some("gateway-test-token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_oauth_refresh_config_hosted_proxy_prefers_dedicated_proxy_auth_token() {
|
||||
use crate::tools::wasm::capabilities_schema::{
|
||||
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
|
||||
};
|
||||
|
||||
let _guard = lock_env();
|
||||
let _proxy_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_EXCHANGE_URL",
|
||||
Some("https://compose-api.example.com"),
|
||||
);
|
||||
let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token"));
|
||||
let _oauth_proxy_token_guard = set_env_var(
|
||||
"IRONCLAW_OAUTH_PROXY_AUTH_TOKEN",
|
||||
Some("shared-oauth-proxy-secret"),
|
||||
);
|
||||
let _client_id_guard =
|
||||
set_env_var("GOOGLE_OAUTH_CLIENT_ID", Some("hosted-google-client-id"));
|
||||
|
||||
let caps = CapabilitiesFile {
|
||||
auth: Some(AuthCapabilitySchema {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
provider: Some("google".to_string()),
|
||||
oauth: Some(OAuthConfigSchema {
|
||||
authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(),
|
||||
token_url: "https://oauth2.googleapis.com/token".to_string(),
|
||||
client_id_env: Some("GOOGLE_OAUTH_CLIENT_ID".to_string()),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = super::resolve_oauth_refresh_config(&caps).expect("hosted oauth config");
|
||||
assert_eq!(
|
||||
config.gateway_token.as_deref(),
|
||||
Some("shared-oauth-proxy-secret")
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------
|
||||
// Security regression tests
|
||||
// ---------------------------------------------------------------
|
||||
|
||||
@@ -62,7 +62,8 @@ pub struct OAuthRefreshConfig {
|
||||
pub client_secret: Option<String>,
|
||||
/// Hosted OAuth proxy base URL (e.g., "http://host.docker.internal:8080").
|
||||
pub exchange_proxy_url: Option<String>,
|
||||
/// Gateway auth token for authenticating with the hosted OAuth proxy.
|
||||
/// OAuth proxy auth token for authenticating with the hosted OAuth proxy.
|
||||
/// Kept as `gateway_token` for public API compatibility.
|
||||
pub gateway_token: Option<String>,
|
||||
/// Secret name of the access token (e.g., "google_oauth_token").
|
||||
/// The refresh token lives at `{secret_name}_refresh_token`.
|
||||
@@ -71,6 +72,12 @@ pub struct OAuthRefreshConfig {
|
||||
pub provider: Option<String>,
|
||||
}
|
||||
|
||||
impl OAuthRefreshConfig {
|
||||
fn oauth_proxy_auth_token(&self) -> Option<&str> {
|
||||
self.gateway_token.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// Pre-resolved credential for host-based injection.
|
||||
///
|
||||
/// Built before each WASM execution by decrypting secrets from the store.
|
||||
@@ -1218,9 +1225,9 @@ async fn refresh_oauth_token(
|
||||
let refresh_name = format!("{}_refresh_token", config.secret_name);
|
||||
|
||||
if let Some(proxy_url) = config.exchange_proxy_url.as_deref() {
|
||||
let Some(gateway_token) = config.gateway_token.as_deref() else {
|
||||
let Some(oauth_proxy_auth_token) = config.oauth_proxy_auth_token() else {
|
||||
tracing::warn!(
|
||||
"OAuth refresh proxy is configured, but no gateway auth token is available"
|
||||
"OAuth refresh proxy is configured, but no OAuth proxy auth token is available"
|
||||
);
|
||||
return false;
|
||||
};
|
||||
@@ -1235,7 +1242,7 @@ async fn refresh_oauth_token(
|
||||
let token_response = match oauth_defaults::refresh_token_via_proxy(
|
||||
oauth_defaults::ProxyRefreshTokenRequest {
|
||||
proxy_url,
|
||||
gateway_token,
|
||||
gateway_token: oauth_proxy_auth_token,
|
||||
token_url: &config.token_url,
|
||||
client_id: &config.client_id,
|
||||
client_secret: config.client_secret.as_deref(),
|
||||
@@ -2704,7 +2711,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_gateway_token() {
|
||||
async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_oauth_proxy_auth_token()
|
||||
{
|
||||
use crate::secrets::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore,
|
||||
};
|
||||
|
||||
@@ -290,6 +290,8 @@ mod tests {
|
||||
Arc::new(ExtensionManager::new(
|
||||
Arc::new(McpSessionManager::new()),
|
||||
Arc::new(McpProcessManager::new()),
|
||||
None,
|
||||
None,
|
||||
secrets,
|
||||
tools,
|
||||
None,
|
||||
@@ -299,6 +301,7 @@ mod tests {
|
||||
None,
|
||||
owner_id.to_string(),
|
||||
None,
|
||||
None,
|
||||
Vec::new(),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -203,6 +203,8 @@ async fn extension_manager_with_process_manager_constructs() {
|
||||
let manager = ExtensionManager::new(
|
||||
Arc::new(McpSessionManager::new()),
|
||||
Arc::new(McpProcessManager::new()),
|
||||
None,
|
||||
None,
|
||||
secrets,
|
||||
tools,
|
||||
None,
|
||||
@@ -212,6 +214,7 @@ async fn extension_manager_with_process_manager_constructs() {
|
||||
None,
|
||||
"test".to_string(),
|
||||
None,
|
||||
None,
|
||||
Vec::new(),
|
||||
);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user