use anyhow::{Context as _, Result};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use ai_session::context::{MessageRole, SessionContext};
use ai_session::coordination::{AgentId, AgentMessage, MessageBus};
use ai_session::core::SessionId;
use ai_session::output::{OutputParser, ParsedOutput};
use ai_session::persistence::PersistenceManager;
use crate::identity::AgentIdentity;
use crate::providers::claude_code::ClaudeCodeExecutor;
use crate::providers::{ClaudeCodeConfig, ProviderExecutor};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BridgeResult {
pub raw: String,
pub parsed: ParsedOutput,
pub success: bool,
}
pub struct AISessionBridge {
executor: ClaudeCodeExecutor,
context_histories: DashMap<String, SessionContext>,
message_bus: Arc<MessageBus>,
output_parser: OutputParser,
persistence: PersistenceManager,
agent_mappings: DashMap<String, AgentId>,
}
impl AISessionBridge {
pub fn new(config: ClaudeCodeConfig, storage_path: PathBuf) -> Self {
Self {
executor: ClaudeCodeExecutor::new(config),
context_histories: DashMap::new(),
message_bus: Arc::new(MessageBus::new()),
output_parser: OutputParser::new(),
persistence: PersistenceManager::new(storage_path),
agent_mappings: DashMap::new(),
}
}
pub fn with_executor(executor: ClaudeCodeExecutor, storage_path: PathBuf) -> Self {
Self {
executor,
context_histories: DashMap::new(),
message_bus: Arc::new(MessageBus::new()),
output_parser: OutputParser::new(),
persistence: PersistenceManager::new(storage_path),
agent_mappings: DashMap::new(),
}
}
pub fn register_agent(&self, agent_id: &str) -> Result<()> {
let session_id = SessionId::new();
let context = SessionContext::new(session_id);
self.context_histories.insert(agent_id.to_string(), context);
let ai_agent_id = AgentId::new();
self.message_bus
.register_agent(ai_agent_id.clone())
.context("Failed to register agent on message bus")?;
self.agent_mappings
.insert(agent_id.to_string(), ai_agent_id);
tracing::info!("Registered agent '{}' with AISessionBridge", agent_id);
Ok(())
}
pub async fn execute_task(
&self,
agent_id: &str,
prompt: &str,
identity: &AgentIdentity,
working_dir: &Path,
) -> Result<BridgeResult> {
let raw_output: String =
ProviderExecutor::execute_prompt(&self.executor, prompt, identity, working_dir)
.await
.context("Claude Code CLI execution failed")?;
let parsed = self
.output_parser
.parse(&raw_output)
.unwrap_or(ParsedOutput::PlainText(raw_output.clone()));
let success = is_parsed_success(&parsed);
if let Some(mut context) = self.context_histories.get_mut(agent_id) {
context.add_message_raw(MessageRole::User, prompt.to_string());
context.add_message_raw(MessageRole::Assistant, raw_output.clone());
context.compress_context().await;
}
if let Some(ai_agent_id) = self.agent_mappings.get(agent_id) {
let task_id = ai_session::coordination::TaskId::new();
let msg = AgentMessage::TaskCompleted {
agent_id: ai_agent_id.clone(),
task_id,
result: serde_json::json!({
"agent": agent_id,
"success": success,
"output_preview": truncate_output(&raw_output, 500),
}),
};
let _ = self.message_bus.broadcast(
ai_agent_id.clone(),
ai_session::BroadcastMessage {
id: uuid::Uuid::new_v4(),
from: ai_agent_id.clone(),
content: serde_json::to_string(&msg).unwrap_or_default(),
priority: ai_session::MessagePriority::Normal,
timestamp: chrono::Utc::now(),
},
);
}
if let Some(context) = self.context_histories.get(agent_id) {
let session_id = context.session_id.clone();
let state = ai_session::persistence::SessionState {
session_id: session_id.clone(),
config: ai_session::SessionConfig::default(),
status: ai_session::SessionStatus::Running,
context: context.clone(),
command_history: Vec::new(),
metadata: ai_session::persistence::SessionMetadata::default(),
};
if let Err(e) = self.persistence.save_session(&session_id, &state).await {
tracing::warn!("Failed to persist session state for {}: {}", agent_id, e);
}
}
Ok(BridgeResult {
raw: raw_output,
parsed,
success,
})
}
pub fn get_compression_stats(
&self,
agent_id: &str,
) -> Option<ai_session::context::CompressionStats> {
self.context_histories
.get(agent_id)
.map(|ctx| ctx.get_compression_stats())
}
pub fn message_bus(&self) -> &Arc<MessageBus> {
&self.message_bus
}
pub fn agent_count(&self) -> usize {
self.context_histories.len()
}
pub fn get_recent_context(&self, agent_id: &str, n: usize) -> Vec<String> {
self.context_histories
.get(agent_id)
.map(|ctx| {
ctx.get_recent_messages(n)
.into_iter()
.map(|m| format!("[{:?}] {}", m.role, m.content))
.collect()
})
.unwrap_or_default()
}
}
fn is_parsed_success(parsed: &ParsedOutput) -> bool {
match parsed {
ParsedOutput::PlainText(_) => true, ParsedOutput::CodeExecution { .. } => true,
ParsedOutput::BuildOutput { status, .. } => {
matches!(status, ai_session::output::BuildStatus::Success)
}
ParsedOutput::TestResults { failed, .. } => *failed == 0,
ParsedOutput::StructuredLog { level, .. } => {
!matches!(level, ai_session::output::LogLevel::Error)
}
}
}
fn truncate_output(output: &str, max_len: usize) -> String {
if output.len() <= max_len {
output.to_string()
} else {
format!("{}...", &output[..max_len])
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
#[test]
fn test_bridge_creation() {
let config = ClaudeCodeConfig::default();
let bridge = AISessionBridge::new(config, PathBuf::from("/tmp/ccswarm-test"));
assert_eq!(bridge.agent_count(), 0);
}
#[test]
fn test_agent_registration() {
let config = ClaudeCodeConfig::default();
let bridge = AISessionBridge::new(config, PathBuf::from("/tmp/ccswarm-test"));
bridge.register_agent("frontend-agent").unwrap();
assert_eq!(bridge.agent_count(), 1);
}
#[test]
fn test_truncate_output() {
assert_eq!(truncate_output("short", 10), "short");
assert_eq!(truncate_output("this is longer text", 10), "this is lo...");
}
#[test]
fn test_is_parsed_success() {
assert!(is_parsed_success(&ParsedOutput::PlainText(
"ok".to_string()
)));
assert!(!is_parsed_success(&ParsedOutput::TestResults {
passed: 5,
failed: 1,
details: ai_session::output::TestDetails {
suite: Some("cargo".to_string()),
duration: Some(std::time::Duration::from_secs(0)),
failed_tests: vec!["test_one".to_string()],
},
}));
}
}