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//! Starts coding tasks on a capable planner, then hands execution to an efficient model.
5
6use std::collections::HashSet;
7use std::sync::Arc;
8
9use parking_lot::Mutex;
10use switchyard_protocol::{ModelId, Request};
11
12use super::util::prompts::{SystemPromptProcessor, TargetPrompts};
13use super::util::tool_signals::ToolSignals;
14use crate::core::algorithm::{Algorithm, Driver, RoutingIdentity};
15use crate::core::processor::{Event, Processor};
16use crate::{LibsyError, Result, RoutingOutcome};
17
18/// Default instruction prepended while the capable model is planning.
19pub const DEFAULT_PLANNING_PROMPT: &str =
20    include_str!("../prompts/plan-execute/planning-system-prompt.md");
21
22/// Maximum session latches retained by one router instance.
23const MAX_EXECUTING_SESSIONS: usize = 4_096;
24
25/// Configuration for [`PlanExecute`].
26#[derive(Clone, Debug)]
27pub struct PlanExecuteConfig {
28    /// System instruction prepended until the first edit or write tool call.
29    pub planning_prompt: String,
30}
31
32impl Default for PlanExecuteConfig {
33    fn default() -> Self {
34        Self {
35            planning_prompt: DEFAULT_PLANNING_PROMPT.trim().to_string(),
36        }
37    }
38}
39
40/// Routes planning turns to a capable model and all turns after the first edit
41/// to an efficient model while preserving the caller's full trajectory.
42pub struct PlanExecute {
43    capable: ModelId,
44    efficient: ModelId,
45    planning_prompt: SystemPromptProcessor,
46    executing_sessions: Mutex<HashSet<RoutingIdentity>>,
47}
48
49impl PlanExecute {
50    /// Creates a plan/execute router.
51    ///
52    /// Returns an error when the planning prompt is empty.
53    pub fn new(capable: ModelId, efficient: ModelId, config: PlanExecuteConfig) -> Result<Self> {
54        if config.planning_prompt.trim().is_empty() {
55            return Err(LibsyError::AlgorithmError {
56                message: "planning_prompt must not be empty".to_string(),
57            });
58        }
59        let planning_prompt = SystemPromptProcessor::new(
60            TargetPrompts::default().with(capable.clone(), config.planning_prompt),
61        );
62        Ok(Self {
63            capable,
64            efficient,
65            planning_prompt,
66            executing_sessions: Mutex::new(HashSet::new()),
67        })
68    }
69
70    /// Whether this request is in execution, latching the transition for keyed sessions.
71    fn is_executing(&self, request: &Request) -> bool {
72        let signals = ToolSignals::from_request(request, None);
73        let mutation_seen = signals.edit_count > 0 || signals.write_count > 0;
74        let Some(identity) = RoutingIdentity::from_request(request) else {
75            return mutation_seen;
76        };
77
78        let mut sessions = self.executing_sessions.lock();
79        let executing = if mutation_seen {
80            if sessions.len() >= MAX_EXECUTING_SESSIONS
81                && !sessions.contains(&identity)
82                && let Some(evicted) = sessions.iter().next().cloned()
83            {
84                sessions.remove(&evicted);
85            }
86            sessions.insert(identity.clone());
87            true
88        } else {
89            sessions.contains(&identity)
90        };
91        if request
92            .metadata
93            .as_ref()
94            .and_then(|metadata| metadata.session_final)
95            == Some(true)
96        {
97            sessions.remove(&identity);
98        }
99        executing
100    }
101}
102
103#[async_trait::async_trait]
104impl Algorithm for PlanExecute {
105    fn name(&self) -> &str {
106        "plan_execute"
107    }
108
109    async fn route(
110        self: Arc<Self>,
111        _driver: Driver,
112        mut request: Request,
113    ) -> Result<RoutingOutcome> {
114        if self.is_executing(&request) {
115            tracing::info!(target = %self.efficient, phase = "execute", "plan-execute selected target");
116            Ok(RoutingOutcome::route_to(
117                self.efficient.clone(),
118                Vec::new(),
119                request,
120            ))
121        } else {
122            self.planning_prompt
123                .process(
124                    &mut (),
125                    Event::Decision {
126                        request: &mut request,
127                        selected_model_id: &self.capable,
128                    },
129                )
130                .await?;
131            tracing::info!(target = %self.capable, phase = "plan", "plan-execute selected target");
132            Ok(RoutingOutcome::route_to(
133                self.capable.clone(),
134                Vec::new(),
135                request,
136            ))
137        }
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use std::sync::{Arc, Mutex};
144
145    use serde_json::json;
146    use switchyard_protocol::{
147        ContentBlock, InstructionBlock, LlmRequest, Message, Metadata, Request, Role, ToolCall,
148    };
149
150    use super::*;
151    use crate::core::testing::{reply, test_drive};
152
153    fn algorithm() -> Arc<dyn Algorithm> {
154        Arc::new(
155            PlanExecute::new(
156                ModelId::from("model/capable"),
157                ModelId::from("model/efficient"),
158                PlanExecuteConfig::default(),
159            )
160            .expect("default config should be valid"),
161        )
162    }
163
164    fn request(messages: Vec<Message>, session_id: Option<&str>) -> Request {
165        Request {
166            llm_request: LlmRequest {
167                model: Some("switchyard/plan-execute".to_string()),
168                messages,
169                ..LlmRequest::default()
170            },
171            metadata: session_id.map(|session_id| Metadata {
172                session_id: Some(session_id.to_string()),
173                ..Metadata::default()
174            }),
175            ..Request::default()
176        }
177    }
178
179    fn tool_call(name: &str, arguments: serde_json::Value) -> Message {
180        Message {
181            role: Role::Assistant,
182            content: vec![ContentBlock::ToolCall(ToolCall {
183                id: "call-1".to_string(),
184                name: name.to_string(),
185                arguments,
186            })],
187        }
188    }
189
190    async fn route_and_capture(
191        algorithm: Arc<dyn Algorithm>,
192        request: Request,
193    ) -> (ModelId, Request) {
194        let captured = Arc::new(Mutex::new(None));
195        let capture = Arc::clone(&captured);
196        let (selected, _) = test_drive(algorithm, request, move |_target, request| {
197            let capture = Arc::clone(&capture);
198            async move {
199                *capture.lock().expect("capture lock should be available") = Some(request);
200                Ok(reply("ok"))
201            }
202        })
203        .await
204        .expect("routing should succeed");
205        let request = captured
206            .lock()
207            .expect("capture lock should be available")
208            .take()
209            .expect("answer request should be captured");
210        (selected, request)
211    }
212
213    #[tokio::test]
214    async fn initial_turn_uses_capable_model_with_planning_prefix() {
215        let messages = vec![Message::text(Role::User, "fix the parser")];
216        let (selected, routed) =
217            route_and_capture(algorithm(), request(messages.clone(), Some("task-1"))).await;
218
219        assert_eq!(selected, "model/capable");
220        assert_eq!(routed.llm_request.messages, messages);
221        assert_eq!(routed.llm_request.instructions.len(), 1);
222        assert_eq!(routed.llm_request.instructions[0].role, Role::System);
223        assert_eq!(
224            routed.llm_request.instructions[0].content,
225            vec![ContentBlock::Text {
226                text: DEFAULT_PLANNING_PROMPT.trim().to_string()
227            }]
228        );
229    }
230
231    #[tokio::test]
232    async fn read_only_tool_calls_remain_in_planning() {
233        let messages = vec![
234            Message::text(Role::User, "fix the parser"),
235            tool_call("exec_command", json!({"cmd": "rg parser crates"})),
236        ];
237
238        let (selected, routed) =
239            route_and_capture(algorithm(), request(messages.clone(), None)).await;
240
241        assert_eq!(selected, "model/capable");
242        assert_eq!(routed.llm_request.messages, messages);
243        assert_eq!(routed.llm_request.instructions.len(), 1);
244    }
245
246    #[tokio::test]
247    async fn first_edit_switches_to_efficient_and_keeps_the_trajectory() {
248        let messages = vec![
249            Message::text(Role::User, "fix the parser"),
250            Message::text(Role::Assistant, "I will update the parser now."),
251            tool_call("apply_patch", json!({"patch": "*** Begin Patch"})),
252        ];
253        let mut input = request(messages.clone(), Some("task-2"));
254        input.llm_request.instructions.push(InstructionBlock {
255            role: Role::Developer,
256            content: vec![ContentBlock::Text {
257                text: "keep the public API stable".to_string(),
258            }],
259        });
260
261        let (selected, routed) = route_and_capture(algorithm(), input).await;
262
263        assert_eq!(selected, "model/efficient");
264        assert_eq!(routed.llm_request.messages, messages);
265        assert_eq!(routed.llm_request.instructions.len(), 1);
266        assert_eq!(routed.llm_request.instructions[0].role, Role::Developer);
267    }
268
269    #[tokio::test]
270    async fn shell_file_write_switches_to_execution() {
271        let messages = vec![tool_call(
272            "exec_command",
273            json!({"cmd": "python -c 'from pathlib import Path; Path(\"x\").write_text(\"y\")'"}),
274        )];
275
276        let (selected, routed) = route_and_capture(algorithm(), request(messages, None)).await;
277
278        assert_eq!(selected, "model/efficient");
279        assert!(routed.llm_request.instructions.is_empty());
280    }
281
282    #[tokio::test]
283    async fn shell_redirection_switches_to_execution() {
284        let messages = vec![tool_call(
285            "exec_command",
286            json!({"cmd": "printf 'completed\\n' > task.txt"}),
287        )];
288
289        let (selected, routed) = route_and_capture(algorithm(), request(messages, None)).await;
290
291        assert_eq!(selected, "model/efficient");
292        assert!(routed.llm_request.instructions.is_empty());
293    }
294
295    #[tokio::test]
296    async fn execution_latches_by_session_after_history_compaction() {
297        let algorithm = algorithm();
298        let edit = request(
299            vec![tool_call("write_file", json!({"path": "src/lib.rs"}))],
300            Some("task-3"),
301        );
302        let (selected, _) = route_and_capture(Arc::clone(&algorithm), edit).await;
303        assert_eq!(selected, "model/efficient");
304
305        let compacted = request(
306            vec![Message::text(
307                Role::User,
308                "Continue from the compacted summary",
309            )],
310            Some("task-3"),
311        );
312        let (selected, routed) = route_and_capture(algorithm, compacted).await;
313
314        assert_eq!(selected, "model/efficient");
315        assert!(routed.llm_request.instructions.is_empty());
316    }
317
318    #[tokio::test]
319    async fn final_request_uses_then_releases_the_session_latch() {
320        let algorithm = algorithm();
321        let edit = request(
322            vec![tool_call("write_file", json!({"path": "src/lib.rs"}))],
323            Some("task-4"),
324        );
325        let (selected, _) = route_and_capture(Arc::clone(&algorithm), edit).await;
326        assert_eq!(selected, "model/efficient");
327
328        let mut final_request = request(vec![Message::text(Role::User, "Finish")], Some("task-4"));
329        final_request
330            .metadata
331            .as_mut()
332            .expect("session metadata should exist")
333            .session_final = Some(true);
334        let (selected, _) = route_and_capture(Arc::clone(&algorithm), final_request).await;
335        assert_eq!(selected, "model/efficient");
336
337        let reused = request(vec![Message::text(Role::User, "New task")], Some("task-4"));
338        let (selected, _) = route_and_capture(algorithm, reused).await;
339        assert_eq!(selected, "model/capable");
340    }
341
342    #[test]
343    fn empty_planning_prompt_is_rejected() {
344        let result = PlanExecute::new(
345            ModelId::from("model/capable"),
346            ModelId::from("model/efficient"),
347            PlanExecuteConfig {
348                planning_prompt: "  ".to_string(),
349            },
350        );
351
352        assert!(matches!(result, Err(LibsyError::AlgorithmError { .. })));
353    }
354}
switchyard_libsy/algorithms/plan_execute.rs

switchyard_libsy/algorithms/
plan_execute.rs
