mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-31 16:49:34 +00:00
Compare commits
29
Commits
main
...
feat/nearai-mcp
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
98e6e2601c | ||
|
|
b5c9c9f5d6 | ||
|
|
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
|
# Python
|
||||||
__pycache__/
|
__pycache__/
|
||||||
*.pyc
|
*.pyc
|
||||||
|
*.pyo
|
||||||
|
*.pyd
|
||||||
|
|
||||||
# Benchmark results (local runs, not committed)
|
# Benchmark results (local runs, not committed)
|
||||||
bench-results/
|
bench-results/
|
||||||
@@ -34,8 +36,5 @@ trace_*.json
|
|||||||
.claude/settings.local.json
|
.claude/settings.local.json
|
||||||
.worktrees/
|
.worktrees/
|
||||||
|
|
||||||
# Python cache
|
# JetBrains IDE
|
||||||
__pycache__/
|
.idea
|
||||||
*.pyc
|
|
||||||
*.pyo
|
|
||||||
*.pyd
|
|
||||||
|
|||||||
Generated
+1
-1
@@ -269,7 +269,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "whatsapp-channel"
|
name = "whatsapp-channel"
|
||||||
version = "0.1.0"
|
version = "0.2.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
|
|||||||
@@ -47,6 +47,15 @@ impl Agent {
|
|||||||
thread_id: Uuid,
|
thread_id: Uuid,
|
||||||
initial_messages: Vec<ChatMessage>,
|
initial_messages: Vec<ChatMessage>,
|
||||||
) -> Result<AgenticLoopResult, Error> {
|
) -> 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)
|
// Detect group chat from channel metadata (needed before loading system prompt)
|
||||||
let is_group_chat = message
|
let is_group_chat = message
|
||||||
.metadata
|
.metadata
|
||||||
|
|||||||
+20
-1
@@ -449,6 +449,8 @@ impl AppBuilder {
|
|||||||
|
|
||||||
let mcp_session_manager = Arc::new(McpSessionManager::new());
|
let mcp_session_manager = Arc::new(McpSessionManager::new());
|
||||||
let mcp_process_manager = Arc::new(McpProcessManager::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
|
// 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
|
// (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 mcp_sm = Arc::clone(&mcp_session_manager);
|
||||||
let pm = Arc::clone(&mcp_process_manager);
|
let pm = Arc::clone(&mcp_process_manager);
|
||||||
let owner_id = self.config.owner_id.clone();
|
let owner_id = self.config.owner_id.clone();
|
||||||
|
let companion_mcp_server = companion_mcp_server.clone();
|
||||||
async move {
|
async move {
|
||||||
let servers_result = if let Some(ref d) = db {
|
let servers_result = if let Some(ref d) = db {
|
||||||
load_mcp_servers_from_db(d.as_ref(), &owner_id).await
|
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
|
crate::tools::mcp::config::load_mcp_servers().await
|
||||||
};
|
};
|
||||||
match servers_result {
|
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();
|
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
|
||||||
if !enabled.is_empty() {
|
if !enabled.is_empty() {
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
@@ -545,6 +557,8 @@ impl AppBuilder {
|
|||||||
let mut join_set = tokio::task::JoinSet::new();
|
let mut join_set = tokio::task::JoinSet::new();
|
||||||
for server in enabled {
|
for server in enabled {
|
||||||
let mcp_sm = Arc::clone(&mcp_sm);
|
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 secrets = secrets_store.clone();
|
||||||
let tools = Arc::clone(&tools);
|
let tools = Arc::clone(&tools);
|
||||||
let pm = Arc::clone(&pm);
|
let pm = Arc::clone(&pm);
|
||||||
@@ -556,6 +570,8 @@ impl AppBuilder {
|
|||||||
let client = match crate::tools::mcp::create_client_from_config(
|
let client = match crate::tools::mcp::create_client_from_config(
|
||||||
server,
|
server,
|
||||||
&mcp_sm,
|
&mcp_sm,
|
||||||
|
Some(nearai_session),
|
||||||
|
nearai_api_key,
|
||||||
&pm,
|
&pm,
|
||||||
secrets,
|
secrets,
|
||||||
&owner_id,
|
&owner_id,
|
||||||
@@ -712,6 +728,8 @@ impl AppBuilder {
|
|||||||
let manager = Arc::new(ExtensionManager::new(
|
let manager = Arc::new(ExtensionManager::new(
|
||||||
Arc::clone(&mcp_session_manager),
|
Arc::clone(&mcp_session_manager),
|
||||||
Arc::clone(&mcp_process_manager),
|
Arc::clone(&mcp_process_manager),
|
||||||
|
Some(Arc::clone(&self.session)),
|
||||||
|
self.config.llm.nearai.api_key.clone(),
|
||||||
ext_secrets,
|
ext_secrets,
|
||||||
Arc::clone(tools),
|
Arc::clone(tools),
|
||||||
Some(Arc::clone(hooks)),
|
Some(Arc::clone(hooks)),
|
||||||
@@ -721,6 +739,7 @@ impl AppBuilder {
|
|||||||
self.config.tunnel.public_url.clone(),
|
self.config.tunnel.public_url.clone(),
|
||||||
self.config.owner_id.clone(),
|
self.config.owner_id.clone(),
|
||||||
self.db.clone(),
|
self.db.clone(),
|
||||||
|
companion_mcp_server,
|
||||||
catalog_entries.clone(),
|
catalog_entries.clone(),
|
||||||
));
|
));
|
||||||
tools.register_extension_tools(Arc::clone(&manager));
|
tools.register_extension_tools(Arc::clone(&manager));
|
||||||
|
|||||||
@@ -70,6 +70,7 @@ pub async fn extensions_list_handler(
|
|||||||
tools: ext.tools,
|
tools: ext.tools,
|
||||||
needs_setup: ext.needs_setup,
|
needs_setup: ext.needs_setup,
|
||||||
has_auth: ext.has_auth,
|
has_auth: ext.has_auth,
|
||||||
|
derived: ext.derived,
|
||||||
activation_status,
|
activation_status,
|
||||||
activation_error: ext.activation_error,
|
activation_error: ext.activation_error,
|
||||||
version: ext.version,
|
version: ext.version,
|
||||||
|
|||||||
@@ -2092,6 +2092,7 @@ async fn extensions_list_handler(
|
|||||||
tools: ext.tools,
|
tools: ext.tools,
|
||||||
needs_setup: ext.needs_setup,
|
needs_setup: ext.needs_setup,
|
||||||
has_auth: ext.has_auth,
|
has_auth: ext.has_auth,
|
||||||
|
derived: ext.derived,
|
||||||
activation_status,
|
activation_status,
|
||||||
activation_error: ext.activation_error,
|
activation_error: ext.activation_error,
|
||||||
version: ext.version,
|
version: ext.version,
|
||||||
@@ -2948,6 +2949,7 @@ mod tests {
|
|||||||
tools: Vec::new(),
|
tools: Vec::new(),
|
||||||
needs_setup: true,
|
needs_setup: true,
|
||||||
has_auth: false,
|
has_auth: false,
|
||||||
|
derived: false,
|
||||||
installed: true,
|
installed: true,
|
||||||
activation_error: None,
|
activation_error: None,
|
||||||
version: None,
|
version: None,
|
||||||
@@ -2985,6 +2987,7 @@ mod tests {
|
|||||||
tools: Vec::new(),
|
tools: Vec::new(),
|
||||||
needs_setup: true,
|
needs_setup: true,
|
||||||
has_auth: false,
|
has_auth: false,
|
||||||
|
derived: false,
|
||||||
installed: true,
|
installed: true,
|
||||||
activation_error: None,
|
activation_error: None,
|
||||||
version: None,
|
version: None,
|
||||||
@@ -4177,6 +4180,8 @@ mod tests {
|
|||||||
let ext_mgr = Arc::new(ExtensionManager::new(
|
let ext_mgr = Arc::new(ExtensionManager::new(
|
||||||
mcp_sm,
|
mcp_sm,
|
||||||
mcp_pm,
|
mcp_pm,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
secrets,
|
secrets,
|
||||||
tool_registry,
|
tool_registry,
|
||||||
None,
|
None,
|
||||||
@@ -4186,6 +4191,7 @@ mod tests {
|
|||||||
None,
|
None,
|
||||||
"test".to_string(),
|
"test".to_string(),
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
vec![],
|
vec![],
|
||||||
));
|
));
|
||||||
(ext_mgr, wasm_tools_dir, wasm_channels_dir)
|
(ext_mgr, wasm_tools_dir, wasm_channels_dir)
|
||||||
|
|||||||
@@ -1160,7 +1160,7 @@ function addToolCard(name) {
|
|||||||
|
|
||||||
const toolName = document.createElement('span');
|
const toolName = document.createElement('span');
|
||||||
toolName.className = 'activity-tool-name';
|
toolName.className = 'activity-tool-name';
|
||||||
toolName.textContent = name;
|
toolName.textContent = humanizeToolName(name);
|
||||||
|
|
||||||
const duration = document.createElement('span');
|
const duration = document.createElement('span');
|
||||||
duration.className = 'activity-tool-duration';
|
duration.className = 'activity-tool-duration';
|
||||||
@@ -1344,7 +1344,7 @@ function finalizeActivityGroup() {
|
|||||||
|
|
||||||
function humanizeToolName(rawName) {
|
function humanizeToolName(rawName) {
|
||||||
if (!rawName) return '';
|
if (!rawName) return '';
|
||||||
return String(rawName)
|
return stripDerivedCompanionToolPrefix(String(rawName))
|
||||||
.replace(/[_-]+/g, ' ')
|
.replace(/[_-]+/g, ' ')
|
||||||
.replace(/([a-z0-9])([A-Z])/g, '$1 $2')
|
.replace(/([a-z0-9])([A-Z])/g, '$1 $2')
|
||||||
.replace(/^tool([a-zA-Z])/, 'tool $1')
|
.replace(/^tool([a-zA-Z])/, 'tool $1')
|
||||||
@@ -1352,6 +1352,12 @@ function humanizeToolName(rawName) {
|
|||||||
.trim();
|
.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) {
|
function shouldShowChannelConnectedMessage(extensionName, success) {
|
||||||
if (!success || !extensionName) return false;
|
if (!success || !extensionName) return false;
|
||||||
return String(extensionName).toLowerCase().includes('telegram');
|
return String(extensionName).toLowerCase().includes('telegram');
|
||||||
@@ -1871,7 +1877,7 @@ function createToolCallsSummaryElement(toolCalls) {
|
|||||||
const icon = tc.has_error ? '\u2717' : '\u2713';
|
const icon = tc.has_error ? '\u2717' : '\u2713';
|
||||||
const nameSpan = document.createElement('span');
|
const nameSpan = document.createElement('span');
|
||||||
nameSpan.className = 'tool-call-name';
|
nameSpan.className = 'tool-call-name';
|
||||||
nameSpan.textContent = icon + ' ' + tc.name;
|
nameSpan.textContent = icon + ' ' + humanizeToolName(tc.name);
|
||||||
item.appendChild(nameSpan);
|
item.appendChild(nameSpan);
|
||||||
|
|
||||||
if (tc.result_preview) {
|
if (tc.result_preview) {
|
||||||
@@ -2906,7 +2912,12 @@ function renderExtensionCard(ext) {
|
|||||||
if (ext.tools && ext.tools.length > 0) {
|
if (ext.tools && ext.tools.length > 0) {
|
||||||
const tools = document.createElement('div');
|
const tools = document.createElement('div');
|
||||||
tools.className = 'ext-tools';
|
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);
|
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 —
|
// 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
|
// this means OAuth credentials resolve automatically (builtin/env) and the user
|
||||||
// just needs to complete the OAuth flow, not fill in a config form.
|
// 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');
|
const configBtn = document.createElement('button');
|
||||||
configBtn.className = 'btn-ext configure';
|
configBtn.className = 'btn-ext configure';
|
||||||
configBtn.textContent = ext.authenticated ? I18n.t('ext.reconfigure') : I18n.t('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');
|
if (!ext.derived) {
|
||||||
removeBtn.className = 'btn-ext remove';
|
const removeBtn = document.createElement('button');
|
||||||
removeBtn.textContent = I18n.t('ext.remove');
|
removeBtn.className = 'btn-ext remove';
|
||||||
removeBtn.addEventListener('click', () => removeExtension(ext.name));
|
removeBtn.textContent = I18n.t('ext.remove');
|
||||||
actions.appendChild(removeBtn);
|
removeBtn.addEventListener('click', () => removeExtension(ext.name));
|
||||||
|
actions.appendChild(removeBtn);
|
||||||
|
}
|
||||||
|
|
||||||
card.appendChild(actions);
|
card.appendChild(actions);
|
||||||
|
|
||||||
|
|||||||
@@ -344,6 +344,9 @@ pub struct ExtensionInfo {
|
|||||||
/// Whether this extension has an auth configuration (OAuth or manual token).
|
/// Whether this extension has an auth configuration (OAuth or manual token).
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub has_auth: bool,
|
pub has_auth: bool,
|
||||||
|
/// Whether this extension is derived from runtime/provider state.
|
||||||
|
#[serde(default)]
|
||||||
|
pub derived: bool,
|
||||||
/// WASM channel activation status.
|
/// WASM channel activation status.
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub activation_status: Option<ExtensionActivationStatus>,
|
pub activation_status: Option<ExtensionActivationStatus>,
|
||||||
|
|||||||
+325
-41
@@ -8,7 +8,7 @@ use std::sync::Arc;
|
|||||||
|
|
||||||
use clap::{Args, Subcommand};
|
use clap::{Args, Subcommand};
|
||||||
|
|
||||||
use crate::config::Config;
|
use crate::config::{Config, LlmConfig};
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
use crate::tools::mcp::{
|
use crate::tools::mcp::{
|
||||||
@@ -173,6 +173,13 @@ async fn add_server(args: McpAddArgs) -> anyhow::Result<()> {
|
|||||||
description,
|
description,
|
||||||
} = args;
|
} = 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 transport_lower = transport.to_lowercase();
|
||||||
|
|
||||||
let mut config = match transport_lower.as_str() {
|
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)
|
// Save (DB if available, else disk)
|
||||||
let db = connect_db().await;
|
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);
|
servers.upsert(config);
|
||||||
save_servers(db.as_deref(), &servers).await?;
|
save_servers(db.as_deref(), &servers).await?;
|
||||||
|
|
||||||
@@ -281,8 +288,15 @@ async fn add_server(args: McpAddArgs) -> anyhow::Result<()> {
|
|||||||
|
|
||||||
/// Remove an MCP server.
|
/// Remove an MCP server.
|
||||||
async fn remove_server(name: String) -> anyhow::Result<()> {
|
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 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) {
|
if !servers.remove(&name) {
|
||||||
anyhow::bail!("Server '{}' not found", name);
|
anyhow::bail!("Server '{}' not found", name);
|
||||||
}
|
}
|
||||||
@@ -298,7 +312,7 @@ async fn remove_server(name: String) -> anyhow::Result<()> {
|
|||||||
/// List configured MCP servers.
|
/// List configured MCP servers.
|
||||||
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||||
let db = connect_db().await;
|
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() {
|
if servers.servers.is_empty() {
|
||||||
println!();
|
println!();
|
||||||
@@ -404,12 +418,23 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
|||||||
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||||
// Get server config
|
// Get server config
|
||||||
let db = connect_db().await;
|
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
|
let server = servers
|
||||||
.get(&name)
|
.get(&name)
|
||||||
.cloned()
|
.cloned()
|
||||||
.ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?;
|
.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
|
// Initialize secrets store
|
||||||
let secrets = get_secrets_store().await?;
|
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<()> {
|
async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||||
// Get server config
|
// Get server config
|
||||||
let db = connect_db().await;
|
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
|
let server = servers
|
||||||
.get(&name)
|
.get(&name)
|
||||||
.cloned()
|
.cloned()
|
||||||
@@ -488,35 +513,66 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
|||||||
|
|
||||||
// Create client
|
// Create client
|
||||||
let session_manager = Arc::new(McpSessionManager::new());
|
let session_manager = Arc::new(McpSessionManager::new());
|
||||||
|
let (client, has_tokens) = if server.uses_runtime_auth_source() {
|
||||||
// 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 process_manager = Arc::new(McpProcessManager::new());
|
let process_manager = Arc::new(McpProcessManager::new());
|
||||||
create_client_from_config(
|
let llm = resolve_llm_for_cli(as_settings_store(db.as_deref())).await?;
|
||||||
server.clone(),
|
let nearai_session = crate::llm::create_session_manager(llm.session.clone()).await;
|
||||||
&session_manager,
|
(
|
||||||
&process_manager,
|
create_client_from_config(
|
||||||
None,
|
server.clone(),
|
||||||
"default",
|
&session_manager,
|
||||||
|
Some(nearai_session),
|
||||||
|
llm.nearai.api_key.clone(),
|
||||||
|
&process_manager,
|
||||||
|
None,
|
||||||
|
"default",
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| anyhow::anyhow!("{}", e))?,
|
||||||
|
false,
|
||||||
)
|
)
|
||||||
.await
|
} else {
|
||||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
// 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
|
// Test connection
|
||||||
@@ -581,8 +637,15 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
|||||||
|
|
||||||
/// Toggle server enabled/disabled state.
|
/// Toggle server enabled/disabled state.
|
||||||
async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> {
|
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 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
|
let server = servers
|
||||||
.get_mut(&name)
|
.get_mut(&name)
|
||||||
@@ -615,13 +678,30 @@ async fn connect_db() -> Option<Arc<dyn Database>> {
|
|||||||
crate::db::connect_from_config(&config.database).await.ok()
|
crate::db::connect_from_config(&config.database).await.ok()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Load MCP servers (DB if available, else disk).
|
/// Load only persisted MCP servers (DB if available, else disk).
|
||||||
async fn load_servers(db: Option<&dyn Database>) -> Result<McpServersFile, config::ConfigError> {
|
async fn load_persisted_servers(
|
||||||
if let Some(db) = db {
|
db: Option<&dyn Database>,
|
||||||
config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await
|
) -> Result<McpServersFile, config::ConfigError> {
|
||||||
|
Ok(if let Some(db) = db {
|
||||||
|
config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await?
|
||||||
} else {
|
} 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).
|
/// Save MCP servers (DB if available, else disk).
|
||||||
@@ -629,10 +709,15 @@ async fn save_servers(
|
|||||||
db: Option<&dyn Database>,
|
db: Option<&dyn Database>,
|
||||||
servers: &McpServersFile,
|
servers: &McpServersFile,
|
||||||
) -> Result<(), config::ConfigError> {
|
) -> 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 {
|
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 {
|
} 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
|
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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
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]
|
#[test]
|
||||||
fn test_mcp_command_parsing() {
|
fn test_mcp_command_parsing() {
|
||||||
// Just verify the command structure is valid
|
// Just verify the command structure is valid
|
||||||
@@ -701,4 +860,129 @@ mod tests {
|
|||||||
assert!(result.is_err());
|
assert!(result.is_err());
|
||||||
assert!(result.unwrap_err().contains("invalid env var format"));
|
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");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+472
-18
@@ -366,6 +366,8 @@ pub struct ExtensionManager {
|
|||||||
// MCP infrastructure
|
// MCP infrastructure
|
||||||
mcp_session_manager: Arc<McpSessionManager>,
|
mcp_session_manager: Arc<McpSessionManager>,
|
||||||
mcp_process_manager: Arc<crate::tools::mcp::process::McpProcessManager>,
|
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.
|
/// Active MCP clients keyed by server name.
|
||||||
mcp_clients: RwLock<HashMap<String, Arc<McpClient>>>,
|
mcp_clients: RwLock<HashMap<String, Arc<McpClient>>>,
|
||||||
|
|
||||||
@@ -389,6 +391,8 @@ pub struct ExtensionManager {
|
|||||||
user_id: String,
|
user_id: String,
|
||||||
/// Optional database store for DB-backed MCP config.
|
/// Optional database store for DB-backed MCP config.
|
||||||
store: Option<Arc<dyn crate::db::Database>>,
|
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.
|
/// Names of WASM channels that were successfully loaded at startup.
|
||||||
active_channel_names: RwLock<HashSet<String>>,
|
active_channel_names: RwLock<HashSet<String>>,
|
||||||
/// Installed channel-relay extensions (no on-disk artifact, tracked in memory).
|
/// Installed channel-relay extensions (no on-disk artifact, tracked in memory).
|
||||||
@@ -497,6 +501,8 @@ impl ExtensionManager {
|
|||||||
pub fn new(
|
pub fn new(
|
||||||
mcp_session_manager: Arc<McpSessionManager>,
|
mcp_session_manager: Arc<McpSessionManager>,
|
||||||
mcp_process_manager: Arc<crate::tools::mcp::process::McpProcessManager>,
|
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>,
|
secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||||
tool_registry: Arc<ToolRegistry>,
|
tool_registry: Arc<ToolRegistry>,
|
||||||
hooks: Option<Arc<HookRegistry>>,
|
hooks: Option<Arc<HookRegistry>>,
|
||||||
@@ -506,6 +512,7 @@ impl ExtensionManager {
|
|||||||
tunnel_url: Option<String>,
|
tunnel_url: Option<String>,
|
||||||
user_id: String,
|
user_id: String,
|
||||||
store: Option<Arc<dyn crate::db::Database>>,
|
store: Option<Arc<dyn crate::db::Database>>,
|
||||||
|
companion_mcp_server: Option<McpServerConfig>,
|
||||||
catalog_entries: Vec<RegistryEntry>,
|
catalog_entries: Vec<RegistryEntry>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
let registry = if catalog_entries.is_empty() {
|
let registry = if catalog_entries.is_empty() {
|
||||||
@@ -518,6 +525,8 @@ impl ExtensionManager {
|
|||||||
discovery: OnlineDiscovery::new(),
|
discovery: OnlineDiscovery::new(),
|
||||||
mcp_session_manager,
|
mcp_session_manager,
|
||||||
mcp_process_manager,
|
mcp_process_manager,
|
||||||
|
nearai_session_manager,
|
||||||
|
nearai_api_key,
|
||||||
mcp_clients: RwLock::new(HashMap::new()),
|
mcp_clients: RwLock::new(HashMap::new()),
|
||||||
wasm_tool_runtime,
|
wasm_tool_runtime,
|
||||||
wasm_tools_dir,
|
wasm_tools_dir,
|
||||||
@@ -531,6 +540,7 @@ impl ExtensionManager {
|
|||||||
tunnel_url,
|
tunnel_url,
|
||||||
user_id,
|
user_id,
|
||||||
store,
|
store,
|
||||||
|
companion_mcp_server,
|
||||||
active_channel_names: RwLock::new(HashSet::new()),
|
active_channel_names: RwLock::new(HashSet::new()),
|
||||||
installed_relay_extensions: RwLock::new(HashSet::new()),
|
installed_relay_extensions: RwLock::new(HashSet::new()),
|
||||||
activation_errors: RwLock::new(HashMap::new()),
|
activation_errors: RwLock::new(HashMap::new()),
|
||||||
@@ -1269,6 +1279,12 @@ impl ExtensionManager {
|
|||||||
tracing::info!(extension = %name, url = ?sanitized_url, kind = ?kind_hint, "Installing extension");
|
tracing::info!(extension = %name, url = ?sanitized_url, kind = ?kind_hint, "Installing extension");
|
||||||
Self::validate_extension_name(name)?;
|
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 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 {
|
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| {
|
return self.install_from_entry(&entry, user_id).await.map_err(|e| {
|
||||||
@@ -1344,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.
|
/// List extensions with their status.
|
||||||
///
|
///
|
||||||
/// When `include_available` is `true`, registry entries that are not yet
|
/// When `include_available` is `true`, registry entries that are not yet
|
||||||
@@ -1361,7 +1403,11 @@ impl ExtensionManager {
|
|||||||
match self.load_mcp_servers(user_id).await {
|
match self.load_mcp_servers(user_id).await {
|
||||||
Ok(servers) => {
|
Ok(servers) => {
|
||||||
for server in &servers.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 clients = self.mcp_clients.read().await;
|
||||||
let active = clients.contains_key(&server.name);
|
let active = clients.contains_key(&server.name);
|
||||||
|
|
||||||
@@ -1377,11 +1423,17 @@ impl ExtensionManager {
|
|||||||
Vec::new()
|
Vec::new()
|
||||||
};
|
};
|
||||||
|
|
||||||
let display_name = self
|
let display_name =
|
||||||
.registry
|
if crate::tools::mcp::config::is_nearai_companion_server_name(
|
||||||
.get_with_kind(&server.name, Some(ExtensionKind::McpServer))
|
&server.name,
|
||||||
.await
|
) {
|
||||||
.map(|e| e.display_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 {
|
extensions.push(InstalledExtension {
|
||||||
name: server.name.clone(),
|
name: server.name.clone(),
|
||||||
kind: ExtensionKind::McpServer,
|
kind: ExtensionKind::McpServer,
|
||||||
@@ -1392,7 +1444,10 @@ impl ExtensionManager {
|
|||||||
active,
|
active,
|
||||||
tools,
|
tools,
|
||||||
needs_setup: false,
|
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,
|
installed: true,
|
||||||
activation_error: None,
|
activation_error: None,
|
||||||
version: None,
|
version: None,
|
||||||
@@ -1444,6 +1499,7 @@ impl ExtensionManager {
|
|||||||
tools: if active { vec![name] } else { Vec::new() },
|
tools: if active { vec![name] } else { Vec::new() },
|
||||||
needs_setup: auth_state == ToolAuthState::NeedsSetup,
|
needs_setup: auth_state == ToolAuthState::NeedsSetup,
|
||||||
has_auth: auth_state != ToolAuthState::NoAuth,
|
has_auth: auth_state != ToolAuthState::NoAuth,
|
||||||
|
derived: false,
|
||||||
installed: true,
|
installed: true,
|
||||||
activation_error: None,
|
activation_error: None,
|
||||||
version,
|
version,
|
||||||
@@ -1500,6 +1556,7 @@ impl ExtensionManager {
|
|||||||
tools: Vec::new(),
|
tools: Vec::new(),
|
||||||
needs_setup: auth_state == ToolAuthState::NeedsSetup,
|
needs_setup: auth_state == ToolAuthState::NeedsSetup,
|
||||||
has_auth: auth_state != ToolAuthState::NoAuth,
|
has_auth: auth_state != ToolAuthState::NoAuth,
|
||||||
|
derived: false,
|
||||||
installed: true,
|
installed: true,
|
||||||
activation_error,
|
activation_error,
|
||||||
version,
|
version,
|
||||||
@@ -1538,6 +1595,7 @@ impl ExtensionManager {
|
|||||||
tools: Vec::new(),
|
tools: Vec::new(),
|
||||||
needs_setup: false,
|
needs_setup: false,
|
||||||
has_auth: true,
|
has_auth: true,
|
||||||
|
derived: false,
|
||||||
installed: true,
|
installed: true,
|
||||||
activation_error,
|
activation_error,
|
||||||
version: None,
|
version: None,
|
||||||
@@ -1572,6 +1630,7 @@ impl ExtensionManager {
|
|||||||
tools: Vec::new(),
|
tools: Vec::new(),
|
||||||
needs_setup: false,
|
needs_setup: false,
|
||||||
has_auth: false,
|
has_auth: false,
|
||||||
|
derived: false,
|
||||||
installed: false,
|
installed: false,
|
||||||
activation_error: None,
|
activation_error: None,
|
||||||
version: entry.version,
|
version: entry.version,
|
||||||
@@ -1602,6 +1661,12 @@ impl ExtensionManager {
|
|||||||
|
|
||||||
match kind {
|
match kind {
|
||||||
ExtensionKind::McpServer => {
|
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
|
// Unregister tools with this server's prefix
|
||||||
let tool_names: Vec<String> = self
|
let tool_names: Vec<String> = self
|
||||||
.tool_registry
|
.tool_registry
|
||||||
@@ -2059,10 +2124,39 @@ impl ExtensionManager {
|
|||||||
user_id: &str,
|
user_id: &str,
|
||||||
) -> Result<crate::tools::mcp::config::McpServersFile, crate::tools::mcp::config::ConfigError>
|
) -> Result<crate::tools::mcp::config::McpServersFile, crate::tools::mcp::config::ConfigError>
|
||||||
{
|
{
|
||||||
if let Some(ref store) = self.store {
|
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
|
crate::tools::mcp::config::load_mcp_servers_from_db(store.as_ref(), user_id).await?
|
||||||
} else {
|
} 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,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2624,6 +2718,20 @@ impl ExtensionManager {
|
|||||||
.await
|
.await
|
||||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
.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
|
// Check if already authenticated
|
||||||
if is_authenticated(&server, &self.secrets, user_id).await {
|
if is_authenticated(&server, &self.secrets, user_id).await {
|
||||||
return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer));
|
return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer));
|
||||||
@@ -3675,6 +3783,8 @@ impl ExtensionManager {
|
|||||||
let client = crate::tools::mcp::create_client_from_config(
|
let client = crate::tools::mcp::create_client_from_config(
|
||||||
server.clone(),
|
server.clone(),
|
||||||
&self.mcp_session_manager,
|
&self.mcp_session_manager,
|
||||||
|
self.nearai_session_manager.clone(),
|
||||||
|
self.nearai_api_key.clone(),
|
||||||
&self.mcp_process_manager,
|
&self.mcp_process_manager,
|
||||||
Some(Arc::clone(&self.secrets)),
|
Some(Arc::clone(&self.secrets)),
|
||||||
user_id,
|
user_id,
|
||||||
@@ -5295,6 +5405,12 @@ impl ExtensionManager {
|
|||||||
.get_mcp_server(name, user_id)
|
.get_mcp_server(name, user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
.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();
|
let mut names = std::collections::HashSet::new();
|
||||||
names.insert(server.token_secret_name());
|
names.insert(server.token_secret_name());
|
||||||
(names, Vec::new())
|
(names, Vec::new())
|
||||||
@@ -6299,6 +6415,8 @@ mod tests {
|
|||||||
tools_dir: std::path::PathBuf,
|
tools_dir: std::path::PathBuf,
|
||||||
channels_dir: std::path::PathBuf,
|
channels_dir: std::path::PathBuf,
|
||||||
store: Option<Arc<dyn crate::db::Database>>,
|
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 {
|
) -> crate::extensions::manager::ExtensionManager {
|
||||||
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
|
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
|
||||||
use crate::tools::mcp::process::McpProcessManager;
|
use crate::tools::mcp::process::McpProcessManager;
|
||||||
@@ -6317,15 +6435,18 @@ mod tests {
|
|||||||
crate::extensions::manager::ExtensionManager::new(
|
crate::extensions::manager::ExtensionManager::new(
|
||||||
mcp,
|
mcp,
|
||||||
Arc::new(McpProcessManager::new()),
|
Arc::new(McpProcessManager::new()),
|
||||||
|
nearai_session_manager,
|
||||||
|
None,
|
||||||
secrets,
|
secrets,
|
||||||
tools,
|
tools,
|
||||||
None, // hooks
|
None, // hooks
|
||||||
wasm_runtime,
|
wasm_runtime,
|
||||||
tools_dir,
|
tools_dir,
|
||||||
channels_dir,
|
channels_dir,
|
||||||
None, // tunnel_url
|
None, // tunnel_url
|
||||||
"test".to_string(), // user_id
|
"test".to_string(),
|
||||||
store,
|
store,
|
||||||
|
companion_mcp_server,
|
||||||
vec![],
|
vec![],
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -6334,7 +6455,7 @@ mod tests {
|
|||||||
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
|
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
|
||||||
tools_dir: std::path::PathBuf,
|
tools_dir: std::path::PathBuf,
|
||||||
) -> crate::extensions::manager::ExtensionManager {
|
) -> 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(
|
fn write_test_tool(
|
||||||
@@ -6395,6 +6516,8 @@ mod tests {
|
|||||||
dir.path().join("tools"),
|
dir.path().join("tools"),
|
||||||
dir.path().join("channels"),
|
dir.path().join("channels"),
|
||||||
Some(Arc::clone(&store)),
|
Some(Arc::clone(&store)),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
let field = crate::tools::wasm::ToolFieldSetupSchema {
|
let field = crate::tools::wasm::ToolFieldSetupSchema {
|
||||||
name: "provider".to_string(),
|
name: "provider".to_string(),
|
||||||
@@ -6436,8 +6559,14 @@ mod tests {
|
|||||||
);
|
);
|
||||||
let channels_dir = dir.path().join("channels");
|
let channels_dir = dir.path().join("channels");
|
||||||
|
|
||||||
let mgr =
|
let mgr = make_test_manager_with_dirs(
|
||||||
make_test_manager_with_dirs(None, tools_dir, channels_dir, Some(Arc::clone(&store)));
|
None,
|
||||||
|
tools_dir,
|
||||||
|
channels_dir,
|
||||||
|
Some(Arc::clone(&store)),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
);
|
||||||
let mut fields = std::collections::HashMap::new();
|
let mut fields = std::collections::HashMap::new();
|
||||||
fields.insert("llm_backend".to_string(), "openai".to_string());
|
fields.insert("llm_backend".to_string(), "openai".to_string());
|
||||||
|
|
||||||
@@ -6489,8 +6618,14 @@ mod tests {
|
|||||||
);
|
);
|
||||||
let channels_dir = dir.path().join("channels");
|
let channels_dir = dir.path().join("channels");
|
||||||
|
|
||||||
let mgr =
|
let mgr = make_test_manager_with_dirs(
|
||||||
make_test_manager_with_dirs(None, tools_dir, channels_dir, Some(Arc::clone(&store)));
|
None,
|
||||||
|
tools_dir,
|
||||||
|
channels_dir,
|
||||||
|
Some(Arc::clone(&store)),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
);
|
||||||
let mut fields = std::collections::HashMap::new();
|
let mut fields = std::collections::HashMap::new();
|
||||||
fields.insert("session".to_string(), "overwrite".to_string());
|
fields.insert("session".to_string(), "overwrite".to_string());
|
||||||
|
|
||||||
@@ -6563,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]
|
#[test]
|
||||||
fn test_capabilities_files_also_separate() {
|
fn test_capabilities_files_also_separate() {
|
||||||
// capabilities.json files for tools and channels should also be separate.
|
// capabilities.json files for tools and channels should also be separate.
|
||||||
@@ -6685,6 +7086,8 @@ mod tests {
|
|||||||
ExtensionManager::new(
|
ExtensionManager::new(
|
||||||
Arc::new(McpSessionManager::new()),
|
Arc::new(McpSessionManager::new()),
|
||||||
Arc::new(McpProcessManager::new()),
|
Arc::new(McpProcessManager::new()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||||
Arc::new(ToolRegistry::new()),
|
Arc::new(ToolRegistry::new()),
|
||||||
None,
|
None,
|
||||||
@@ -6694,6 +7097,7 @@ mod tests {
|
|||||||
None,
|
None,
|
||||||
"test".to_string(),
|
"test".to_string(),
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
Vec::new(),
|
Vec::new(),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -6835,6 +7239,8 @@ mod tests {
|
|||||||
ExtensionManager::new(
|
ExtensionManager::new(
|
||||||
Arc::new(McpSessionManager::new()),
|
Arc::new(McpSessionManager::new()),
|
||||||
Arc::new(McpProcessManager::new()),
|
Arc::new(McpProcessManager::new()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||||
Arc::new(ToolRegistry::new()),
|
Arc::new(ToolRegistry::new()),
|
||||||
None,
|
None,
|
||||||
@@ -6844,6 +7250,7 @@ mod tests {
|
|||||||
None,
|
None,
|
||||||
"test".to_string(),
|
"test".to_string(),
|
||||||
Some(db),
|
Some(db),
|
||||||
|
None,
|
||||||
Vec::new(),
|
Vec::new(),
|
||||||
)
|
)
|
||||||
};
|
};
|
||||||
@@ -7097,6 +7504,8 @@ mod tests {
|
|||||||
let manager = ExtensionManager::new(
|
let manager = ExtensionManager::new(
|
||||||
Arc::new(McpSessionManager::new()),
|
Arc::new(McpSessionManager::new()),
|
||||||
Arc::new(McpProcessManager::new()),
|
Arc::new(McpProcessManager::new()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||||
Arc::new(ToolRegistry::new()),
|
Arc::new(ToolRegistry::new()),
|
||||||
None,
|
None,
|
||||||
@@ -7106,6 +7515,7 @@ mod tests {
|
|||||||
None,
|
None,
|
||||||
"test".to_string(),
|
"test".to_string(),
|
||||||
Some(db.clone() as Arc<dyn crate::db::Database>),
|
Some(db.clone() as Arc<dyn crate::db::Database>),
|
||||||
|
None,
|
||||||
Vec::new(),
|
Vec::new(),
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -7590,7 +8000,8 @@ mod tests {
|
|||||||
let dir = tempfile::tempdir().expect("temp dir");
|
let dir = tempfile::tempdir().expect("temp dir");
|
||||||
let tools_dir = dir.path().join("tools");
|
let tools_dir = dir.path().join("tools");
|
||||||
let channels_dir = dir.path().join("channels");
|
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 wasm_path = channels_dir.join("telegram.wasm");
|
||||||
let cap_path = channels_dir.join("telegram.capabilities.json");
|
let cap_path = channels_dir.join("telegram.capabilities.json");
|
||||||
@@ -7734,6 +8145,8 @@ mod tests {
|
|||||||
ExtensionManager::new(
|
ExtensionManager::new(
|
||||||
mcp,
|
mcp,
|
||||||
Arc::new(McpProcessManager::new()),
|
Arc::new(McpProcessManager::new()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
secrets,
|
secrets,
|
||||||
tools,
|
tools,
|
||||||
None,
|
None,
|
||||||
@@ -7743,6 +8156,7 @@ mod tests {
|
|||||||
tunnel_url,
|
tunnel_url,
|
||||||
"test".to_string(),
|
"test".to_string(),
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
vec![],
|
vec![],
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -8025,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]
|
#[tokio::test]
|
||||||
async fn test_auth_is_read_only_for_wasm_channel() {
|
async fn test_auth_is_read_only_for_wasm_channel() {
|
||||||
// Regression: auth() must be a pure status check — it must not store
|
// 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).
|
/// Whether this extension has an auth configuration (OAuth or manual token).
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub has_auth: bool,
|
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).
|
/// Whether this extension is installed locally (false = available in registry but not installed).
|
||||||
#[serde(default = "default_true")]
|
#[serde(default = "default_true")]
|
||||||
pub installed: bool,
|
pub installed: bool,
|
||||||
@@ -936,6 +940,7 @@ mod tests {
|
|||||||
assert!(ext.installed, "installed should default to true");
|
assert!(ext.installed, "installed should default to true");
|
||||||
assert!(!ext.needs_setup, "needs_setup should default to false");
|
assert!(!ext.needs_setup, "needs_setup should default to false");
|
||||||
assert!(!ext.has_auth);
|
assert!(!ext.has_auth);
|
||||||
|
assert!(!ext.derived);
|
||||||
assert!(ext.tools.is_empty());
|
assert!(ext.tools.is_empty());
|
||||||
assert!(ext.display_name.is_none());
|
assert!(ext.display_name.is_none());
|
||||||
assert!(ext.description.is_none());
|
assert!(ext.description.is_none());
|
||||||
@@ -956,6 +961,7 @@ mod tests {
|
|||||||
tools: vec!["send_email".to_string(), "read_inbox".to_string()],
|
tools: vec!["send_email".to_string(), "read_inbox".to_string()],
|
||||||
needs_setup: true,
|
needs_setup: true,
|
||||||
has_auth: true,
|
has_auth: true,
|
||||||
|
derived: true,
|
||||||
installed: false,
|
installed: false,
|
||||||
activation_error: Some("token expired".to_string()),
|
activation_error: Some("token expired".to_string()),
|
||||||
version: None,
|
version: None,
|
||||||
@@ -965,6 +971,7 @@ mod tests {
|
|||||||
assert_eq!(json["description"], "Read and send emails");
|
assert_eq!(json["description"], "Read and send emails");
|
||||||
assert_eq!(json["url"], "https://gmail.example.com");
|
assert_eq!(json["url"], "https://gmail.example.com");
|
||||||
assert_eq!(json["needs_setup"], true);
|
assert_eq!(json["needs_setup"], true);
|
||||||
|
assert_eq!(json["derived"], true);
|
||||||
assert_eq!(json["installed"], false);
|
assert_eq!(json["installed"], false);
|
||||||
assert_eq!(json["activation_error"], "token expired");
|
assert_eq!(json["activation_error"], "token expired");
|
||||||
|
|
||||||
@@ -972,6 +979,7 @@ mod tests {
|
|||||||
assert_eq!(back.name, "gmail");
|
assert_eq!(back.name, "gmail");
|
||||||
assert_eq!(back.tools.len(), 2);
|
assert_eq!(back.tools.len(), 2);
|
||||||
assert!(back.needs_setup);
|
assert!(back.needs_setup);
|
||||||
|
assert!(back.derived);
|
||||||
assert!(!back.installed);
|
assert!(!back.installed);
|
||||||
assert_eq!(back.activation_error.as_deref(), Some("token expired"));
|
assert_eq!(back.activation_error.as_deref(), Some("token expired"));
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ pub mod failover;
|
|||||||
pub mod gemini_oauth;
|
pub mod gemini_oauth;
|
||||||
mod github_copilot;
|
mod github_copilot;
|
||||||
pub(crate) mod github_copilot_auth;
|
pub(crate) mod github_copilot_auth;
|
||||||
|
pub mod nearai_auth;
|
||||||
mod nearai_chat;
|
mod nearai_chat;
|
||||||
pub mod oauth_helpers;
|
pub mod oauth_helpers;
|
||||||
pub mod openai_codex_provider;
|
pub mod openai_codex_provider;
|
||||||
@@ -53,6 +54,7 @@ pub use config::{
|
|||||||
pub use error::LlmError;
|
pub use error::LlmError;
|
||||||
pub use failover::{CooldownConfig, FailoverProvider};
|
pub use failover::{CooldownConfig, FailoverProvider};
|
||||||
pub use gemini_oauth::GeminiOauthProvider;
|
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 nearai_chat::{DEFAULT_MODEL, ModelInfo, NearAiChatProvider, default_models};
|
||||||
pub use openai_codex_provider::OpenAiCodexProvider;
|
pub use openai_codex_provider::OpenAiCodexProvider;
|
||||||
pub use openai_codex_session::{OpenAiCodexSession, OpenAiCodexSessionManager};
|
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"));
|
||||||
|
}
|
||||||
|
}
|
||||||
+1
-30
@@ -173,36 +173,7 @@ impl NearAiChatProvider {
|
|||||||
/// The env var fallback (#3) only triggers after `ensure_authenticated()`
|
/// The env var fallback (#3) only triggers after `ensure_authenticated()`
|
||||||
/// runs, because `api_key_login()` sets the env var but not a session token.
|
/// runs, because `api_key_login()` sets the env var but not a session token.
|
||||||
async fn resolve_bearer_token(&self) -> Result<String, LlmError> {
|
async fn resolve_bearer_token(&self) -> Result<String, LlmError> {
|
||||||
// 1. Config-level API key takes priority
|
crate::llm::resolve_nearai_bearer_token(self.config.api_key.as_ref(), &self.session).await
|
||||||
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(),
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Send a single request to the chat completions API.
|
/// Send a single request to the chat completions API.
|
||||||
|
|||||||
@@ -139,6 +139,8 @@ mod tests {
|
|||||||
Arc::new(ExtensionManager::new(
|
Arc::new(ExtensionManager::new(
|
||||||
Arc::new(McpSessionManager::new()),
|
Arc::new(McpSessionManager::new()),
|
||||||
Arc::new(McpProcessManager::new()),
|
Arc::new(McpProcessManager::new()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
secrets,
|
secrets,
|
||||||
tools,
|
tools,
|
||||||
Some(Arc::new(HookRegistry::default())),
|
Some(Arc::new(HookRegistry::default())),
|
||||||
@@ -148,6 +150,7 @@ mod tests {
|
|||||||
None,
|
None,
|
||||||
owner_id.to_string(),
|
owner_id.to_string(),
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
Vec::new(),
|
Vec::new(),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -800,6 +800,8 @@ mod tests {
|
|||||||
Arc::new(ExtensionManager::new(
|
Arc::new(ExtensionManager::new(
|
||||||
Arc::new(McpSessionManager::new()),
|
Arc::new(McpSessionManager::new()),
|
||||||
Arc::new(crate::tools::mcp::process::McpProcessManager::new()),
|
Arc::new(crate::tools::mcp::process::McpProcessManager::new()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||||
Arc::new(ToolRegistry::new()),
|
Arc::new(ToolRegistry::new()),
|
||||||
None,
|
None,
|
||||||
@@ -809,6 +811,7 @@ mod tests {
|
|||||||
None,
|
None,
|
||||||
"test".to_string(),
|
"test".to_string(),
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
Vec::new(),
|
Vec::new(),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|||||||
+311
-8
@@ -8,12 +8,13 @@ use std::sync::Arc;
|
|||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use secrecy::SecretString;
|
||||||
use tokio::sync::RwLock;
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
use crate::context::JobContext;
|
use crate::context::JobContext;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
use crate::tools::mcp::auth::refresh_access_token;
|
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::http_transport::HttpMcpTransport;
|
||||||
use crate::tools::mcp::protocol::{
|
use crate::tools::mcp::protocol::{
|
||||||
CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool,
|
CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool,
|
||||||
@@ -46,6 +47,13 @@ pub struct McpClient {
|
|||||||
/// Session manager (shared across clients).
|
/// Session manager (shared across clients).
|
||||||
session_manager: Option<Arc<McpSessionManager>>,
|
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 store for retrieving access tokens.
|
||||||
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
|
|
||||||
@@ -80,6 +88,8 @@ impl McpClient {
|
|||||||
next_id: AtomicU64::new(1),
|
next_id: AtomicU64::new(1),
|
||||||
tools_cache: RwLock::new(None),
|
tools_cache: RwLock::new(None),
|
||||||
session_manager: None,
|
session_manager: None,
|
||||||
|
nearai_session_manager: None,
|
||||||
|
nearai_api_key: None,
|
||||||
secrets: None,
|
secrets: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
server_config: None,
|
server_config: None,
|
||||||
@@ -103,6 +113,8 @@ impl McpClient {
|
|||||||
next_id: AtomicU64::new(1),
|
next_id: AtomicU64::new(1),
|
||||||
tools_cache: RwLock::new(None),
|
tools_cache: RwLock::new(None),
|
||||||
session_manager: None,
|
session_manager: None,
|
||||||
|
nearai_session_manager: None,
|
||||||
|
nearai_api_key: None,
|
||||||
secrets: None,
|
secrets: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
server_config: None,
|
server_config: None,
|
||||||
@@ -123,6 +135,9 @@ impl McpClient {
|
|||||||
/// the transport with session tracking.
|
/// the transport with session tracking.
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
pub fn new_with_config(config: McpServerConfig) -> Result<Self, ToolError> {
|
pub fn new_with_config(config: McpServerConfig) -> Result<Self, ToolError> {
|
||||||
|
config
|
||||||
|
.validate()
|
||||||
|
.map_err(|e| ToolError::InvalidParameters(e.to_string()))?;
|
||||||
if !matches!(
|
if !matches!(
|
||||||
config.effective_transport(),
|
config.effective_transport(),
|
||||||
crate::tools::mcp::config::EffectiveTransport::Http
|
crate::tools::mcp::config::EffectiveTransport::Http
|
||||||
@@ -144,6 +159,8 @@ impl McpClient {
|
|||||||
next_id: AtomicU64::new(1),
|
next_id: AtomicU64::new(1),
|
||||||
tools_cache: RwLock::new(None),
|
tools_cache: RwLock::new(None),
|
||||||
session_manager: None,
|
session_manager: None,
|
||||||
|
nearai_session_manager: None,
|
||||||
|
nearai_api_key: None,
|
||||||
secrets: None,
|
secrets: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
custom_headers: config.headers.clone(),
|
custom_headers: config.headers.clone(),
|
||||||
@@ -175,6 +192,8 @@ impl McpClient {
|
|||||||
next_id: AtomicU64::new(1),
|
next_id: AtomicU64::new(1),
|
||||||
tools_cache: RwLock::new(None),
|
tools_cache: RwLock::new(None),
|
||||||
session_manager: Some(session_manager),
|
session_manager: Some(session_manager),
|
||||||
|
nearai_session_manager: None,
|
||||||
|
nearai_api_key: None,
|
||||||
secrets: Some(secrets),
|
secrets: Some(secrets),
|
||||||
user_id: user_id.into(),
|
user_id: user_id.into(),
|
||||||
server_config: Some(config),
|
server_config: Some(config),
|
||||||
@@ -211,6 +230,8 @@ impl McpClient {
|
|||||||
next_id: AtomicU64::new(1),
|
next_id: AtomicU64::new(1),
|
||||||
tools_cache: RwLock::new(None),
|
tools_cache: RwLock::new(None),
|
||||||
session_manager,
|
session_manager,
|
||||||
|
nearai_session_manager: None,
|
||||||
|
nearai_api_key: None,
|
||||||
secrets,
|
secrets,
|
||||||
user_id: user_id.into(),
|
user_id: user_id.into(),
|
||||||
server_config,
|
server_config,
|
||||||
@@ -232,6 +253,21 @@ impl McpClient {
|
|||||||
self
|
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.
|
/// Get the server name.
|
||||||
pub fn server_name(&self) -> &str {
|
pub fn server_name(&self) -> &str {
|
||||||
&self.server_name
|
&self.server_name
|
||||||
@@ -266,6 +302,9 @@ impl McpClient {
|
|||||||
let Some(ref config) = self.server_config else {
|
let Some(ref config) = self.server_config else {
|
||||||
return Ok(None);
|
return Ok(None);
|
||||||
};
|
};
|
||||||
|
if config.uses_runtime_auth_source() {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
match secrets
|
match secrets
|
||||||
.get_decrypted(&self.user_id, &config.token_secret_name())
|
.get_decrypted(&self.user_id, &config.token_secret_name())
|
||||||
.await
|
.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).
|
/// Build the headers map for a request (auth, session-id, custom headers).
|
||||||
///
|
///
|
||||||
/// Custom headers are applied first. OAuth token injection is skipped if the
|
/// Custom headers are applied first. OAuth token injection is skipped if the
|
||||||
@@ -292,6 +361,9 @@ impl McpClient {
|
|||||||
.custom_headers
|
.custom_headers
|
||||||
.keys()
|
.keys()
|
||||||
.any(|k| k.eq_ignore_ascii_case("authorization"));
|
.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? {
|
if !has_custom_auth && let Some(token) = self.get_access_token().await? {
|
||||||
let trimmed = token.trim();
|
let trimmed = token.trim();
|
||||||
if !trimmed.is_empty() {
|
if !trimmed.is_empty() {
|
||||||
@@ -512,13 +584,12 @@ impl McpClient {
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
|
|
||||||
response
|
let raw_result = response
|
||||||
.result
|
.result
|
||||||
.ok_or_else(|| ToolError::ExternalService("No result in MCP response".to_string()))
|
.ok_or_else(|| ToolError::ExternalService("No result in MCP response".to_string()))?;
|
||||||
.and_then(|r| {
|
|
||||||
serde_json::from_value(r)
|
serde_json::from_value(raw_result)
|
||||||
.map_err(|e| ToolError::ExternalService(format!("Invalid tool result: {}", e)))
|
.map_err(|e| ToolError::ExternalService(format!("Invalid tool result: {}", e)))
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Clear the tools cache.
|
/// Clear the tools cache.
|
||||||
@@ -565,6 +636,8 @@ impl Clone for McpClient {
|
|||||||
next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)),
|
next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)),
|
||||||
tools_cache: RwLock::new(None),
|
tools_cache: RwLock::new(None),
|
||||||
session_manager: self.session_manager.clone(),
|
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(),
|
secrets: self.secrets.clone(),
|
||||||
user_id: self.user_id.clone(),
|
user_id: self.user_id.clone(),
|
||||||
server_config: self.server_config.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
|
// Strip top-level null values before forwarding — LLMs often emit
|
||||||
// `"field": null` for optional params, but many MCP servers reject
|
// `"field": null` for optional params, but many MCP servers reject
|
||||||
// explicit nulls for fields that should simply be absent.
|
// 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 result = self.client.call_tool(&self.tool.name, params).await?;
|
||||||
let content: String = result
|
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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -823,6 +921,138 @@ mod tests {
|
|||||||
assert!(client.has_session_manager());
|
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]
|
#[test]
|
||||||
fn test_next_request_id_monotonically_increasing() {
|
fn test_next_request_id_monotonically_increasing() {
|
||||||
let client = McpClient::new("http://localhost:1234");
|
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 ---
|
// --- Issue 13: McpToolWrapper unit tests ---
|
||||||
|
|
||||||
fn make_test_mcp_tool(destructive: bool) -> McpTool {
|
fn make_test_mcp_tool(destructive: bool) -> McpTool {
|
||||||
@@ -1434,4 +1678,63 @@ mod tests {
|
|||||||
"Token must be trimmed before use in Authorization header"
|
"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")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub oauth: Option<OAuthConfig>,
|
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.
|
/// Whether this server is enabled.
|
||||||
#[serde(default = "default_true")]
|
#[serde(default = "default_true")]
|
||||||
pub enabled: bool,
|
pub enabled: bool,
|
||||||
@@ -60,6 +70,14 @@ pub struct McpServerConfig {
|
|||||||
pub description: Option<String>,
|
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 {
|
fn default_true() -> bool {
|
||||||
true
|
true
|
||||||
}
|
}
|
||||||
@@ -73,6 +91,7 @@ impl McpServerConfig {
|
|||||||
transport: None,
|
transport: None,
|
||||||
headers: HashMap::new(),
|
headers: HashMap::new(),
|
||||||
oauth: None,
|
oauth: None,
|
||||||
|
auth_source: None,
|
||||||
enabled: true,
|
enabled: true,
|
||||||
description: None,
|
description: None,
|
||||||
}
|
}
|
||||||
@@ -95,6 +114,7 @@ impl McpServerConfig {
|
|||||||
}),
|
}),
|
||||||
headers: HashMap::new(),
|
headers: HashMap::new(),
|
||||||
oauth: None,
|
oauth: None,
|
||||||
|
auth_source: None,
|
||||||
enabled: true,
|
enabled: true,
|
||||||
description: None,
|
description: None,
|
||||||
}
|
}
|
||||||
@@ -110,6 +130,7 @@ impl McpServerConfig {
|
|||||||
}),
|
}),
|
||||||
headers: HashMap::new(),
|
headers: HashMap::new(),
|
||||||
oauth: None,
|
oauth: None,
|
||||||
|
auth_source: None,
|
||||||
enabled: true,
|
enabled: true,
|
||||||
description: None,
|
description: None,
|
||||||
}
|
}
|
||||||
@@ -121,6 +142,12 @@ impl McpServerConfig {
|
|||||||
self
|
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.
|
/// Set description.
|
||||||
pub fn with_description(mut self, description: impl Into<String>) -> Self {
|
pub fn with_description(mut self, description: impl Into<String>) -> Self {
|
||||||
self.description = Some(description.into());
|
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() {
|
match self.effective_transport() {
|
||||||
EffectiveTransport::Http => {
|
EffectiveTransport::Http => {
|
||||||
if self.url.is_empty() {
|
if self.url.is_empty() {
|
||||||
@@ -222,6 +258,11 @@ impl McpServerConfig {
|
|||||||
.any(|k| k.eq_ignore_ascii_case("authorization"))
|
.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.
|
/// Check if this server requires authentication.
|
||||||
///
|
///
|
||||||
/// Returns true if OAuth is pre-configured OR if this is a remote HTTPS server
|
/// Returns true if OAuth is pre-configured OR if this is a remote HTTPS server
|
||||||
@@ -234,7 +275,7 @@ impl McpServerConfig {
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
if self.oauth.is_some() {
|
if self.oauth.is_some() || self.uses_runtime_auth_source() {
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
// Remote HTTPS servers need auth handling (DCR, token refresh, 401 detection).
|
// 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.
|
/// OAuth 2.1 configuration for an MCP server.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct OAuthConfig {
|
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.
|
/// Remove a server by name.
|
||||||
pub fn remove(&mut self, name: &str) -> bool {
|
pub fn remove(&mut self, name: &str) -> bool {
|
||||||
let len_before = self.servers.len();
|
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 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
|
// Validate every server on load so corrupted configs are caught early
|
||||||
for server in &config.servers {
|
for server in &config.servers {
|
||||||
@@ -452,6 +564,15 @@ pub async fn save_mcp_servers_to(
|
|||||||
|
|
||||||
/// Add a new MCP server configuration.
|
/// Add a new MCP server configuration.
|
||||||
pub async fn add_mcp_server(config: McpServerConfig) -> Result<(), ConfigError> {
|
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()?;
|
config.validate()?;
|
||||||
|
|
||||||
let mut servers = load_mcp_servers().await?;
|
let mut servers = load_mcp_servers().await?;
|
||||||
@@ -499,7 +620,8 @@ pub async fn load_mcp_servers_from_db(
|
|||||||
) -> Result<McpServersFile, ConfigError> {
|
) -> Result<McpServersFile, ConfigError> {
|
||||||
match store.get_setting(user_id, "mcp_servers").await {
|
match store.get_setting(user_id, "mcp_servers").await {
|
||||||
Ok(Some(value)) => {
|
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
|
// Validate every server on load so corrupted DB configs are caught early
|
||||||
for server in &config.servers {
|
for server in &config.servers {
|
||||||
server.validate().map_err(|e| ConfigError::InvalidConfig {
|
server.validate().map_err(|e| ConfigError::InvalidConfig {
|
||||||
@@ -542,6 +664,15 @@ pub async fn add_mcp_server_db(
|
|||||||
user_id: &str,
|
user_id: &str,
|
||||||
config: McpServerConfig,
|
config: McpServerConfig,
|
||||||
) -> Result<(), ConfigError> {
|
) -> 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()?;
|
config.validate()?;
|
||||||
|
|
||||||
let mut servers = load_mcp_servers_from_db(store, user_id).await?;
|
let mut servers = load_mcp_servers_from_db(store, user_id).await?;
|
||||||
@@ -718,6 +849,69 @@ mod tests {
|
|||||||
assert!(config.servers.is_empty());
|
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]
|
#[tokio::test]
|
||||||
async fn test_load_rejects_corrupted_headers() {
|
async fn test_load_rejects_corrupted_headers() {
|
||||||
let dir = tempdir().unwrap();
|
let dir = tempdir().unwrap();
|
||||||
@@ -763,6 +957,20 @@ mod tests {
|
|||||||
assert!(config.requires_auth());
|
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]
|
#[test]
|
||||||
fn test_requires_auth_remote_https_without_oauth() {
|
fn test_requires_auth_remote_https_without_oauth() {
|
||||||
// Remote HTTPS servers need auth even without pre-configured OAuth (DCR)
|
// Remote HTTPS servers need auth even without pre-configured OAuth (DCR)
|
||||||
|
|||||||
@@ -21,6 +21,8 @@ pub enum McpFactoryError {
|
|||||||
UnixNotSupported { name: String },
|
UnixNotSupported { name: String },
|
||||||
#[error("Invalid configuration for MCP server '{name}': {reason}")]
|
#[error("Invalid configuration for MCP server '{name}': {reason}")]
|
||||||
InvalidConfig { name: String, reason: String },
|
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
|
/// Create an `McpClient` from a server configuration, dispatching on the
|
||||||
@@ -28,6 +30,8 @@ pub enum McpFactoryError {
|
|||||||
pub async fn create_client_from_config(
|
pub async fn create_client_from_config(
|
||||||
server: McpServerConfig,
|
server: McpServerConfig,
|
||||||
session_manager: &Arc<McpSessionManager>,
|
session_manager: &Arc<McpSessionManager>,
|
||||||
|
nearai_session_manager: Option<Arc<crate::llm::SessionManager>>,
|
||||||
|
nearai_api_key: Option<secrecy::SecretString>,
|
||||||
process_manager: &Arc<McpProcessManager>,
|
process_manager: &Arc<McpProcessManager>,
|
||||||
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
user_id: &str,
|
user_id: &str,
|
||||||
@@ -79,7 +83,31 @@ pub async fn create_client_from_config(
|
|||||||
Err(McpFactoryError::UnixNotSupported { name: server_name })
|
Err(McpFactoryError::UnixNotSupported { name: server_name })
|
||||||
}
|
}
|
||||||
EffectiveTransport::Http => {
|
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 {
|
if let Some(ref secrets) = secrets {
|
||||||
let has_tokens =
|
let has_tokens =
|
||||||
crate::tools::mcp::is_authenticated(&server, secrets, user_id).await;
|
crate::tools::mcp::is_authenticated(&server, secrets, user_id).await;
|
||||||
@@ -127,6 +155,8 @@ mod tests {
|
|||||||
let client = create_client_from_config(
|
let client = create_client_from_config(
|
||||||
server,
|
server,
|
||||||
&session_manager,
|
&session_manager,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
&process_manager,
|
&process_manager,
|
||||||
None,
|
None,
|
||||||
"test-user",
|
"test-user",
|
||||||
@@ -185,6 +215,8 @@ mod tests {
|
|||||||
let client = create_client_from_config(
|
let client = create_client_from_config(
|
||||||
server,
|
server,
|
||||||
&session_manager,
|
&session_manager,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
&process_manager,
|
&process_manager,
|
||||||
None,
|
None,
|
||||||
"test-user",
|
"test-user",
|
||||||
|
|||||||
@@ -290,6 +290,8 @@ mod tests {
|
|||||||
Arc::new(ExtensionManager::new(
|
Arc::new(ExtensionManager::new(
|
||||||
Arc::new(McpSessionManager::new()),
|
Arc::new(McpSessionManager::new()),
|
||||||
Arc::new(McpProcessManager::new()),
|
Arc::new(McpProcessManager::new()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
secrets,
|
secrets,
|
||||||
tools,
|
tools,
|
||||||
None,
|
None,
|
||||||
@@ -299,6 +301,7 @@ mod tests {
|
|||||||
None,
|
None,
|
||||||
owner_id.to_string(),
|
owner_id.to_string(),
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
Vec::new(),
|
Vec::new(),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -203,6 +203,8 @@ async fn extension_manager_with_process_manager_constructs() {
|
|||||||
let manager = ExtensionManager::new(
|
let manager = ExtensionManager::new(
|
||||||
Arc::new(McpSessionManager::new()),
|
Arc::new(McpSessionManager::new()),
|
||||||
Arc::new(McpProcessManager::new()),
|
Arc::new(McpProcessManager::new()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
secrets,
|
secrets,
|
||||||
tools,
|
tools,
|
||||||
None,
|
None,
|
||||||
@@ -212,6 +214,7 @@ async fn extension_manager_with_process_manager_constructs() {
|
|||||||
None,
|
None,
|
||||||
"test".to_string(),
|
"test".to_string(),
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
Vec::new(),
|
Vec::new(),
|
||||||
);
|
);
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user