a3s-code-core 3.4.0

A3S Code Core - Embeddable AI agent library with tool execution
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

//! Live session extension runtime.
//!
//! A3S Code sessions can be extended after creation with agent definitions and
//! MCP servers. This module owns those dynamic capability mutations so the
//! facade does not need to know how registries, managers, and executors stay in
//! sync.

use super::AgentSession;
use crate::error::Result;
use crate::mcp::{McpServerConfig, McpServerStatus};
use crate::subagent::{AgentDefinition, WorkerAgentSpec};
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;

pub(super) struct SessionExtensionRuntime<'a> {
    session: &'a AgentSession,
}

impl<'a> SessionExtensionRuntime<'a> {
    pub(super) fn from_session(session: &'a AgentSession) -> Self {
        Self { session }
    }

    pub(super) fn register_agent_dir(&self, dir: &Path) -> usize {
        let agents = crate::subagent::load_agents_from_dir(dir);
        let count = agents.len();
        for agent in agents {
            tracing::info!(
                session_id = %self.session.session_id,
                agent = agent.name,
                dir = %dir.display(),
                "Dynamically registered agent"
            );
            self.session.agent_registry.register(agent);
        }
        count
    }

    pub(super) fn register_worker_agent(&self, spec: WorkerAgentSpec) -> AgentDefinition {
        let kind = spec.kind;
        let agent = self.session.agent_registry.register_worker(spec);
        tracing::info!(
            session_id = %self.session.session_id,
            agent = %agent.name,
            kind = ?kind,
            "Dynamically registered worker agent"
        );
        agent
    }

    pub(super) fn register_worker_agents<I>(&self, specs: I) -> Vec<AgentDefinition>
    where
        I: IntoIterator<Item = WorkerAgentSpec>,
    {
        specs
            .into_iter()
            .map(|spec| self.register_worker_agent(spec))
            .collect()
    }

    pub(super) async fn add_mcp_server(&self, config: McpServerConfig) -> Result<usize> {
        let server_name = config.name.clone();
        self.session.mcp_manager.register_server(config).await;
        self.session
            .mcp_manager
            .connect(&server_name)
            .await
            .map_err(|e| crate::error::CodeError::Tool {
                tool: server_name.clone(),
                message: format!("Failed to connect MCP server: {}", e),
            })?;

        let tools = self
            .session
            .mcp_manager
            .get_server_tools(&server_name)
            .await;
        let count = tools.len();

        for tool in crate::mcp::tools::create_mcp_tools(
            &server_name,
            tools,
            Arc::clone(&self.session.mcp_manager),
        ) {
            self.session.tool_executor.register_dynamic_tool(tool);
        }

        tracing::info!(
            session_id = %self.session.session_id,
            server = server_name,
            tools = count,
            "MCP server added to live session"
        );

        Ok(count)
    }

    pub(super) async fn remove_mcp_server(&self, server_name: &str) -> Result<()> {
        self.session
            .tool_executor
            .unregister_tools_by_prefix(&format!("mcp__{server_name}__"));
        self.session
            .mcp_manager
            .disconnect(server_name)
            .await
            .map_err(|e| crate::error::CodeError::Tool {
                tool: server_name.to_string(),
                message: format!("Failed to disconnect MCP server: {}", e),
            })?;
        tracing::info!(
            session_id = %self.session.session_id,
            server = server_name,
            "MCP server removed from live session"
        );
        Ok(())
    }

    pub(super) async fn mcp_status(&self) -> HashMap<String, McpServerStatus> {
        self.session.mcp_manager.get_status().await
    }
}