Skip to main content

switchyard_libsy/algorithms/
plan_execute.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Plans coding tasks on a capable model, then hands execution to an efficient model.
5
6use std::collections::HashSet;
7use std::sync::Arc;
8
9use parking_lot::Mutex;
10use switchyard_protocol::{Category, ContentBlock, Request};
11
12use super::util::prompts::{append_note, drop_exact_replay, prepend_system_prompt};
13use super::util::tool_signals::ToolSignals;
14use crate::core::algorithm::{Algorithm, Driver, RoutingIdentity};
15use crate::{LibsyError, Result, RoutingOutcome};
16
17/// Default instruction added while the capable model is planning.
18pub const DEFAULT_PLANNING_PROMPT: &str =
19    include_str!("../prompts/plan-execute/planning-system-prompt.md");
20
21const MAX_EXECUTING_SESSIONS: usize = 4_096;
22
23/// Configuration for [`PlanExecute`].
24#[derive(Clone, Debug)]
25pub struct PlanExecuteConfig {
26    /// System instruction added until the first edit or write tool call.
27    pub planning_prompt: String,
28    /// Optional instruction appended to the handoff request.
29    pub handoff_prompt: Option<String>,
30    /// Replays visible planner reasoning summaries as assistant text at handoff.
31    pub planner_reasoning_as_text: bool,
32}
33
34impl Default for PlanExecuteConfig {
35    fn default() -> Self {
36        Self {
37            planning_prompt: DEFAULT_PLANNING_PROMPT.trim().to_string(),
38            handoff_prompt: None,
39            planner_reasoning_as_text: false,
40        }
41    }
42}
43
44#[derive(Clone, Copy, Debug, Eq, PartialEq)]
45enum Phase {
46    Plan,
47    Handoff,
48    Execute,
49}
50
51/// Routes planning turns to the runtime capable model, then latches execution
52/// to the runtime efficient model after the first recorded mutation.
53pub struct PlanExecute {
54    config: PlanExecuteConfig,
55    executing_sessions: Mutex<HashSet<RoutingIdentity>>,
56}
57
58impl PlanExecute {
59    /// Creates a plan/execute router.
60    ///
61    /// Returns an error when either configured prompt is blank.
62    pub fn new(config: PlanExecuteConfig) -> Result<Self> {
63        if config.planning_prompt.trim().is_empty() {
64            return Err(LibsyError::AlgorithmError {
65                message: "planning_prompt must not be empty".to_string(),
66            });
67        }
68        if config
69            .handoff_prompt
70            .as_deref()
71            .is_some_and(|prompt| prompt.trim().is_empty())
72        {
73            return Err(LibsyError::AlgorithmError {
74                message: "handoff_prompt must not be empty".to_string(),
75            });
76        }
77        Ok(Self {
78            config,
79            executing_sessions: Mutex::new(HashSet::new()),
80        })
81    }
82
83    fn phase(&self, request: &Request) -> Phase {
84        let signals = ToolSignals::from_request(request, None);
85        let mutation_seen = signals.edit_count > 0 || signals.write_count > 0;
86        let Some(identity) = RoutingIdentity::from_request(request) else {
87            return if mutation_seen {
88                Phase::Handoff
89            } else {
90                Phase::Plan
91            };
92        };
93
94        let mut sessions = self.executing_sessions.lock();
95        let phase = if sessions.contains(&identity) {
96            Phase::Execute
97        } else if mutation_seen {
98            if sessions.len() >= MAX_EXECUTING_SESSIONS
99                && let Some(evicted) = sessions.iter().next().cloned()
100            {
101                sessions.remove(&evicted);
102            }
103            sessions.insert(identity.clone());
104            Phase::Handoff
105        } else {
106            Phase::Plan
107        };
108        if request
109            .metadata
110            .as_ref()
111            .and_then(|metadata| metadata.session_final)
112            == Some(true)
113        {
114            sessions.remove(&identity);
115        }
116        phase
117    }
118
119    fn replay_planner_reasoning_as_text(request: &mut Request) -> usize {
120        let mut converted = 0;
121        for message in &mut request.llm_request.messages {
122            message.content = std::mem::take(&mut message.content)
123                .into_iter()
124                .filter_map(|block| match block {
125                    ContentBlock::Reasoning { text, .. } => {
126                        converted += 1;
127                        (!text.is_empty()).then_some(ContentBlock::Text { text })
128                    }
129                    other => Some(other),
130                })
131                .collect();
132        }
133        request
134            .llm_request
135            .messages
136            .retain(|message| !message.content.is_empty());
137        if converted > 0 {
138            drop_exact_replay(request);
139        }
140        converted
141    }
142
143    fn route_to(driver: &Driver, category: Category, request: Request) -> Result<RoutingOutcome> {
144        let models = driver.models_for(&category);
145        let Some((selected, fallbacks)) = models.split_first() else {
146            return Err(LibsyError::AlgorithmError {
147                message: format!("no models available for category {}", category.as_str()),
148            });
149        };
150        Ok(RoutingOutcome::route_to(
151            selected.clone(),
152            fallbacks.to_vec(),
153            request,
154        ))
155    }
156}
157
158#[async_trait::async_trait]
159impl Algorithm for PlanExecute {
160    fn name(&self) -> &str {
161        "plan_execute"
162    }
163
164    fn needs_history_replay(&self, _request: &Request) -> bool {
165        true
166    }
167
168    async fn route(
169        self: Arc<Self>,
170        driver: Driver,
171        mut request: Request,
172    ) -> Result<RoutingOutcome> {
173        match self.phase(&request) {
174            Phase::Plan => {
175                prepend_system_prompt(&mut request, &self.config.planning_prompt);
176                tracing::debug!(phase = "plan", "plan-execute selected capable tier");
177                Self::route_to(&driver, Category::Capable, request)
178            }
179            Phase::Handoff => {
180                let reasoning_converted = if self.config.planner_reasoning_as_text {
181                    Self::replay_planner_reasoning_as_text(&mut request)
182                } else {
183                    0
184                };
185                let prompt_applied = if let Some(prompt) = &self.config.handoff_prompt {
186                    append_note(&mut request, prompt);
187                    true
188                } else {
189                    false
190                };
191                tracing::debug!(
192                    phase = "handoff",
193                    handoff_prompt_applied = prompt_applied,
194                    planner_reasoning_converted = reasoning_converted,
195                    "plan-execute selected efficient tier"
196                );
197                Self::route_to(&driver, Category::Efficient, request)
198            }
199            Phase::Execute => {
200                tracing::debug!(phase = "execute", "plan-execute selected efficient tier");
201                Self::route_to(&driver, Category::Efficient, request)
202            }
203        }
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use std::collections::HashMap;
210    use std::sync::Arc;
211
212    use serde_json::json;
213    use switchyard_protocol::{
214        ContentBlock, LlmRequest, Message, Metadata, ModelId, Request, Role, ToolCall,
215    };
216
217    use super::*;
218    use crate::RuntimeModels;
219    use crate::core::testing::{reply, test_drive_with_models};
220
221    const CAPABLE: &str = "model/capable";
222    const EFFICIENT: &str = "model/efficient";
223
224    fn algorithm(config: PlanExecuteConfig) -> Arc<dyn Algorithm> {
225        Arc::new(PlanExecute::new(config).expect("config should be valid"))
226    }
227
228    fn request(messages: Vec<Message>, session_id: Option<&str>) -> Request {
229        Request {
230            llm_request: LlmRequest {
231                model: Some("switchyard/plan-execute".to_string()),
232                messages,
233                ..LlmRequest::default()
234            },
235            metadata: session_id.map(|session_id| Metadata {
236                session_id: Some(session_id.to_string()),
237                ..Metadata::default()
238            }),
239            ..Request::default()
240        }
241    }
242
243    fn tool_call(name: &str, arguments: serde_json::Value) -> Message {
244        Message {
245            role: Role::Assistant,
246            content: vec![ContentBlock::ToolCall(ToolCall {
247                id: "call-1".to_string(),
248                name: name.to_string(),
249                arguments,
250            })],
251        }
252    }
253
254    fn models() -> RuntimeModels {
255        RuntimeModels::new(HashMap::from([
256            (Category::Capable, vec![ModelId::from(CAPABLE)]),
257            (Category::Efficient, vec![ModelId::from(EFFICIENT)]),
258        ]))
259    }
260
261    async fn route_and_capture(
262        algorithm: Arc<dyn Algorithm>,
263        request: Request,
264    ) -> (ModelId, Request) {
265        let captured = Arc::new(Mutex::new(None));
266        let capture = Arc::clone(&captured);
267        let (selected, _) =
268            test_drive_with_models(algorithm, request, models(), move |_target, request| {
269                let capture = Arc::clone(&capture);
270                async move {
271                    *capture.lock() = Some(request);
272                    Ok(reply("ok"))
273                }
274            })
275            .await
276            .expect("routing should succeed");
277        let request = captured
278            .lock()
279            .take()
280            .expect("answer request should be captured");
281        (selected, request)
282    }
283
284    #[tokio::test]
285    async fn plans_then_hands_off_and_latches_execution() {
286        const HANDOFF: &str = "Continue from the plan and repository evidence.";
287        let algorithm = algorithm(PlanExecuteConfig {
288            handoff_prompt: Some(HANDOFF.to_string()),
289            planner_reasoning_as_text: true,
290            ..PlanExecuteConfig::default()
291        });
292
293        let read_only = request(
294            vec![tool_call(
295                "exec_command",
296                json!({"cmd": "rg parser crates"}),
297            )],
298            Some("task-1"),
299        );
300        let (selected, routed) = route_and_capture(Arc::clone(&algorithm), read_only).await;
301        assert_eq!(selected, CAPABLE);
302        assert_eq!(routed.llm_request.instructions.len(), 1);
303
304        let first_edit = request(
305            vec![Message {
306                role: Role::Assistant,
307                content: vec![
308                    ContentBlock::Reasoning {
309                        text: "The parser needs a boundary check.".to_string(),
310                        signature: Some("planner-signature".to_string()),
311                        details: vec![json!({"type": "reasoning.encrypted", "data": "opaque"})],
312                    },
313                    ContentBlock::ToolCall(ToolCall {
314                        id: "call-1".to_string(),
315                        name: "apply_patch".to_string(),
316                        arguments: json!({"patch": "*** Begin Patch"}),
317                    }),
318                ],
319            }],
320            Some("task-1"),
321        );
322        let (selected, routed) = route_and_capture(Arc::clone(&algorithm), first_edit).await;
323        assert_eq!(selected, EFFICIENT);
324        assert_eq!(
325            routed.llm_request.messages[0].content[0],
326            ContentBlock::Text {
327                text: "The parser needs a boundary check.".to_string()
328            }
329        );
330        assert_eq!(
331            routed.llm_request.messages.last(),
332            Some(&Message::text(Role::User, HANDOFF))
333        );
334
335        let mut final_request = request(
336            vec![Message::text(Role::User, "Continue after compaction")],
337            Some("task-1"),
338        );
339        final_request
340            .metadata
341            .as_mut()
342            .expect("session metadata should exist")
343            .session_final = Some(true);
344        let (selected, routed) = route_and_capture(Arc::clone(&algorithm), final_request).await;
345        assert_eq!(selected, EFFICIENT);
346        assert!(routed.llm_request.instructions.is_empty());
347        assert_eq!(routed.llm_request.messages.len(), 1);
348
349        let reused = request(vec![Message::text(Role::User, "New task")], Some("task-1"));
350        let (selected, _) = route_and_capture(algorithm, reused).await;
351        assert_eq!(selected, CAPABLE);
352    }
353
354    #[tokio::test]
355    async fn mutation_without_a_session_uses_the_efficient_tier() {
356        let messages = vec![tool_call(
357            "exec_command",
358            json!({"cmd": "printf 'done\\n' > task.txt"}),
359        )];
360
361        let (selected, routed) = route_and_capture(
362            algorithm(PlanExecuteConfig::default()),
363            request(messages.clone(), None),
364        )
365        .await;
366
367        assert_eq!(selected, EFFICIENT);
368        assert_eq!(routed.llm_request.messages, messages);
369        assert!(routed.llm_request.instructions.is_empty());
370    }
371
372    #[tokio::test]
373    async fn editor_view_keeps_planning() {
374        let messages = vec![tool_call(
375            "str_replace_based_edit_tool",
376            json!({"command": "view", "path": "/app/main.py"}),
377        )];
378
379        let (selected, _) = route_and_capture(
380            algorithm(PlanExecuteConfig::default()),
381            request(messages, None),
382        )
383        .await;
384
385        assert_eq!(selected, CAPABLE);
386    }
387
388    #[test]
389    fn rejects_blank_prompts() {
390        for config in [
391            PlanExecuteConfig {
392                planning_prompt: "  ".to_string(),
393                ..PlanExecuteConfig::default()
394            },
395            PlanExecuteConfig {
396                handoff_prompt: Some("  ".to_string()),
397                ..PlanExecuteConfig::default()
398            },
399        ] {
400            assert!(matches!(
401                PlanExecute::new(config),
402                Err(LibsyError::AlgorithmError { .. })
403            ));
404        }
405    }
406}