telos_agent/orchestration/subagent/
fork.rs

1//! Fork engine — lightweight concurrent multi-perspective execution.
2//!
3//! Each "lens" shares the parent session's provider, tools, and messages but
4//! gets its own system prompt + task. All lenses run concurrently via a
5//! tokio [`Semaphore`]-bounded [`Synapse`].
6//!
7//! Fork is NOT a subprocess. It is an in-process concurrent provider call.
8
9use futures_util::future::join_all;
10use serde_json::Value;
11use std::sync::Arc;
12use tokio::sync::Semaphore;
13
14use crate::config::AgentConfig;
15use crate::knowledge::tasks::{Task, TaskManager, TaskStatus};
16use crate::model::message::Message;
17use crate::model::provider::{CompletionRequest, ModelProvider};
18use crate::tools::api::ToolRegistry;
19
20/// Shared state across all fork lenses — cheap to clone.
21pub struct ForkShared {
22    pub provider: Arc<dyn ModelProvider + Send + Sync>,
23    pub tool_registry: ToolRegistry,
24    pub messages: Arc<Vec<Message>>,
25    pub config: AgentConfig,
26}
27
28/// A single lens — one perspective on the shared context.
29#[derive(Debug, Clone)]
30pub struct ForkLens {
31    /// Label for logging and tracking.
32    pub lens: String,
33    /// Injected as the system prompt for this lens.
34    pub system_prompt: String,
35    /// The specific task for this lens.
36    pub task: String,
37    /// Optional JSON Schema for structured output.
38    pub output_schema: Option<Value>,
39    /// Tools available to this lens (if empty, uses all registry tools).
40    pub allowed_tools: Vec<String>,
41}
42
43/// Result from a single lens execution.
44#[derive(Debug, Clone)]
45pub enum ForkResult {
46    Text(String),
47    Structured(Value),
48}
49
50/// Result of a complete fork execution.
51pub struct ForkExecution {
52    pub results: Vec<Option<ForkResult>>,
53    pub task_ids: Vec<String>,
54}
55
56/// Lightweight concurrency limiter for fork lens execution.
57pub struct Synapse {
58    semaphore: Arc<Semaphore>,
59}
60
61impl Synapse {
62    pub fn new(max_concurrent: usize) -> Self {
63        Synapse { semaphore: Arc::new(Semaphore::new(max_concurrent.max(1))) }
64    }
65
66    /// Run all lenses concurrently, respecting the concurrency limit.
67    /// Each lens gets a single provider call (not a full turn loop).
68    /// If a `TaskManager` is provided, one task is created per lens and
69    /// updated on completion.
70    pub async fn run_all(
71        &self,
72        shared: &ForkShared,
73        lenses: Vec<ForkLens>,
74        task_manager: Option<&TaskManager>,
75    ) -> ForkExecution {
76        let task_ids: Vec<String> = if let Some(tm) = &task_manager {
77            lenses
78                .iter()
79                .map(|lens| {
80                    let id = uuid_v4();
81                    let task = Task {
82                        id: id.clone(),
83                        subject: format!("fork lens: {}", lens.lens),
84                        description: lens.task.clone(),
85                        status: TaskStatus::InProgress,
86                        blocked_by: vec![],
87                        blocks: vec![],
88                        output: None,
89                        kind: None,
90                        agent_id: None,
91                        agent_type: None,
92                        worktree_path: None,
93                        error: None,
94                    };
95                    tm.create(task);
96                    id
97                })
98                .collect()
99        } else {
100            vec![]
101        };
102
103        let results = join_all(lenses.into_iter().map(|lens| {
104            let sem = self.semaphore.clone();
105            let shared = shared.clone();
106            async move {
107                let _permit = sem.acquire().await.ok()?;
108                execute_lens(&shared, &lens).await
109            }
110        }))
111        .await;
112
113        // Update task status based on execution results
114        if let Some(tm) = &task_manager {
115            for (i, _result) in results.iter().enumerate() {
116                let status = TaskStatus::Completed;
117                if let Some(task_id) = task_ids.get(i) {
118                    tm.update(task_id, status);
119                }
120            }
121        }
122
123        ForkExecution { results, task_ids }
124    }
125}
126
127/// Generate a unique task ID using the current timestamp.
128fn uuid_v4() -> String {
129    use std::time::{SystemTime, UNIX_EPOCH};
130    let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default();
131    format!("task_{:x}", now.as_nanos())
132}
133
134impl Clone for ForkShared {
135    fn clone(&self) -> Self {
136        Self {
137            provider: self.provider.clone(),
138            tool_registry: self.tool_registry.clone(),
139            messages: self.messages.clone(),
140            config: self.config.clone(),
141        }
142    }
143}
144
145/// Execute a single lens: build context summary, call provider, extract result.
146async fn execute_lens(shared: &ForkShared, lens: &ForkLens) -> Option<ForkResult> {
147    // Build context summary from shared messages (last few messages only)
148    let context_summary = build_context_summary(&shared.messages);
149
150    let fork_messages = vec![
151        Message::system(&lens.system_prompt),
152        Message::user(format!("## Context\n\n{context_summary}\n\n## Task\n\n{}", lens.task)),
153    ];
154
155    // Filter tools if allowed_tools specified
156    let tools = if lens.allowed_tools.is_empty() {
157        shared.tool_registry.definitions()
158    } else {
159        shared
160            .tool_registry
161            .definitions()
162            .into_iter()
163            .filter(|d| lens.allowed_tools.contains(&d.name))
164            .collect()
165    };
166
167    let request = CompletionRequest {
168        system_prompt_blocks: vec![], // system message already in fork_messages
169        messages: fork_messages,
170        tools,
171        model_hint: None,
172        max_tokens: None,
173    };
174
175    let response = match shared.provider.complete(request).await {
176        Ok(r) => r,
177        Err(_) => return None,
178    };
179
180    let text = response.message.text_content();
181
182    // If output_schema specified, try to extract structured JSON
183    if let Some(_schema) = &lens.output_schema {
184        if let Ok(val) = serde_json::from_str::<Value>(&text) {
185            Some(ForkResult::Structured(val))
186        } else {
187            // Fallback: wrap text
188            Some(ForkResult::Text(text))
189        }
190    } else {
191        Some(ForkResult::Text(text))
192    }
193}
194
195/// Build a compact context summary from recent messages.
196fn build_context_summary(messages: &[Message]) -> String {
197    let recent: Vec<&Message> = messages.iter().rev().take(6).collect();
198    if recent.is_empty() {
199        return String::new();
200    }
201    let mut parts = vec!["## Recent Conversation".to_string()];
202    for msg in recent.iter().rev() {
203        let text = msg.text_content();
204        if text.is_empty() {
205            continue;
206        }
207        let truncated: String = text.chars().take(2000).collect();
208        parts.push(format!("[{}]: {}", role_str(&msg.role), truncated));
209    }
210    parts.join("\n")
211}
212
213fn role_str(role: &crate::model::message::Role) -> &str {
214    match role {
215        crate::model::message::Role::System => "system",
216        crate::model::message::Role::User => "user",
217        crate::model::message::Role::Assistant => "assistant",
218        crate::model::message::Role::Tool => "tool",
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use crate::config::AgentConfig;
226    use crate::model::mock::MockProvider;
227
228    #[tokio::test]
229    async fn synapse_runs_multiple_lenses_concurrently() {
230        let provider = Arc::new(MockProvider::new(vec![]));
231        let tool_registry = ToolRegistry::new();
232        let messages = Arc::new(vec![Message::user("Original request")]);
233
234        let shared = ForkShared {
235            provider,
236            tool_registry: tool_registry.clone(),
237            messages,
238            config: AgentConfig::default(),
239        };
240
241        let lenses = vec![
242            ForkLens {
243                lens: "a".into(),
244                system_prompt: "You are lens A".into(),
245                task: "Task A".into(),
246                output_schema: None,
247                allowed_tools: vec![],
248            },
249            ForkLens {
250                lens: "b".into(),
251                system_prompt: "You are lens B".into(),
252                task: "Task B".into(),
253                output_schema: None,
254                allowed_tools: vec![],
255            },
256        ];
257
258        let synapse = Synapse::new(2);
259        let execution = synapse.run_all(&shared, lenses, None).await;
260        assert_eq!(execution.results.len(), 2);
261        // Both lenses should fail gracefully (no mock responses)
262        assert!(execution.results.iter().all(|r| r.is_none()));
263    }
264
265    #[test]
266    fn build_context_summary_truncates_long_messages() {
267        let long_text = "x".repeat(5000);
268        let msgs = vec![Message::user(&long_text), Message::assistant("short response")];
269        let summary = build_context_summary(&msgs);
270        assert!(summary.contains("short response"));
271        assert!(!summary.contains(&long_text));
272    }
273}
telos_agent/orchestration/subagent/fork.rs

telos_agent/orchestration/subagent/
fork.rs
