Some checks failed
Test / rust-fmt-check (pull_request) Failing after 0s
Test / rust-clippy (pull_request) Failing after 1s
Test / rust-tests (pull_request) Failing after 0s
Test / frontend-typecheck (pull_request) Failing after 16s
Test / frontend-tests (pull_request) Failing after 18s
PR Review Automation / review (pull_request) Failing after 4m13s
Complete backport of all features from apollo_nxt-trcaa repository: - Three-tier shell execution safety system (Tier 1: auto, Tier 2: approve, Tier 3: deny) - Ollama function calling with tool use support - AI provider tool calling auto-detection - kubectl binary bundling and management - kubeconfig upload and context management - Shell approval modal with real-time UI - MCP protocol HTTP transport with custom headers - Enhanced security audit logging - Comprehensive test coverage (275+ tests) - Updated CI/CD workflows for Gitea Actions - Complete documentation (ADRs, wiki, release notes) Sanitization applied to all files: - Removed all MSI, Motorola, VNXT, Vesta references - Replaced internal infrastructure references with TFTSR equivalents - Updated all URLs and API endpoints - Sanitized commit history references in documentation Technical changes: - New modules: shell/classifier, shell/executor, shell/kubectl, shell/kubeconfig - Enhanced AI providers: ollama.rs, openai.rs with function calling - New Tauri commands: shell execution, kubeconfig management, tool calling detection - Database migrations: shell_execution_audit table - Frontend: ShellApprovalModal, ShellExecution, KubeconfigManager pages - CI/CD: kubectl bundling, multi-platform builds, Gitea Actions integration Version: 1.0.8 Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
1887 lines
70 KiB
Rust
1887 lines
70 KiB
Rust
use rusqlite::OptionalExtension;
|
|
use tauri::{Manager, State};
|
|
use tracing::warn;
|
|
|
|
use crate::ai::agents::create_agent_registry;
|
|
use crate::ai::provider::create_provider;
|
|
use crate::ai::{AnalysisResult, ChatResponse, Message, ProviderInfo};
|
|
use crate::db::models::{AiConversation, AiMessage, AuditEntry};
|
|
use crate::state::{AppState, ProviderConfig};
|
|
|
|
#[tauri::command]
|
|
pub async fn analyze_logs(
|
|
issue_id: String,
|
|
log_file_ids: Vec<String>,
|
|
provider_config: ProviderConfig,
|
|
state: State<'_, AppState>,
|
|
) -> Result<AnalysisResult, String> {
|
|
// Load log file contents — only redacted files may be sent to an AI provider
|
|
let mut log_contents = String::new();
|
|
{
|
|
let db = state.db.lock().map_err(|e| e.to_string())?;
|
|
for file_id in &log_file_ids {
|
|
let mut stmt = db
|
|
.prepare("SELECT file_name, file_path, redacted FROM log_files WHERE id = ?1")
|
|
.map_err(|e| e.to_string())?;
|
|
if let Ok((name, path, redacted)) = stmt.query_row([file_id], |row| {
|
|
Ok((
|
|
row.get::<_, String>(0)?,
|
|
row.get::<_, String>(1)?,
|
|
row.get::<_, i32>(2)? != 0,
|
|
))
|
|
}) {
|
|
let redacted_path = redacted_path_for(&name, &path, redacted)?;
|
|
log_contents.push_str(&format!("--- {name} ---\n"));
|
|
if let Ok(content) = std::fs::read_to_string(&redacted_path) {
|
|
log_contents.push_str(&content);
|
|
} else {
|
|
log_contents.push_str("[Could not read redacted file]\n");
|
|
}
|
|
log_contents.push('\n');
|
|
}
|
|
}
|
|
}
|
|
|
|
let provider = create_provider(&provider_config);
|
|
|
|
let messages = vec![
|
|
Message {
|
|
role: "system".into(),
|
|
content: "You are an expert IT engineer. Analyze the following log content and \
|
|
identify key issues, errors, and anomalies. Structure your response as: \
|
|
SUMMARY: (one paragraph), KEY_FINDINGS: (bullet list), \
|
|
FIRST_WHY: (initial why question for 5-whys analysis), \
|
|
SEVERITY: (critical/high/medium/low)"
|
|
.into(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
},
|
|
Message {
|
|
role: "user".into(),
|
|
content: format!("Analyze logs for issue {issue_id}:\n\n{log_contents}"),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
},
|
|
];
|
|
|
|
let response = provider
|
|
.chat(messages, &provider_config, None)
|
|
.await
|
|
.map_err(|e| {
|
|
warn!(error = %e, "ai analyze_logs provider request failed");
|
|
"AI analysis request failed".to_string()
|
|
})?;
|
|
|
|
let content = &response.content;
|
|
let summary = extract_section(content, "SUMMARY:").unwrap_or_else(|| {
|
|
content
|
|
.lines()
|
|
.next()
|
|
.unwrap_or("Analysis complete")
|
|
.to_string()
|
|
});
|
|
let key_findings = extract_list(content, "KEY_FINDINGS:");
|
|
let suggested_why1 = extract_section(content, "FIRST_WHY:")
|
|
.unwrap_or_else(|| "Why did this issue occur?".to_string());
|
|
let severity_assessment =
|
|
extract_section(content, "SEVERITY:").unwrap_or_else(|| "medium".to_string());
|
|
|
|
// Write audit entry
|
|
{
|
|
let db = state.db.lock().map_err(|e| e.to_string())?;
|
|
let entry = AuditEntry::new(
|
|
"ai_analyze_logs".to_string(),
|
|
"issue".to_string(),
|
|
issue_id.clone(),
|
|
serde_json::json!({ "log_file_ids": log_file_ids, "provider": provider_config.name })
|
|
.to_string(),
|
|
);
|
|
crate::audit::log::write_audit_event(
|
|
&db,
|
|
&entry.action,
|
|
&entry.entity_type,
|
|
&entry.entity_id,
|
|
&entry.details,
|
|
)
|
|
.map_err(|_| "Failed to write security audit entry".to_string())?;
|
|
}
|
|
|
|
Ok(AnalysisResult {
|
|
summary,
|
|
key_findings,
|
|
suggested_why1,
|
|
severity_assessment,
|
|
})
|
|
}
|
|
|
|
/// Returns the path to the `.redacted` file, or an error if the file has not been redacted.
|
|
fn redacted_path_for(name: &str, path: &str, redacted: bool) -> Result<String, String> {
|
|
if !redacted {
|
|
return Err(format!(
|
|
"Log file '{name}' has not been scanned and redacted. \
|
|
Run PII detection and apply redactions before sending to AI."
|
|
));
|
|
}
|
|
Ok(format!("{path}.redacted"))
|
|
}
|
|
|
|
fn extract_section(text: &str, header: &str) -> Option<String> {
|
|
let start = text.find(header)?;
|
|
let after = &text[start + header.len()..];
|
|
let end = after.find('\n').unwrap_or(after.len());
|
|
let section = after[..end].trim().to_string();
|
|
if section.is_empty() {
|
|
None
|
|
} else {
|
|
Some(section)
|
|
}
|
|
}
|
|
|
|
fn extract_list(text: &str, header: &str) -> Vec<String> {
|
|
let start = match text.find(header) {
|
|
Some(s) => s + header.len(),
|
|
None => return vec![],
|
|
};
|
|
let section = &text[start..];
|
|
section
|
|
.lines()
|
|
.skip(1)
|
|
.take_while(|l| {
|
|
!l.chars()
|
|
.next()
|
|
.map(|c| c.is_uppercase() && l.contains(':'))
|
|
.unwrap_or(false)
|
|
})
|
|
.filter(|l| l.trim_start().starts_with('-') || l.trim_start().starts_with('*'))
|
|
.map(|l| {
|
|
l.trim_start_matches(|c: char| c == '-' || c == '*' || c.is_whitespace())
|
|
.to_string()
|
|
})
|
|
.filter(|s| !s.is_empty())
|
|
.collect()
|
|
}
|
|
|
|
/// Sanitize messages for final call when tool iteration limit is reached.
|
|
/// Converts tool role messages to assistant role with clear labeling as untrusted data.
|
|
fn sanitize_messages_for_final_call(messages: Vec<Message>) -> Vec<Message> {
|
|
messages
|
|
.into_iter()
|
|
.map(|mut msg| {
|
|
if msg.role == "tool" {
|
|
// Convert tool output to assistant role with clear labeling as untrusted data
|
|
msg.role = "assistant".into();
|
|
msg.content = format!("[UNTRUSTED TOOL OUTPUT]: {}", msg.content);
|
|
msg.tool_call_id = None;
|
|
}
|
|
msg.tool_calls = None; // Strip tool_calls from all messages
|
|
msg
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub async fn chat_message(
|
|
issue_id: String,
|
|
message: String,
|
|
log_file_ids: Option<Vec<String>>,
|
|
provider_config: ProviderConfig,
|
|
system_prompt: Option<String>,
|
|
app_handle: tauri::AppHandle,
|
|
state: State<'_, AppState>,
|
|
) -> Result<ChatResponse, String> {
|
|
// Find or create a conversation for this issue + provider
|
|
let conversation_id = {
|
|
let db = state.db.lock().map_err(|e| e.to_string())?;
|
|
let existing: Option<String> = db
|
|
.prepare("SELECT id FROM ai_conversations WHERE issue_id = ?1 AND provider = ?2 AND model = ?3 ORDER BY created_at DESC LIMIT 1")
|
|
.and_then(|mut stmt| {
|
|
stmt.query_row(
|
|
rusqlite::params![issue_id, provider_config.name, provider_config.model],
|
|
|row| row.get(0),
|
|
)
|
|
})
|
|
.ok();
|
|
|
|
match existing {
|
|
Some(id) => id,
|
|
None => {
|
|
let conv = AiConversation::new(
|
|
issue_id.clone(),
|
|
provider_config.name.clone(),
|
|
provider_config.model.clone(),
|
|
);
|
|
db.execute(
|
|
"INSERT INTO ai_conversations (id, issue_id, provider, model, created_at, title) \
|
|
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params![
|
|
conv.id, conv.issue_id, conv.provider, conv.model, conv.created_at, conv.title
|
|
],
|
|
)
|
|
.map_err(|e| e.to_string())?;
|
|
conv.id
|
|
}
|
|
}
|
|
};
|
|
|
|
// Load conversation history across ALL conversations for this issue
|
|
let history: Vec<Message> = {
|
|
let db = state.db.lock().map_err(|e| e.to_string())?;
|
|
let raw: Vec<(String, String)> = db
|
|
.prepare(
|
|
"SELECT am.role, am.content \
|
|
FROM ai_messages am \
|
|
JOIN ai_conversations ac ON ac.id = am.conversation_id \
|
|
WHERE ac.issue_id = ?1 \
|
|
ORDER BY am.created_at ASC",
|
|
)
|
|
.and_then(|mut stmt| {
|
|
stmt.query_map([&issue_id], |row| {
|
|
Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
|
|
})
|
|
.map(|rows| rows.filter_map(|r| r.ok()).collect::<Vec<_>>())
|
|
})
|
|
.unwrap_or_default();
|
|
drop(db);
|
|
raw.into_iter()
|
|
.map(|(role, content)| Message {
|
|
role,
|
|
content,
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
})
|
|
.collect()
|
|
};
|
|
|
|
// Auto-redact PII in both the typed message and any file attachments.
|
|
// The backend is the sole authority for redaction; the frontend sends original content.
|
|
let mut was_pii_redacted = false;
|
|
let mut redacted_pii_types: Vec<String> = Vec::new();
|
|
|
|
let full_message = {
|
|
// Step 1: redact the typed user message text.
|
|
let base = {
|
|
let spans = crate::pii::PiiDetector::new().detect(&message);
|
|
if spans.is_empty() {
|
|
message.clone()
|
|
} else {
|
|
let types: std::collections::HashSet<&str> =
|
|
spans.iter().map(|s| s.pii_type.as_str()).collect();
|
|
let mut type_list: Vec<&str> = types.into_iter().collect();
|
|
type_list.sort_unstable();
|
|
warn!(
|
|
pii_types = ?type_list,
|
|
pii_count = spans.len(),
|
|
"PII detected in typed chat message — auto-redacting before AI send"
|
|
);
|
|
was_pii_redacted = true;
|
|
redacted_pii_types.extend(type_list.iter().map(|s| s.to_string()));
|
|
crate::pii::apply_redactions(&message, &spans)
|
|
}
|
|
};
|
|
|
|
// Step 2: load attachment files from DB, scan, and embed clean content.
|
|
let files: Vec<(String, String)> = if let Some(ref ids) = log_file_ids {
|
|
let db = state.db.lock().map_err(|e| e.to_string())?;
|
|
let mut v = Vec::new();
|
|
for file_id in ids {
|
|
if let Ok((name, path)) = db
|
|
.prepare("SELECT file_name, file_path FROM log_files WHERE id = ?1")
|
|
.and_then(|mut s| {
|
|
s.query_row([file_id], |row| {
|
|
Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
|
|
})
|
|
})
|
|
{
|
|
v.push((name, path));
|
|
}
|
|
}
|
|
v
|
|
} else {
|
|
vec![]
|
|
};
|
|
|
|
// 8 KB embed limit; detect + redact on full content so PII at the boundary is caught.
|
|
const EMBED_LIMIT: usize = 8_000;
|
|
|
|
let mut msg = base;
|
|
for (file_name, file_path) in &files {
|
|
let content = std::fs::read_to_string(file_path).unwrap_or_default();
|
|
let spans = crate::pii::PiiDetector::new().detect(&content);
|
|
let redacted = if spans.is_empty() {
|
|
content
|
|
} else {
|
|
let types: std::collections::HashSet<&str> =
|
|
spans.iter().map(|s| s.pii_type.as_str()).collect();
|
|
let mut type_list: Vec<&str> = types.into_iter().collect();
|
|
type_list.sort_unstable();
|
|
warn!(
|
|
file_name = %file_name,
|
|
pii_types = ?type_list,
|
|
pii_count = spans.len(),
|
|
"PII detected in chat attachment — auto-redacting before AI send"
|
|
);
|
|
was_pii_redacted = true;
|
|
redacted_pii_types.extend(type_list.iter().map(|s| s.to_string()));
|
|
crate::pii::apply_redactions(&content, &spans)
|
|
};
|
|
// Truncate after redaction so the cut never lands inside a PII span.
|
|
let embed_end = if redacted.len() > EMBED_LIMIT {
|
|
let mut e = EMBED_LIMIT;
|
|
while !redacted.is_char_boundary(e) {
|
|
e -= 1;
|
|
}
|
|
e
|
|
} else {
|
|
redacted.len()
|
|
};
|
|
msg.push_str(&format!(
|
|
"\n\n--- Attached: {file_name} ---\n{}",
|
|
&redacted[..embed_end]
|
|
));
|
|
}
|
|
msg
|
|
};
|
|
|
|
let provider = create_provider(&provider_config);
|
|
|
|
// Search integration sources for relevant context
|
|
let integration_context = search_integration_sources(&message, &app_handle, &state).await;
|
|
|
|
// Load agent system
|
|
let agent_registry = create_agent_registry();
|
|
let devops_agent = agent_registry.get("devops-incident-responder");
|
|
|
|
let mut messages = Vec::new();
|
|
|
|
// Inject devops-incident-responder as primary system prompt (always)
|
|
if let Some(agent) = devops_agent {
|
|
messages.push(Message {
|
|
role: "system".into(),
|
|
content: agent.system_prompt.clone(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
});
|
|
}
|
|
|
|
// Inject domain system prompt if provided
|
|
if let Some(ref prompt) = system_prompt {
|
|
if !prompt.is_empty() {
|
|
messages.push(Message {
|
|
role: "system".into(),
|
|
content: prompt.clone(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
});
|
|
}
|
|
}
|
|
|
|
messages.extend(history);
|
|
|
|
// If we found integration content, add it to the conversation context
|
|
if !integration_context.is_empty() {
|
|
let context_message = Message {
|
|
role: "system".into(),
|
|
content: format!(
|
|
"INTERNAL DOCUMENTATION SOURCES:\n\n{integration_context}\n\n\
|
|
Instructions: The above content is from internal company documentation systems \
|
|
(Confluence, ServiceNow, Azure DevOps). \
|
|
\n\n**IMPORTANT**: First determine if this documentation is RELEVANT to the user's question:\
|
|
\n- If the documentation directly addresses the question → Use it and cite sources with URLs\
|
|
\n- If the documentation is tangentially related but doesn't answer the question → Briefly mention what internal docs exist, then provide a complete answer using general knowledge\
|
|
\n- If the documentation is completely unrelated → Ignore it and answer using general knowledge\
|
|
\n\nDo NOT force irrelevant internal documentation into your answer. The user needs accurate information, not forced citations."
|
|
),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
};
|
|
messages.push(context_message);
|
|
}
|
|
|
|
// Tool execution configuration
|
|
const MAX_TOOL_ITERATIONS: usize = 20; // Allow sufficient iterations for complex diagnostics
|
|
|
|
// Get available tools — static + MCP
|
|
// Only enable tools if the provider explicitly supports tool calling
|
|
let tools = if provider_config.supports_tool_calling.unwrap_or(false) {
|
|
let mut all_tools = crate::ai::tools::get_available_tools();
|
|
let mcp_tools = crate::ai::tools::get_enabled_mcp_tools(&state).await;
|
|
all_tools.extend(mcp_tools);
|
|
if all_tools.is_empty() {
|
|
None
|
|
} else {
|
|
Some(all_tools)
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// If tools are available AND using OpenAI-compatible provider, add explicit JSON format instruction
|
|
// Only OpenAI-compatible providers (default case in create_provider) actually support tool calling.
|
|
// Others (anthropic, gemini, mistral, ollama) either ignore tools or use provider-specific formats.
|
|
let is_openai_compatible = {
|
|
let kind = if provider_config.provider_type.is_empty() {
|
|
provider_config.name.as_str()
|
|
} else {
|
|
provider_config.provider_type.as_str()
|
|
};
|
|
!matches!(kind, "anthropic" | "gemini" | "mistral" | "ollama")
|
|
};
|
|
|
|
if tools.is_some() && is_openai_compatible {
|
|
messages.push(Message {
|
|
role: "system".into(),
|
|
content: "CRITICAL: You have tools available. When calling tools, you MUST use the native JSON function calling format in your API response. DO NOT output XML tags like <tool_name>. DO NOT output text descriptions of tool calls. Use the structured tool_calls field in your response.".into(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
});
|
|
|
|
// Add iteration budget awareness
|
|
messages.push(Message {
|
|
role: "system".into(),
|
|
content: format!(
|
|
"TOOL EXECUTION BUDGET: You have a maximum of {MAX_TOOL_ITERATIONS} rounds (each AI response counts as one round). \
|
|
You can call multiple tools in a single round. \
|
|
Plan your investigation efficiently:\n\
|
|
- Call multiple related tools in the same round when possible\n\
|
|
- Prioritize high-value diagnostic commands first\n\
|
|
- Use comprehensive output formats (e.g., kubectl --output=yaml) to gather more data per call\n\
|
|
- Reserve 1 round for your final summary/answer\n\
|
|
- If you exceed the budget, you'll be cut off mid-investigation\n\
|
|
Current round count is not visible to you, so plan conservatively."
|
|
),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
});
|
|
}
|
|
|
|
messages.push(Message {
|
|
role: "user".into(),
|
|
content: full_message.clone(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
});
|
|
|
|
// Tool-calling loop: keep calling until AI gives final answer
|
|
let final_response;
|
|
let mut iteration = 0;
|
|
|
|
loop {
|
|
iteration += 1;
|
|
|
|
// Warn AI when approaching limit
|
|
if iteration == MAX_TOOL_ITERATIONS - 2 {
|
|
messages.push(Message {
|
|
role: "system".into(),
|
|
content: format!(
|
|
"WARNING: You are on iteration {iteration}/{MAX_TOOL_ITERATIONS} (2 rounds remaining). \
|
|
You MUST provide your final answer in the NEXT round. \
|
|
Do NOT call any more tools. \
|
|
Summarize your findings based on the data you've already gathered."
|
|
),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
});
|
|
}
|
|
|
|
// Force stop at limit with collected data
|
|
if iteration > MAX_TOOL_ITERATIONS {
|
|
let sanitized_messages = sanitize_messages_for_final_call(messages);
|
|
|
|
// Add final instruction
|
|
let mut final_messages = sanitized_messages;
|
|
final_messages.push(Message {
|
|
role: "system".into(),
|
|
content: format!(
|
|
"CRITICAL: Tool iteration limit reached ({iteration}/{MAX_TOOL_ITERATIONS}). \
|
|
TOOLS ARE NOW DISABLED. \
|
|
You MUST respond now with a natural language summary of your findings. \
|
|
DO NOT attempt to call any tools - they will not execute. \
|
|
DO NOT emit tool_calls JSON - it will be ignored. \
|
|
Ignore any earlier instructions about tool calling or JSON formatting. \
|
|
Provide your best answer in plain text based on the diagnostic data already collected."
|
|
),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
});
|
|
|
|
// Make one final call WITHOUT tools to force text response
|
|
let final_attempt = provider
|
|
.chat(final_messages, &provider_config, None) // No tools available
|
|
.await
|
|
.map_err(|e| {
|
|
format!("AI provider request failed after reaching iteration limit: {e}")
|
|
})?;
|
|
|
|
final_response = final_attempt;
|
|
tracing::warn!(
|
|
"Tool iteration limit exceeded, forced final response: {} chars",
|
|
final_response.content.len()
|
|
);
|
|
break;
|
|
}
|
|
|
|
let response = provider
|
|
.chat(messages.clone(), &provider_config, tools.clone())
|
|
.await
|
|
.map_err(|e| {
|
|
let error_msg = format!("AI provider request failed: {e}");
|
|
warn!("{}", error_msg);
|
|
error_msg
|
|
})?;
|
|
|
|
// Check if AI wants to call tools
|
|
if let Some(tool_calls) = &response.tool_calls {
|
|
tracing::info!("AI requested {} tool call(s)", tool_calls.len());
|
|
|
|
// Execute each tool call
|
|
for tool_call in tool_calls {
|
|
tracing::info!("Executing tool: {}", tool_call.name);
|
|
|
|
let tool_result = execute_tool_call(tool_call, &app_handle, &state).await;
|
|
|
|
// Format result
|
|
let result_content = match tool_result {
|
|
Ok(result) => result,
|
|
Err(e) => format!("Error executing tool: {e}"),
|
|
};
|
|
|
|
// Add tool result as a message
|
|
messages.push(Message {
|
|
role: "tool".into(),
|
|
content: result_content,
|
|
tool_call_id: Some(tool_call.id.clone()),
|
|
tool_calls: None,
|
|
});
|
|
}
|
|
|
|
// Continue loop to get AI's next response
|
|
continue;
|
|
}
|
|
|
|
// No tool calls - this is the final answer
|
|
final_response = response;
|
|
break;
|
|
}
|
|
|
|
// Save both user message and response to DB
|
|
let stored_user_message;
|
|
{
|
|
let db = state.db.lock().map_err(|e| e.to_string())?;
|
|
let user_msg = AiMessage::new(conversation_id.clone(), "user".to_string(), full_message);
|
|
stored_user_message = user_msg.content.clone();
|
|
let asst_msg = AiMessage::new(
|
|
conversation_id,
|
|
"assistant".to_string(),
|
|
final_response.content.clone(),
|
|
);
|
|
|
|
db.execute(
|
|
"INSERT INTO ai_messages (id, conversation_id, role, content, token_count, created_at) \
|
|
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params![
|
|
user_msg.id, user_msg.conversation_id, user_msg.role,
|
|
user_msg.content, user_msg.token_count, user_msg.created_at
|
|
],
|
|
)
|
|
.map_err(|e| e.to_string())?;
|
|
|
|
db.execute(
|
|
"INSERT INTO ai_messages (id, conversation_id, role, content, token_count, created_at) \
|
|
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params![
|
|
asst_msg.id, asst_msg.conversation_id, asst_msg.role,
|
|
asst_msg.content, asst_msg.token_count, asst_msg.created_at
|
|
],
|
|
)
|
|
.map_err(|e| e.to_string())?;
|
|
|
|
// Audit - capture full transmission details
|
|
let pii_types_for_audit = {
|
|
use std::collections::HashSet;
|
|
let mut v: Vec<String> = redacted_pii_types
|
|
.into_iter()
|
|
.collect::<HashSet<_>>()
|
|
.into_iter()
|
|
.collect();
|
|
v.sort_unstable();
|
|
v
|
|
};
|
|
let audit_details = serde_json::json!({
|
|
"provider": provider_config.name,
|
|
"model": provider_config.model,
|
|
"api_url": provider_config.api_url,
|
|
"user_message": user_msg.content,
|
|
"was_pii_redacted": was_pii_redacted,
|
|
"pii_types_redacted": pii_types_for_audit,
|
|
"response_preview": if final_response.content.len() > 200 {
|
|
format!("{preview}...", preview = &final_response.content[..200])
|
|
} else {
|
|
final_response.content.clone()
|
|
},
|
|
"token_count": user_msg.token_count,
|
|
});
|
|
let entry = AuditEntry::new(
|
|
"ai_chat".to_string(),
|
|
"issue".to_string(),
|
|
issue_id,
|
|
audit_details.to_string(),
|
|
);
|
|
if let Err(err) = crate::audit::log::write_audit_event(
|
|
&db,
|
|
&entry.action,
|
|
&entry.entity_type,
|
|
&entry.entity_id,
|
|
&entry.details,
|
|
) {
|
|
warn!(error = %err, "failed to write ai_chat audit entry");
|
|
}
|
|
}
|
|
|
|
Ok(crate::ai::ChatResponse {
|
|
user_message: Some(stored_user_message),
|
|
..final_response
|
|
})
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub async fn test_provider_connection(
|
|
provider_config: ProviderConfig,
|
|
) -> Result<ChatResponse, String> {
|
|
let provider = create_provider(&provider_config);
|
|
let messages = vec![Message {
|
|
role: "user".into(),
|
|
content:
|
|
"Reply with exactly: Troubleshooting and RCA Assistant connection test successful."
|
|
.into(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
}];
|
|
provider
|
|
.chat(messages, &provider_config, None)
|
|
.await
|
|
.map_err(|e| {
|
|
warn!(error = %e, "ai test_provider_connection failed");
|
|
"Provider connection test failed".to_string()
|
|
})
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub async fn detect_tool_calling_support(provider_config: ProviderConfig) -> Result<bool, String> {
|
|
use crate::ai::{Tool, ToolParameters};
|
|
use std::collections::HashMap;
|
|
use tracing::info;
|
|
|
|
// Create a simple test tool
|
|
let test_tool = Tool {
|
|
name: "test_tool".to_string(),
|
|
description: "A test tool that returns 'success'. Call this tool with no arguments."
|
|
.to_string(),
|
|
parameters: ToolParameters {
|
|
param_type: "object".to_string(),
|
|
properties: HashMap::new(),
|
|
required: vec![],
|
|
},
|
|
};
|
|
|
|
// Override config with detection-optimized settings
|
|
let mut detection_config = provider_config.clone();
|
|
detection_config.max_tokens = Some(100); // Small budget for capability check
|
|
detection_config.temperature = Some(0.0); // Deterministic for reliability
|
|
|
|
let provider = create_provider(&detection_config);
|
|
let messages = vec![Message {
|
|
role: "user".into(),
|
|
content: "Please call the test_tool function.".into(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
}];
|
|
|
|
match provider
|
|
.chat(messages, &detection_config, Some(vec![test_tool]))
|
|
.await
|
|
{
|
|
Ok(response) => {
|
|
// Check if response contains tool_calls
|
|
if let Some(tool_calls) = response.tool_calls {
|
|
if tool_calls.iter().any(|tc| tc.name == "test_tool") {
|
|
info!(
|
|
"Tool calling support detected for provider {}",
|
|
provider_config.name
|
|
);
|
|
return Ok(true);
|
|
}
|
|
}
|
|
// Provider responded but didn't use tool calls
|
|
info!(
|
|
"Provider {} responded but did not call tool",
|
|
provider_config.name
|
|
);
|
|
Ok(false)
|
|
}
|
|
Err(e) => {
|
|
// Check if error indicates tool calling is not supported
|
|
let error_msg = e.to_string().to_lowercase();
|
|
if error_msg.contains("tool")
|
|
|| error_msg.contains("function")
|
|
|| error_msg.contains("503")
|
|
{
|
|
info!(
|
|
"Tool calling not supported for provider {}: {}",
|
|
provider_config.name, e
|
|
);
|
|
Ok(false)
|
|
} else {
|
|
// Connection or other error
|
|
Err(format!("Failed to test tool calling support: {e}"))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub async fn list_providers() -> Result<Vec<ProviderInfo>, String> {
|
|
Ok(vec![
|
|
ProviderInfo {
|
|
name: "openai".to_string(),
|
|
supports_streaming: true,
|
|
models: vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()],
|
|
},
|
|
ProviderInfo {
|
|
name: "anthropic".to_string(),
|
|
supports_streaming: true,
|
|
models: vec![
|
|
"claude-3-5-sonnet-20241022".to_string(),
|
|
"claude-3-haiku-20240307".to_string(),
|
|
],
|
|
},
|
|
ProviderInfo {
|
|
name: "gemini".to_string(),
|
|
supports_streaming: false,
|
|
models: vec!["gemini-1.5-pro".to_string(), "gemini-1.5-flash".to_string()],
|
|
},
|
|
ProviderInfo {
|
|
name: "mistral".to_string(),
|
|
supports_streaming: true,
|
|
models: vec![
|
|
"mistral-large-latest".to_string(),
|
|
"mistral-small-latest".to_string(),
|
|
],
|
|
},
|
|
ProviderInfo {
|
|
name: "ollama".to_string(),
|
|
supports_streaming: false,
|
|
models: vec!["llama3.2:3b".to_string(), "llama3.1:8b".to_string()],
|
|
},
|
|
])
|
|
}
|
|
|
|
/// Search integration sources (Confluence, ServiceNow, Azure DevOps) for relevant context
|
|
async fn search_integration_sources(
|
|
query: &str,
|
|
app_handle: &tauri::AppHandle,
|
|
state: &State<'_, AppState>,
|
|
) -> String {
|
|
let mut all_results = Vec::new();
|
|
|
|
// Try to get integration configurations
|
|
let configs: Vec<crate::commands::integrations::IntegrationConfig> = {
|
|
let db = match state.db.lock() {
|
|
Ok(db) => db,
|
|
Err(e) => {
|
|
tracing::warn!("Failed to lock database: {}", e);
|
|
return String::new();
|
|
}
|
|
};
|
|
|
|
let mut stmt = match db.prepare(
|
|
"SELECT service, base_url, username, project_name, space_key FROM integration_config",
|
|
) {
|
|
Ok(stmt) => stmt,
|
|
Err(e) => {
|
|
tracing::warn!("Failed to prepare statement: {}", e);
|
|
return String::new();
|
|
}
|
|
};
|
|
|
|
let rows = match stmt.query_map([], |row| {
|
|
Ok(crate::commands::integrations::IntegrationConfig {
|
|
service: row.get(0)?,
|
|
base_url: row.get(1)?,
|
|
username: row.get(2)?,
|
|
project_name: row.get(3)?,
|
|
space_key: row.get(4)?,
|
|
})
|
|
}) {
|
|
Ok(rows) => rows,
|
|
Err(e) => {
|
|
tracing::warn!("Failed to query integration configs: {}", e);
|
|
return String::new();
|
|
}
|
|
};
|
|
|
|
rows.filter_map(|r| r.ok()).collect()
|
|
};
|
|
|
|
// Search each available integration in parallel
|
|
let mut search_tasks = Vec::new();
|
|
|
|
for config in configs {
|
|
// Authentication priority:
|
|
// 1. Try cookies from persistent browser (may fail for HttpOnly)
|
|
// 2. Try stored credentials from database
|
|
// 3. Fall back to webview-based search (uses browser's session directly)
|
|
|
|
let cookies_opt = match crate::commands::integrations::get_fresh_cookies_from_webview(
|
|
&config.service,
|
|
app_handle,
|
|
state,
|
|
)
|
|
.await
|
|
{
|
|
Ok(Some(cookies)) => {
|
|
tracing::info!("Using extracted cookies for {}", config.service);
|
|
Some(cookies)
|
|
}
|
|
_ => {
|
|
// Fallback: check for stored credentials in database
|
|
tracing::info!(
|
|
"Cookie extraction failed for {}, checking stored credentials",
|
|
config.service
|
|
);
|
|
let encrypted_token: Option<String> = {
|
|
let db = match state.db.lock() {
|
|
Ok(db) => db,
|
|
Err(_) => continue,
|
|
};
|
|
db.query_row(
|
|
"SELECT encrypted_token FROM credentials WHERE service = ?1",
|
|
[&config.service],
|
|
|row| row.get::<_, String>(0),
|
|
)
|
|
.optional()
|
|
.ok()
|
|
.flatten()
|
|
};
|
|
|
|
if let Some(token) = encrypted_token {
|
|
if let Ok(decrypted) = crate::integrations::auth::decrypt_token(&token) {
|
|
// Try to parse as cookies JSON
|
|
if let Ok(cookie_list) = serde_json::from_str::<
|
|
Vec<crate::integrations::webview_auth::Cookie>,
|
|
>(&decrypted)
|
|
{
|
|
tracing::info!(
|
|
"Using stored cookies for {} (count: {})",
|
|
config.service,
|
|
cookie_list.len()
|
|
);
|
|
Some(cookie_list)
|
|
} else {
|
|
tracing::warn!(
|
|
"Stored credentials for {} not in cookie format",
|
|
config.service
|
|
);
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
};
|
|
|
|
// If we have cookies (from extraction or database), use standard API search
|
|
if let Some(cookies) = cookies_opt {
|
|
match config.service.as_str() {
|
|
"confluence" => {
|
|
let base_url = config.base_url.clone();
|
|
let query = query.to_string();
|
|
let cookies_clone = cookies.clone();
|
|
search_tasks.push(tokio::spawn(async move {
|
|
crate::integrations::confluence_search::search_confluence(
|
|
&base_url,
|
|
&query,
|
|
&cookies_clone,
|
|
)
|
|
.await
|
|
.unwrap_or_default()
|
|
}));
|
|
}
|
|
"servicenow" => {
|
|
let instance_url = config.base_url.clone();
|
|
let query = query.to_string();
|
|
let cookies_clone = cookies.clone();
|
|
search_tasks.push(tokio::spawn(async move {
|
|
let mut results = Vec::new();
|
|
// Search knowledge base
|
|
if let Ok(kb_results) =
|
|
crate::integrations::servicenow_search::search_servicenow(
|
|
&instance_url,
|
|
&query,
|
|
&cookies_clone,
|
|
)
|
|
.await
|
|
{
|
|
results.extend(kb_results);
|
|
}
|
|
// Search incidents
|
|
if let Ok(incident_results) =
|
|
crate::integrations::servicenow_search::search_incidents(
|
|
&instance_url,
|
|
&query,
|
|
&cookies_clone,
|
|
)
|
|
.await
|
|
{
|
|
results.extend(incident_results);
|
|
}
|
|
results
|
|
}));
|
|
}
|
|
"azuredevops" => {
|
|
let org_url = config.base_url.clone();
|
|
let project = config.project_name.unwrap_or_default();
|
|
let query = query.to_string();
|
|
let cookies_clone = cookies.clone();
|
|
search_tasks.push(tokio::spawn(async move {
|
|
let mut results = Vec::new();
|
|
// Search wiki
|
|
if let Ok(wiki_results) =
|
|
crate::integrations::azuredevops_search::search_wiki(
|
|
&org_url,
|
|
&project,
|
|
&query,
|
|
&cookies_clone,
|
|
)
|
|
.await
|
|
{
|
|
results.extend(wiki_results);
|
|
}
|
|
// Search work items
|
|
if let Ok(wi_results) =
|
|
crate::integrations::azuredevops_search::search_work_items(
|
|
&org_url,
|
|
&project,
|
|
&query,
|
|
&cookies_clone,
|
|
)
|
|
.await
|
|
{
|
|
results.extend(wi_results);
|
|
}
|
|
results
|
|
}));
|
|
}
|
|
_ => {}
|
|
}
|
|
} else {
|
|
// Final fallback: try webview-based fetch (includes HttpOnly cookies automatically)
|
|
// This makes HTTP requests FROM the authenticated webview, which includes all cookies
|
|
tracing::info!(
|
|
"No extracted cookies for {}, trying webview-based fetch",
|
|
config.service
|
|
);
|
|
|
|
// Check if webview exists for this service
|
|
let webview_label = {
|
|
let webviews = match state.integration_webviews.lock() {
|
|
Ok(w) => w,
|
|
Err(_) => continue,
|
|
};
|
|
webviews.get(&config.service).cloned()
|
|
};
|
|
|
|
if let Some(label) = webview_label {
|
|
// Get window handle
|
|
if let Some(webview_window) = app_handle.get_webview_window(&label) {
|
|
let base_url = config.base_url.clone();
|
|
let service = config.service.clone();
|
|
let query_str = query.to_string();
|
|
|
|
match service.as_str() {
|
|
"confluence" => {
|
|
search_tasks.push(tokio::spawn(async move {
|
|
tracing::info!("Executing Confluence search via webview fetch");
|
|
match crate::integrations::webview_fetch::search_confluence_webview(
|
|
&webview_window,
|
|
&base_url,
|
|
&query_str,
|
|
)
|
|
.await
|
|
{
|
|
Ok(results) => {
|
|
tracing::info!(
|
|
"Webview fetch for Confluence returned {} results",
|
|
results.len()
|
|
);
|
|
results
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
"Webview fetch failed for Confluence: {}",
|
|
e
|
|
);
|
|
Vec::new()
|
|
}
|
|
}
|
|
}));
|
|
}
|
|
"servicenow" => {
|
|
search_tasks.push(tokio::spawn(async move {
|
|
tracing::info!("Executing ServiceNow search via webview fetch");
|
|
match crate::integrations::webview_fetch::search_servicenow_webview(
|
|
&webview_window,
|
|
&base_url,
|
|
&query_str,
|
|
)
|
|
.await
|
|
{
|
|
Ok(results) => {
|
|
tracing::info!(
|
|
"Webview fetch for ServiceNow returned {} results",
|
|
results.len()
|
|
);
|
|
results
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
"Webview fetch failed for ServiceNow: {}",
|
|
e
|
|
);
|
|
Vec::new()
|
|
}
|
|
}
|
|
}));
|
|
}
|
|
"azuredevops" => {
|
|
let project = config.project_name.unwrap_or_default();
|
|
search_tasks.push(tokio::spawn(async move {
|
|
tracing::info!("Executing Azure DevOps search via webview fetch");
|
|
let mut results = Vec::new();
|
|
|
|
// Search wiki
|
|
match crate::integrations::webview_fetch::search_azuredevops_wiki_webview(
|
|
&webview_window,
|
|
&base_url,
|
|
&project,
|
|
&query_str
|
|
).await {
|
|
Ok(wiki_results) => {
|
|
tracing::info!("Webview fetch for ADO wiki returned {} results", wiki_results.len());
|
|
results.extend(wiki_results);
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!("Webview fetch failed for ADO wiki: {}", e);
|
|
}
|
|
}
|
|
|
|
// Search work items
|
|
match crate::integrations::webview_fetch::search_azuredevops_workitems_webview(
|
|
&webview_window,
|
|
&base_url,
|
|
&project,
|
|
&query_str
|
|
).await {
|
|
Ok(wi_results) => {
|
|
tracing::info!("Webview fetch for ADO work items returned {} results", wi_results.len());
|
|
results.extend(wi_results);
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!("Webview fetch failed for ADO work items: {}", e);
|
|
}
|
|
}
|
|
|
|
results
|
|
}));
|
|
}
|
|
_ => {}
|
|
}
|
|
} else {
|
|
tracing::warn!("Webview window not found for {}", config.service);
|
|
}
|
|
} else {
|
|
tracing::warn!(
|
|
"No webview open for {} - cannot search. Please open browser window in Settings → Integrations",
|
|
config.service
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Wait for all searches to complete
|
|
for task in search_tasks {
|
|
if let Ok(results) = task.await {
|
|
all_results.extend(results);
|
|
}
|
|
}
|
|
|
|
// Format results for AI context
|
|
if all_results.is_empty() {
|
|
return String::new();
|
|
}
|
|
|
|
let mut context = String::new();
|
|
for (idx, result) in all_results.iter().enumerate() {
|
|
context.push_str(&format!("--- SOURCE {} ({}) ---\n", idx + 1, result.source));
|
|
context.push_str(&format!("Title: {}\n", result.title));
|
|
context.push_str(&format!("URL: {}\n", result.url));
|
|
|
|
if let Some(content) = &result.content {
|
|
context.push_str(&format!("Content:\n{content}\n\n"));
|
|
} else {
|
|
context.push_str(&format!("Excerpt: {}\n\n", result.excerpt));
|
|
}
|
|
}
|
|
|
|
tracing::info!(
|
|
"Found {} integration sources for AI context",
|
|
all_results.len()
|
|
);
|
|
context
|
|
}
|
|
|
|
/// Execute a tool call made by the AI
|
|
async fn execute_tool_call(
|
|
tool_call: &crate::ai::ToolCall,
|
|
app_handle: &tauri::AppHandle,
|
|
app_state: &State<'_, AppState>,
|
|
) -> Result<String, String> {
|
|
match tool_call.name.as_str() {
|
|
"add_ado_comment" => {
|
|
// Parse arguments
|
|
let args: serde_json::Value = serde_json::from_str(&tool_call.arguments)
|
|
.map_err(|e| format!("Failed to parse tool arguments: {e}"))?;
|
|
|
|
let work_item_id = args
|
|
.get("work_item_id")
|
|
.and_then(|v| v.as_i64())
|
|
.ok_or_else(|| "Missing or invalid work_item_id parameter".to_string())?;
|
|
|
|
let comment_text = args
|
|
.get("comment_text")
|
|
.and_then(|v| v.as_str())
|
|
.ok_or_else(|| "Missing or invalid comment_text parameter".to_string())?;
|
|
|
|
// Execute the add_ado_comment command
|
|
tracing::info!(
|
|
"AI executing tool: add_ado_comment({}, \"{}\")",
|
|
work_item_id,
|
|
comment_text
|
|
);
|
|
crate::commands::integrations::add_ado_comment(
|
|
work_item_id,
|
|
comment_text.to_string(),
|
|
app_handle.clone(),
|
|
app_state.clone(),
|
|
)
|
|
.await
|
|
}
|
|
"execute_shell_command" => execute_shell_tool_call(tool_call, app_handle, app_state).await,
|
|
name if name.starts_with("mcp_") => execute_mcp_tool_call(tool_call, app_state).await,
|
|
_ => {
|
|
let error = format!("Unknown tool: {}", tool_call.name);
|
|
tracing::warn!("{error}");
|
|
Err(error)
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn execute_shell_tool_call(
|
|
tool_call: &crate::ai::ToolCall,
|
|
app_handle: &tauri::AppHandle,
|
|
app_state: &State<'_, AppState>,
|
|
) -> Result<String, String> {
|
|
// Parse arguments
|
|
let args: serde_json::Value = serde_json::from_str(&tool_call.arguments)
|
|
.map_err(|e| format!("Failed to parse tool arguments: {e}"))?;
|
|
|
|
let command = args
|
|
.get("command")
|
|
.and_then(|v| v.as_str())
|
|
.ok_or_else(|| "Missing command parameter".to_string())?;
|
|
|
|
let working_directory = args.get("working_directory").and_then(|v| v.as_str());
|
|
let kubeconfig_id = args.get("kubeconfig_id").and_then(|v| v.as_str());
|
|
|
|
// PII detection
|
|
{
|
|
let detector = crate::pii::detector::PiiDetector::new();
|
|
let spans = detector.detect(command);
|
|
if !spans.is_empty() {
|
|
tracing::warn!(
|
|
tool = %tool_call.name,
|
|
pii_spans = spans.len(),
|
|
"PII detected in shell command"
|
|
);
|
|
}
|
|
}
|
|
|
|
// Audit log
|
|
{
|
|
let db = app_state.db.lock().map_err(|e| e.to_string())?;
|
|
let details = serde_json::json!({
|
|
"tool": tool_call.name,
|
|
"command": command,
|
|
});
|
|
crate::audit::log::write_audit_event(
|
|
&db,
|
|
"shell_tool_call",
|
|
"shell_command",
|
|
command,
|
|
&details.to_string(),
|
|
)
|
|
.map_err(|e| format!("Audit log failed: {e}"))?;
|
|
}
|
|
|
|
// Execute with approval flow
|
|
let result = crate::shell::executor::execute_with_approval(
|
|
command,
|
|
app_handle,
|
|
app_state.inner(),
|
|
kubeconfig_id,
|
|
working_directory,
|
|
)
|
|
.await?;
|
|
|
|
// Format output for AI
|
|
Ok(format!(
|
|
"Exit Code: {}\n\nStdout:\n{}\n\nStderr:\n{}",
|
|
result.exit_code, result.stdout, result.stderr
|
|
))
|
|
}
|
|
|
|
async fn execute_mcp_tool_call(
|
|
tool_call: &crate::ai::ToolCall,
|
|
app_state: &State<'_, AppState>,
|
|
) -> Result<String, String> {
|
|
// PII scan — log warning if sensitive data detected (non-blocking)
|
|
{
|
|
let detector = crate::pii::detector::PiiDetector::new();
|
|
let spans = detector.detect(&tool_call.arguments);
|
|
if !spans.is_empty() {
|
|
tracing::warn!(
|
|
tool = %tool_call.name,
|
|
pii_spans = spans.len(),
|
|
"PII detected in MCP tool call arguments"
|
|
);
|
|
}
|
|
}
|
|
|
|
// Audit log — mandatory before any external call
|
|
{
|
|
let db = app_state.db.lock().map_err(|e| e.to_string())?;
|
|
let details = serde_json::json!({
|
|
"tool": tool_call.name,
|
|
"args_preview": if tool_call.arguments.len() > 200 {
|
|
format!("{}...", &tool_call.arguments[..200])
|
|
} else {
|
|
tool_call.arguments.clone()
|
|
},
|
|
});
|
|
crate::audit::log::write_audit_event(
|
|
&db,
|
|
"mcp_tool_call",
|
|
"mcp_tool",
|
|
&tool_call.name,
|
|
&details.to_string(),
|
|
)
|
|
.map_err(|e| format!("Audit log failed: {e}"))?;
|
|
}
|
|
|
|
// Look up the tool → server_id + raw tool name
|
|
let (server_id, raw_tool_name) = {
|
|
let db = app_state.db.lock().map_err(|e| e.to_string())?;
|
|
let tool = crate::mcp::store::get_tool_by_key(&db, &tool_call.name)?
|
|
.ok_or_else(|| format!("MCP tool not found: {}", tool_call.name))?;
|
|
(tool.server_id, tool.name)
|
|
};
|
|
|
|
// Get connection from state — use tokio Mutex, never hold std::Mutex simultaneously
|
|
let conn_arc = {
|
|
let connections = app_state.mcp_connections.lock().await;
|
|
connections
|
|
.get(&server_id)
|
|
.cloned()
|
|
.ok_or_else(|| format!("No active connection for MCP server {server_id}"))?
|
|
};
|
|
|
|
let args: serde_json::Value = serde_json::from_str(&tool_call.arguments)
|
|
.unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
|
|
|
|
let conn = conn_arc.lock().await;
|
|
crate::mcp::client::call_tool(&conn, &raw_tool_name, &args).await
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_extract_section_basic() {
|
|
let text = "SUMMARY: The server crashed due to OOM.\nKEY_FINDINGS:\n- item";
|
|
assert_eq!(
|
|
extract_section(text, "SUMMARY:"),
|
|
Some("The server crashed due to OOM.".to_string())
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_section_missing() {
|
|
let text = "No matching header here";
|
|
assert_eq!(extract_section(text, "SUMMARY:"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_section_empty_value() {
|
|
let text = "SUMMARY:\nKEY_FINDINGS:";
|
|
assert_eq!(extract_section(text, "SUMMARY:"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_section_last_line() {
|
|
let text = "SEVERITY: critical";
|
|
assert_eq!(
|
|
extract_section(text, "SEVERITY:"),
|
|
Some("critical".to_string())
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_list_basic() {
|
|
let text = "KEY_FINDINGS:\n- First issue\n- Second issue\nSEVERITY: high";
|
|
let list = extract_list(text, "KEY_FINDINGS:");
|
|
assert_eq!(list, vec!["First issue", "Second issue"]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_list_with_asterisks() {
|
|
let text = "KEY_FINDINGS:\n* Item one\n* Item two\nFIRST_WHY: Why?";
|
|
let list = extract_list(text, "KEY_FINDINGS:");
|
|
assert_eq!(list, vec!["Item one", "Item two"]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_redacted_path_rejects_unredacted_file() {
|
|
let err = redacted_path_for("app.log", "/data/app.log", false).unwrap_err();
|
|
assert!(err.contains("app.log"));
|
|
assert!(err.contains("redacted"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_redacted_path_returns_dotredacted_suffix() {
|
|
let path = redacted_path_for("app.log", "/data/app.log", true).unwrap();
|
|
assert_eq!(path, "/data/app.log.redacted");
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_list_missing_header() {
|
|
let text = "No findings here";
|
|
let list = extract_list(text, "KEY_FINDINGS:");
|
|
assert!(list.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_list_empty_items_filtered() {
|
|
let text = "KEY_FINDINGS:\n- \n- Actual item\n- \nNEXT:";
|
|
let list = extract_list(text, "KEY_FINDINGS:");
|
|
assert_eq!(list, vec!["Actual item"]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_history_query_same_conversation() {
|
|
let conn = rusqlite::Connection::open_in_memory().unwrap();
|
|
crate::db::migrations::run_migrations(&conn).unwrap();
|
|
|
|
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
|
conn.execute(
|
|
"INSERT INTO issues (id, title, created_at, updated_at) VALUES (?1, ?2, ?3, ?4)",
|
|
rusqlite::params!["issue-1", "Test", &now, &now],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_conversations (id, issue_id, provider, model, created_at, title) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["conv-1", "issue-1", "openai", "gpt-4o", &now, "Conv 1"],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_messages (id, conversation_id, role, content, token_count, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["msg-1", "conv-1", "user", "Hello", 5, "2025-01-01 10:00:00"],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_messages (id, conversation_id, role, content, token_count, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["msg-2", "conv-1", "assistant", "Hi there", 8, "2025-01-01 10:00:01"],
|
|
)
|
|
.unwrap();
|
|
|
|
let issue_id = "issue-1";
|
|
let raw: Vec<(String, String)> = conn
|
|
.prepare(
|
|
"SELECT am.role, am.content \
|
|
FROM ai_messages am \
|
|
JOIN ai_conversations ac ON ac.id = am.conversation_id \
|
|
WHERE ac.issue_id = ?1 \
|
|
ORDER BY am.created_at ASC",
|
|
)
|
|
.and_then(|mut stmt| {
|
|
stmt.query_map([&issue_id], |row| {
|
|
Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
|
|
})
|
|
.map(|rows| rows.filter_map(|r| r.ok()).collect::<Vec<_>>())
|
|
})
|
|
.unwrap();
|
|
|
|
assert_eq!(raw.len(), 2);
|
|
assert_eq!(raw[0], ("user".to_string(), "Hello".to_string()));
|
|
assert_eq!(raw[1], ("assistant".to_string(), "Hi there".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_history_query_across_conversations() {
|
|
let conn = rusqlite::Connection::open_in_memory().unwrap();
|
|
crate::db::migrations::run_migrations(&conn).unwrap();
|
|
|
|
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
|
conn.execute(
|
|
"INSERT INTO issues (id, title, created_at, updated_at) VALUES (?1, ?2, ?3, ?4)",
|
|
rusqlite::params!["issue-1", "Test", &now, &now],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_conversations (id, issue_id, provider, model, created_at, title) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["conv-1", "issue-1", "openai", "gpt-4o", &now, "Conv 1"],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_conversations (id, issue_id, provider, model, created_at, title) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["conv-2", "issue-1", "anthropic", "claude-3", &now, "Conv 2"],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_messages (id, conversation_id, role, content, token_count, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["msg-1", "conv-1", "user", "From conv 1", 5, "2025-01-01 10:00:00"],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_messages (id, conversation_id, role, content, token_count, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["msg-2", "conv-2", "user", "From conv 2", 5, "2025-01-01 11:00:00"],
|
|
)
|
|
.unwrap();
|
|
|
|
let issue_id = "issue-1";
|
|
let raw: Vec<(String, String)> = conn
|
|
.prepare(
|
|
"SELECT am.role, am.content \
|
|
FROM ai_messages am \
|
|
JOIN ai_conversations ac ON ac.id = am.conversation_id \
|
|
WHERE ac.issue_id = ?1 \
|
|
ORDER BY am.created_at ASC",
|
|
)
|
|
.and_then(|mut stmt| {
|
|
stmt.query_map([&issue_id], |row| {
|
|
Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
|
|
})
|
|
.map(|rows| rows.filter_map(|r| r.ok()).collect::<Vec<_>>())
|
|
})
|
|
.unwrap();
|
|
|
|
assert_eq!(raw.len(), 2);
|
|
assert_eq!(raw[0].1, "From conv 1");
|
|
assert_eq!(raw[1].1, "From conv 2");
|
|
}
|
|
|
|
#[test]
|
|
fn test_history_query_empty_for_new_issue() {
|
|
let conn = rusqlite::Connection::open_in_memory().unwrap();
|
|
crate::db::migrations::run_migrations(&conn).unwrap();
|
|
|
|
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
|
conn.execute(
|
|
"INSERT INTO issues (id, title, created_at, updated_at) VALUES (?1, ?2, ?3, ?4)",
|
|
rusqlite::params!["issue-new", "Empty", &now, &now],
|
|
)
|
|
.unwrap();
|
|
|
|
let issue_id = "issue-new";
|
|
let raw: Vec<(String, String)> = conn
|
|
.prepare(
|
|
"SELECT am.role, am.content \
|
|
FROM ai_messages am \
|
|
JOIN ai_conversations ac ON ac.id = am.conversation_id \
|
|
WHERE ac.issue_id = ?1 \
|
|
ORDER BY am.created_at ASC",
|
|
)
|
|
.and_then(|mut stmt| {
|
|
stmt.query_map([&issue_id], |row| {
|
|
Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
|
|
})
|
|
.map(|rows| rows.filter_map(|r| r.ok()).collect::<Vec<_>>())
|
|
})
|
|
.unwrap();
|
|
|
|
assert!(raw.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_history_query_ordered_by_created_at() {
|
|
let conn = rusqlite::Connection::open_in_memory().unwrap();
|
|
crate::db::migrations::run_migrations(&conn).unwrap();
|
|
|
|
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
|
conn.execute(
|
|
"INSERT INTO issues (id, title, created_at, updated_at) VALUES (?1, ?2, ?3, ?4)",
|
|
rusqlite::params!["issue-1", "Test", &now, &now],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_conversations (id, issue_id, provider, model, created_at, title) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["conv-1", "issue-1", "openai", "gpt-4o", &now, "Conv 1"],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_conversations (id, issue_id, provider, model, created_at, title) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["conv-2", "issue-1", "anthropic", "claude-3", &now, "Conv 2"],
|
|
)
|
|
.unwrap();
|
|
// Insert messages out of order: conv-2 message is earlier than conv-1 message
|
|
conn.execute(
|
|
"INSERT INTO ai_messages (id, conversation_id, role, content, token_count, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["msg-1", "conv-1", "user", "Second", 5, "2025-01-01 12:00:00"],
|
|
)
|
|
.unwrap();
|
|
conn.execute(
|
|
"INSERT INTO ai_messages (id, conversation_id, role, content, token_count, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
|
|
rusqlite::params!["msg-2", "conv-2", "user", "First", 5, "2025-01-01 09:00:00"],
|
|
)
|
|
.unwrap();
|
|
|
|
let issue_id = "issue-1";
|
|
let raw: Vec<(String, String)> = conn
|
|
.prepare(
|
|
"SELECT am.role, am.content \
|
|
FROM ai_messages am \
|
|
JOIN ai_conversations ac ON ac.id = am.conversation_id \
|
|
WHERE ac.issue_id = ?1 \
|
|
ORDER BY am.created_at ASC",
|
|
)
|
|
.and_then(|mut stmt| {
|
|
stmt.query_map([&issue_id], |row| {
|
|
Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
|
|
})
|
|
.map(|rows| rows.filter_map(|r| r.ok()).collect::<Vec<_>>())
|
|
})
|
|
.unwrap();
|
|
|
|
assert_eq!(raw.len(), 2);
|
|
assert_eq!(raw[0].1, "First");
|
|
assert_eq!(raw[1].1, "Second");
|
|
}
|
|
|
|
#[test]
|
|
fn test_message_sanitization_converts_tool_role_to_assistant() {
|
|
// Messages with 'tool' role that would cause validation errors
|
|
let messages = vec![
|
|
Message {
|
|
role: "user".into(),
|
|
content: "What pods are running?".into(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
},
|
|
Message {
|
|
role: "assistant".into(),
|
|
content: "".into(),
|
|
tool_call_id: None,
|
|
tool_calls: Some(vec![crate::ai::ToolCall {
|
|
id: "call_123".into(),
|
|
name: "execute_shell_command".into(),
|
|
arguments: r#"{"command":"kubectl get pods"}"#.into(),
|
|
}]),
|
|
},
|
|
Message {
|
|
role: "tool".into(),
|
|
content: "pod1 Running\npod2 Running".into(),
|
|
tool_call_id: Some("call_123".into()),
|
|
tool_calls: None,
|
|
},
|
|
];
|
|
|
|
// Use production sanitization helper
|
|
let sanitized = sanitize_messages_for_final_call(messages);
|
|
|
|
// Verify sanitization
|
|
assert_eq!(sanitized.len(), 3);
|
|
assert_eq!(sanitized[0].role, "user");
|
|
assert_eq!(sanitized[1].role, "assistant");
|
|
assert_eq!(sanitized[1].tool_calls, None); // Stripped
|
|
assert_eq!(sanitized[2].role, "assistant"); // Converted from 'tool' to assistant
|
|
assert!(sanitized[2].content.starts_with("[UNTRUSTED TOOL OUTPUT]:")); // Labeled as untrusted
|
|
assert_eq!(sanitized[2].tool_call_id, None); // Stripped
|
|
}
|
|
|
|
#[test]
|
|
fn test_message_sanitization_preserves_non_tool_messages() {
|
|
let messages = vec![
|
|
Message {
|
|
role: "system".into(),
|
|
content: "You are a helpful assistant".into(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
},
|
|
Message {
|
|
role: "user".into(),
|
|
content: "Hello".into(),
|
|
tool_call_id: None,
|
|
tool_calls: None,
|
|
},
|
|
];
|
|
|
|
// Use production sanitization helper
|
|
let sanitized = sanitize_messages_for_final_call(messages);
|
|
|
|
// Verify non-tool messages preserved
|
|
assert_eq!(sanitized.len(), 2);
|
|
assert_eq!(sanitized[0].role, "system");
|
|
assert_eq!(sanitized[0].content, "You are a helpful assistant");
|
|
assert_eq!(sanitized[1].role, "user");
|
|
assert_eq!(sanitized[1].content, "Hello");
|
|
}
|
|
|
|
#[test]
|
|
fn test_sanitize_messages_strips_tool_calls_from_all_messages() {
|
|
let messages = vec![
|
|
Message {
|
|
role: "assistant".into(),
|
|
content: "".into(),
|
|
tool_call_id: None,
|
|
tool_calls: Some(vec![crate::ai::ToolCall {
|
|
id: "call_1".into(),
|
|
name: "test".into(),
|
|
arguments: "{}".into(),
|
|
}]),
|
|
},
|
|
Message {
|
|
role: "tool".into(),
|
|
content: "output".into(),
|
|
tool_call_id: Some("call_1".into()),
|
|
tool_calls: None,
|
|
},
|
|
];
|
|
|
|
let sanitized = sanitize_messages_for_final_call(messages);
|
|
|
|
assert_eq!(sanitized.len(), 2);
|
|
assert_eq!(sanitized[0].tool_calls, None); // Stripped from assistant
|
|
assert_eq!(sanitized[1].tool_calls, None); // Already None, but verified
|
|
}
|
|
|
|
#[test]
|
|
fn test_supports_tool_calling_flag_false_disables_tools() {
|
|
// Test the logic that enforces supports_tool_calling flag
|
|
// When flag is false, tools should be None even if tools exist
|
|
|
|
let supports_tool_calling = Some(false);
|
|
let has_tools_available = true; // Simulate available tools
|
|
|
|
// Simulate the conditional logic from chat_message command (lines 403-415)
|
|
let tools_enabled = if supports_tool_calling.unwrap_or(false) {
|
|
if has_tools_available {
|
|
Some(vec!["mock_tool"])
|
|
} else {
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
assert!(
|
|
tools_enabled.is_none(),
|
|
"Tools should be None when supports_tool_calling is false"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_supports_tool_calling_flag_true_enables_tools() {
|
|
let supports_tool_calling = Some(true);
|
|
let has_tools_available = true;
|
|
|
|
let tools_enabled = if supports_tool_calling.unwrap_or(false) {
|
|
if has_tools_available {
|
|
Some(vec!["mock_tool"])
|
|
} else {
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
assert!(
|
|
tools_enabled.is_some(),
|
|
"Tools should be enabled when supports_tool_calling is true"
|
|
);
|
|
assert_eq!(tools_enabled.unwrap().len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_supports_tool_calling_flag_none_defaults_to_false() {
|
|
let supports_tool_calling: Option<bool> = None;
|
|
let has_tools_available = true;
|
|
|
|
let tools_enabled = if supports_tool_calling.unwrap_or(false) {
|
|
if has_tools_available {
|
|
Some(vec!["mock_tool"])
|
|
} else {
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
assert!(
|
|
tools_enabled.is_none(),
|
|
"Tools should be None when supports_tool_calling is None (defaults to false)"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_supports_tool_calling_true_but_no_tools_available() {
|
|
let supports_tool_calling = Some(true);
|
|
let has_tools_available = false;
|
|
|
|
let tools_enabled = if supports_tool_calling.unwrap_or(false) {
|
|
if has_tools_available {
|
|
Some(vec!["mock_tool"])
|
|
} else {
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
assert!(
|
|
tools_enabled.is_none(),
|
|
"Tools should be None even when flag is true if no tools are available"
|
|
);
|
|
}
|
|
|
|
// Tests for detect_tool_calling_support
|
|
// NOTE: These are unit tests for the detection logic, not integration tests with real providers.
|
|
// They verify the logical correctness of detection criteria (checking for tool_calls presence,
|
|
// error pattern matching) but do not exercise the full command implementation.
|
|
// Tradeoff: These tests provide fast feedback on detection logic without requiring network calls
|
|
// or mock providers. Integration tests with real providers would be more comprehensive but slower
|
|
// and require test infrastructure setup.
|
|
|
|
#[test]
|
|
fn test_detect_tool_calling_logic_with_tool_calls_in_response() {
|
|
use crate::ai::ToolCall;
|
|
|
|
// Simulate a response that contains tool_calls
|
|
let tool_calls = Some(vec![ToolCall {
|
|
id: "call_1".to_string(),
|
|
name: "test_tool".to_string(),
|
|
arguments: "{}".to_string(),
|
|
}]);
|
|
|
|
// Check if any tool_call has the expected name
|
|
let supports_tools = tool_calls
|
|
.as_ref()
|
|
.map(|calls| calls.iter().any(|tc| tc.name == "test_tool"))
|
|
.unwrap_or(false);
|
|
|
|
assert!(
|
|
supports_tools,
|
|
"Should detect tool support when response contains test_tool call"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_detect_tool_calling_logic_without_tool_calls() {
|
|
use crate::ai::ToolCall;
|
|
|
|
// Simulate a response without tool_calls
|
|
let tool_calls: Option<Vec<ToolCall>> = None;
|
|
|
|
let supports_tools = tool_calls
|
|
.as_ref()
|
|
.map(|calls| calls.iter().any(|tc| tc.name == "test_tool"))
|
|
.unwrap_or(false);
|
|
|
|
assert!(
|
|
!supports_tools,
|
|
"Should not detect tool support when response has no tool_calls"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_detect_tool_calling_logic_with_wrong_tool_name() {
|
|
use crate::ai::ToolCall;
|
|
|
|
// Simulate a response with tool_calls but wrong tool name
|
|
let tool_calls = Some(vec![ToolCall {
|
|
id: "call_1".to_string(),
|
|
name: "different_tool".to_string(),
|
|
arguments: "{}".to_string(),
|
|
}]);
|
|
|
|
let supports_tools = tool_calls
|
|
.as_ref()
|
|
.map(|calls| calls.iter().any(|tc| tc.name == "test_tool"))
|
|
.unwrap_or(false);
|
|
|
|
assert!(
|
|
!supports_tools,
|
|
"Should not detect tool support when tool name doesn't match"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_detect_tool_calling_error_patterns() {
|
|
// Test error message patterns that indicate tool calling is not supported
|
|
let error_cases = vec![
|
|
"503 Service Unavailable: UNEXPECTED_TOOL_CALL",
|
|
"Tool calling not supported",
|
|
"Function calls are not allowed",
|
|
"tools parameter is invalid",
|
|
];
|
|
|
|
for error_msg in error_cases {
|
|
let msg_lower = error_msg.to_lowercase();
|
|
let is_tool_error = msg_lower.contains("tool")
|
|
|| msg_lower.contains("function")
|
|
|| msg_lower.contains("503");
|
|
|
|
assert!(
|
|
is_tool_error,
|
|
"Error message '{}' should be recognized as tool-related",
|
|
error_msg
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_detect_tool_calling_non_tool_errors() {
|
|
// Test error messages that are NOT tool-related
|
|
let error_cases = vec![
|
|
"Connection timeout",
|
|
"Invalid API key",
|
|
"Rate limit exceeded",
|
|
"Network error",
|
|
];
|
|
|
|
for error_msg in error_cases {
|
|
let msg_lower = error_msg.to_lowercase();
|
|
let is_tool_error = msg_lower.contains("tool")
|
|
|| msg_lower.contains("function")
|
|
|| msg_lower.contains("503");
|
|
|
|
assert!(
|
|
!is_tool_error,
|
|
"Error message '{}' should NOT be recognized as tool-related",
|
|
error_msg
|
|
);
|
|
}
|
|
}
|
|
}
|