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::affinity::starts_user_turn;
13use super::util::prompts::{append_note, drop_exact_replay, prepend_system_prompt};
14use super::util::tool_signals::{ToolSemantics, ToolSignals};
15use crate::core::algorithm::{Algorithm, Driver, RoutingIdentity};
16use crate::{
17    CapabilityJudgeConfig, Classifier, LibsyError, LlmClassifierConfig, LlmTaskClassifier, Result,
18    RoutingOutcome, State, TaskClassifierConfig,
19};
20
21/// Default instruction added while the capable model is planning.
22pub const DEFAULT_PLANNING_PROMPT: &str =
23    include_str!("../prompts/plan-execute/planning-system-prompt.md");
24
25const MAX_EXECUTING_SESSIONS: usize = 4_096;
26
27/// Configuration for [`PlanExecute`].
28#[derive(Clone, Debug)]
29pub struct PlanExecuteConfig {
30    /// LLM or decision judge used by the judge trigger.
31    pub judge: Option<TaskClassifierConfig>,
32    /// Additional tool names whose mutations trigger handoff.
33    pub tool_semantics: ToolSemantics,
34    /// System instruction added until handoff.
35    pub planning_prompt: String,
36    /// Optional instruction appended to the handoff request.
37    pub handoff_prompt: Option<String>,
38    /// Replays visible planner reasoning summaries as assistant text at handoff.
39    pub planner_reasoning_as_text: bool,
40}
41
42impl Default for PlanExecuteConfig {
43    fn default() -> Self {
44        Self {
45            judge: None,
46            tool_semantics: ToolSemantics::default(),
47            planning_prompt: DEFAULT_PLANNING_PROMPT.trim().to_string(),
48            handoff_prompt: None,
49            planner_reasoning_as_text: false,
50        }
51    }
52}
53
54#[derive(Clone, Copy, Debug, Eq, PartialEq)]
55enum Phase {
56    Plan,
57    Handoff,
58    Execute,
59}
60
61/// Routes planning turns to the runtime capable model, then latches execution
62/// to the runtime efficient model after an edit or judge approval.
63pub struct PlanExecute {
64    config: PlanExecuteConfig,
65    judge: Option<LlmTaskClassifier>,
66    executing_sessions: Mutex<HashSet<RoutingIdentity>>,
67}
68
69impl PlanExecute {
70    /// Creates a plan/execute router.
71    ///
72    /// Returns an error when a prompt is blank or tool semantics are invalid.
73    pub fn new(config: PlanExecuteConfig) -> Result<Self> {
74        if config.planning_prompt.trim().is_empty() {
75            return Err(LibsyError::AlgorithmError {
76                message: "planning_prompt must not be empty".to_string(),
77            });
78        }
79        if config
80            .handoff_prompt
81            .as_deref()
82            .is_some_and(|prompt| prompt.trim().is_empty())
83        {
84            return Err(LibsyError::AlgorithmError {
85                message: "handoff_prompt must not be empty".to_string(),
86            });
87        }
88        config.tool_semantics.validate()?;
89        let judge = if let Some(mut judge) = config.judge.clone() {
90            judge.recent_turn_window.get_or_insert(10);
91            let prompt = include_str!("../prompts/plan-execute/handoff-decision.md");
92            match &mut judge.judge {
93                CapabilityJudgeConfig::Llm(judge) if judge.contract.prompt().is_none() => {
94                    judge.contract = judge.contract.clone().with_prompt(format!(
95                        "{}\n{prompt}",
96                        include_str!("../prompts/capability-classifier/prompt.md")
97                    ));
98                }
99                CapabilityJudgeConfig::Decision(judge) if judge.instructions.is_none() => {
100                    judge.instructions = Some(serde_json::json!(prompt));
101                }
102                _ => {}
103            }
104            Some(LlmTaskClassifier::new(LlmClassifierConfig::Capability {
105                config: judge,
106            })?)
107        } else {
108            None
109        };
110        Ok(Self {
111            config,
112            judge,
113            executing_sessions: Mutex::new(HashSet::new()),
114        })
115    }
116
117    fn phase(&self, request: &Request, handoff: bool) -> Result<Phase> {
118        let Some(identity) = RoutingIdentity::from_request(request) else {
119            return Ok(if handoff { Phase::Handoff } else { Phase::Plan });
120        };
121        let is_session_final = request
122            .metadata
123            .as_ref()
124            .and_then(|metadata| metadata.session_final)
125            == Some(true);
126
127        let mut sessions = self.executing_sessions.lock();
128        if self.judge.is_some() && starts_user_turn(request) {
129            sessions.remove(&identity);
130        }
131        let phase = if sessions.contains(&identity) {
132            Phase::Execute
133        } else if handoff {
134            if !is_session_final {
135                if sessions.len() >= MAX_EXECUTING_SESSIONS {
136                    return Err(LibsyError::AlgorithmError {
137                        message: format!(
138                            "plan_execute reached its limit of {MAX_EXECUTING_SESSIONS} executing sessions. \
139                             Finish an existing session with session_final before retrying the handoff"
140                        ),
141                    });
142                }
143                sessions.insert(identity.clone());
144            }
145            Phase::Handoff
146        } else {
147            Phase::Plan
148        };
149        if is_session_final {
150            sessions.remove(&identity);
151        }
152        Ok(phase)
153    }
154
155    fn replay_planner_reasoning_as_text(request: &mut Request) -> usize {
156        let mut converted = 0;
157        for message in &mut request.llm_request.messages {
158            message.content = std::mem::take(&mut message.content)
159                .into_iter()
160                .filter_map(|block| match block {
161                    ContentBlock::Reasoning { text, .. } => {
162                        converted += 1;
163                        (!text.is_empty()).then_some(ContentBlock::Text { text })
164                    }
165                    other => Some(other),
166                })
167                .collect();
168        }
169        request
170            .llm_request
171            .messages
172            .retain(|message| !message.content.is_empty());
173        if converted > 0 {
174            drop_exact_replay(request);
175        }
176        converted
177    }
178
179    fn route_to(driver: &Driver, category: Category, request: Request) -> Result<RoutingOutcome> {
180        let models = driver.models_for(&category);
181        let Some((selected, fallbacks)) = models.split_first() else {
182            return Err(LibsyError::AlgorithmError {
183                message: format!("no models available for category {}", category.as_str()),
184            });
185        };
186        Ok(RoutingOutcome::route_to(
187            selected.clone(),
188            fallbacks.to_vec(),
189            request,
190        ))
191    }
192}
193
194#[async_trait::async_trait]
195impl Algorithm for PlanExecute {
196    fn name(&self) -> &str {
197        "plan_execute"
198    }
199
200    async fn route(
201        self: Arc<Self>,
202        driver: Driver,
203        mut request: Request,
204    ) -> Result<RoutingOutcome> {
205        let mut phase = self.phase(&request, false)?;
206        if phase == Phase::Plan {
207            let handoff = if let Some(judge) = &self.judge {
208                if starts_user_turn(&request)
209                    || (driver.models_for(&Category::Judge).is_empty()
210                        && self
211                            .config
212                            .judge
213                            .as_ref()
214                            .is_some_and(|judge| judge.fail_open))
215                {
216                    false
217                } else {
218                    judge
219                        .score(&mut State::default(), &mut request, &driver)
220                        .await?
221                        .0
222                        .argmax(false)?
223                        .is_some_and(|score| score.category == Some(Category::Efficient))
224                }
225            } else {
226                let signals = ToolSignals::from_request_with_semantics(
227                    &request,
228                    None,
229                    &self.config.tool_semantics,
230                );
231                signals.edit_count > 0 || signals.write_count > 0
232            };
233            if handoff {
234                phase = self.phase(&request, true)?;
235            }
236        }
237        match phase {
238            Phase::Plan => {
239                prepend_system_prompt(&mut request, &self.config.planning_prompt);
240                tracing::debug!(phase = "plan", "plan-execute selected capable tier");
241                Self::route_to(&driver, Category::Capable, request)
242            }
243            Phase::Handoff => {
244                let reasoning_converted = if self.config.planner_reasoning_as_text {
245                    Self::replay_planner_reasoning_as_text(&mut request)
246                } else {
247                    0
248                };
249                let prompt_applied = if let Some(prompt) = &self.config.handoff_prompt {
250                    append_note(&mut request, prompt);
251                    true
252                } else {
253                    false
254                };
255                tracing::debug!(
256                    phase = "handoff",
257                    handoff_prompt_applied = prompt_applied,
258                    planner_reasoning_converted = reasoning_converted,
259                    "plan-execute selected efficient tier"
260                );
261                Self::route_to(&driver, Category::Efficient, request)
262            }
263            Phase::Execute => {
264                tracing::debug!(phase = "execute", "plan-execute selected efficient tier");
265                Self::route_to(&driver, Category::Efficient, request)
266            }
267        }
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use std::collections::HashMap;
274    use std::sync::Arc;
275
276    use serde_json::json;
277    use switchyard_protocol::{
278        ContentBlock, LlmRequest, Message, Metadata, ModelId, Request, Role, ToolCall,
279    };
280
281    use super::*;
282    use crate::RuntimeModels;
283    use crate::core::testing::{reply, test_drive_with_models};
284
285    const CAPABLE: &str = "model/capable";
286    const EFFICIENT: &str = "model/efficient";
287
288    fn algorithm(config: PlanExecuteConfig) -> Arc<dyn Algorithm> {
289        Arc::new(PlanExecute::new(config).expect("config should be valid"))
290    }
291
292    fn request(messages: Vec<Message>, session_id: Option<&str>) -> Request {
293        Request {
294            llm_request: LlmRequest {
295                model: Some("switchyard/plan-execute".to_string()),
296                messages,
297                ..LlmRequest::default()
298            },
299            metadata: session_id.map(|session_id| Metadata {
300                session_id: Some(session_id.to_string()),
301                ..Metadata::default()
302            }),
303            ..Request::default()
304        }
305    }
306
307    fn tool_call(name: &str, arguments: serde_json::Value) -> Message {
308        Message {
309            role: Role::Assistant,
310            content: vec![ContentBlock::ToolCall(ToolCall {
311                id: "call-1".to_string(),
312                name: name.to_string(),
313                arguments,
314            })],
315        }
316    }
317
318    fn models() -> RuntimeModels {
319        RuntimeModels::new(HashMap::from([
320            (Category::Capable, vec![ModelId::from(CAPABLE)]),
321            (Category::Efficient, vec![ModelId::from(EFFICIENT)]),
322            (Category::Judge, vec![ModelId::from("judge")]),
323        ]))
324    }
325
326    async fn route_and_capture(
327        algorithm: Arc<dyn Algorithm>,
328        request: Request,
329    ) -> (ModelId, Request) {
330        let captured = Arc::new(Mutex::new(None));
331        let capture = Arc::clone(&captured);
332        let (selected, _) =
333            test_drive_with_models(algorithm, request, models(), move |_target, request| {
334                let capture = Arc::clone(&capture);
335                async move {
336                    *capture.lock() = Some(request);
337                    Ok(reply("ok"))
338                }
339            })
340            .await
341            .expect("routing should succeed");
342        let request = captured
343            .lock()
344            .take()
345            .expect("answer request should be captured");
346        (selected, request)
347    }
348
349    #[tokio::test]
350    async fn judge_handoff_ignores_edits_and_resets_on_user_turns() -> Result<()> {
351        let router = algorithm(PlanExecuteConfig {
352            judge: Some(serde_json::from_value(json!({"base_threshold": 0.5})).unwrap()),
353            ..PlanExecuteConfig::default()
354        });
355        let user = Message::text(Role::User, "Fix the parser");
356        for (turn, probability, expected) in [
357            (user.clone(), 0.9, CAPABLE),
358            (tool_call("apply_patch", json!({})), 0.1, CAPABLE),
359            (tool_call("read_file", json!({})), 0.9, EFFICIENT),
360            (tool_call("read_file", json!({})), 0.1, EFFICIENT),
361            (user.clone(), 0.9, CAPABLE),
362        ] {
363            let request = request(vec![user.clone(), turn], Some("judged"));
364            let verdict = json!({"crux":"remaining work", "primary_rule":"SUP-1",
365                "capability_boundary":"supported", "p_solve":probability})
366            .to_string();
367            let (selected, _) =
368                test_drive_with_models(router.clone(), request, models(), move |_, _| {
369                    std::future::ready(Ok(reply(verdict.clone())))
370                })
371                .await?;
372            assert_eq!(selected, expected);
373        }
374        Ok(())
375    }
376
377    #[tokio::test]
378    async fn wrapped_patch_handoff() {
379        for (input, expected) in [
380            (
381                r#"text(await tools.apply_patch("*** Begin Patch\n*** Add File: a.txt\n+hello\n*** End Patch"))"#,
382                EFFICIENT,
383            ),
384            (
385                r#"const patch = "patch contents"; text(await tools.apply_patch(patch));"#,
386                EFFICIENT,
387            ),
388            ("const patch = 'example only'; text(patch)", CAPABLE),
389            ("// await tools.apply_patch('contents')", CAPABLE),
390            (
391                "if (false) { await tools.apply_patch('contents'); }",
392                CAPABLE,
393            ),
394        ] {
395            let request = request(vec![tool_call("exec", json!({"input": input}))], None);
396            let signals = ToolSignals::from_request(&request, None);
397            let dimensions = crate::algorithms::util::stage::dimensions_from_signal(&signals);
398            assert_eq!(dimensions.production_intensity > 0.0, expected == EFFICIENT);
399            let (selected, _) =
400                route_and_capture(algorithm(PlanExecuteConfig::default()), request).await;
401            assert_eq!(selected, expected, "{input}");
402        }
403    }
404
405    #[tokio::test]
406    async fn plans_then_hands_off_and_latches_execution() {
407        const HANDOFF: &str = "Continue from the plan and repository evidence.";
408        let algorithm = algorithm(PlanExecuteConfig {
409            handoff_prompt: Some(HANDOFF.to_string()),
410            planner_reasoning_as_text: true,
411            ..PlanExecuteConfig::default()
412        });
413
414        let read_only = request(
415            vec![tool_call(
416                "exec_command",
417                json!({"cmd": "rg parser crates"}),
418            )],
419            Some("task-1"),
420        );
421        let (selected, routed) = route_and_capture(Arc::clone(&algorithm), read_only).await;
422        assert_eq!(selected, CAPABLE);
423        assert_eq!(routed.llm_request.instructions.len(), 1);
424
425        let first_edit = request(
426            vec![Message {
427                role: Role::Assistant,
428                content: vec![
429                    ContentBlock::Reasoning {
430                        text: "The parser needs a boundary check.".to_string(),
431                        signature: Some("planner-signature".to_string()),
432                        details: vec![json!({"type": "reasoning.encrypted", "data": "opaque"})],
433                    },
434                    ContentBlock::ToolCall(ToolCall {
435                        id: "call-1".to_string(),
436                        name: "apply_patch".to_string(),
437                        arguments: json!({"patch": "*** Begin Patch"}),
438                    }),
439                ],
440            }],
441            Some("task-1"),
442        );
443        let (selected, routed) = route_and_capture(Arc::clone(&algorithm), first_edit).await;
444        assert_eq!(selected, EFFICIENT);
445        assert_eq!(
446            routed.llm_request.messages[0].content[0],
447            ContentBlock::Text {
448                text: "The parser needs a boundary check.".to_string()
449            }
450        );
451        assert_eq!(
452            routed.llm_request.messages.last(),
453            Some(&Message::text(Role::User, HANDOFF))
454        );
455
456        let mut final_request = request(
457            vec![Message::text(Role::User, "Continue after compaction")],
458            Some("task-1"),
459        );
460        final_request
461            .metadata
462            .as_mut()
463            .expect("session metadata should exist")
464            .session_final = Some(true);
465        let (selected, routed) = route_and_capture(Arc::clone(&algorithm), final_request).await;
466        assert_eq!(selected, EFFICIENT);
467        assert!(routed.llm_request.instructions.is_empty());
468        assert_eq!(routed.llm_request.messages.len(), 1);
469
470        let reused = request(vec![Message::text(Role::User, "New task")], Some("task-1"));
471        let (selected, _) = route_and_capture(algorithm, reused).await;
472        assert_eq!(selected, CAPABLE);
473    }
474
475    #[tokio::test]
476    async fn capacity_preserves_execution_after_compaction() {
477        let algorithm = algorithm(PlanExecuteConfig::default());
478        let sessions: Vec<_> = (0..MAX_EXECUTING_SESSIONS)
479            .map(|index| format!("task-{index}"))
480            .collect();
481        for session in &sessions {
482            let first_edit = request(
483                vec![tool_call("Write", json!({"file_path": "task.py"}))],
484                Some(session),
485            );
486            let (selected, _) = route_and_capture(Arc::clone(&algorithm), first_edit).await;
487            assert_eq!(selected, EFFICIENT);
488        }
489
490        let overflow = request(
491            vec![tool_call("Write", json!({"file_path": "task.py"}))],
492            Some("overflow"),
493        );
494        let (driver, _) = Driver::new("plan_execute", Arc::new(models()));
495        let result = Arc::clone(&algorithm).route(driver, overflow.clone()).await;
496        assert!(matches!(
497            result,
498            Err(LibsyError::AlgorithmError { message })
499                if message.contains("limit of 4096 executing sessions")
500        ));
501
502        for session in &sessions {
503            let mut compacted = request(
504                vec![Message::text(Role::User, "Continue after compaction")],
505                Some(session),
506            );
507            if Some(session) == sessions.last() {
508                compacted.metadata.as_mut().unwrap().session_final = Some(true);
509            }
510            let (selected, routed) = route_and_capture(Arc::clone(&algorithm), compacted).await;
511            assert_eq!(selected, EFFICIENT, "session {session} lost its latch");
512            assert!(routed.llm_request.instructions.is_empty());
513        }
514
515        let (selected, _) = route_and_capture(Arc::clone(&algorithm), overflow).await;
516        assert_eq!(selected, EFFICIENT);
517        let compacted = request(
518            vec![Message::text(Role::User, "Continue after compaction")],
519            Some("overflow"),
520        );
521        let (selected, routed) = route_and_capture(algorithm, compacted).await;
522        assert_eq!(selected, EFFICIENT);
523        assert!(routed.llm_request.instructions.is_empty());
524    }
525
526    #[tokio::test]
527    async fn final_handoff_does_not_evict_an_executing_session_at_capacity() {
528        let algorithm = Arc::new(PlanExecute::new(PlanExecuteConfig::default()).unwrap());
529        algorithm.executing_sessions.lock().extend(
530            (0..MAX_EXECUTING_SESSIONS)
531                .map(|index| RoutingIdentity::Session(format!("task-{index}"))),
532        );
533        let mut final_request = request(
534            vec![tool_call("Write", json!({"file_path": "task.py"}))],
535            Some("final-handoff"),
536        );
537        final_request.metadata.as_mut().unwrap().session_final = Some(true);
538
539        let (selected, routed) = route_and_capture(algorithm.clone(), final_request).await;
540        assert_eq!(selected, EFFICIENT);
541        assert!(routed.llm_request.instructions.is_empty());
542        let sessions = algorithm.executing_sessions.lock();
543        assert_eq!(sessions.len(), MAX_EXECUTING_SESSIONS);
544        assert!(!sessions.contains(&RoutingIdentity::Session("final-handoff".to_string())));
545    }
546
547    #[tokio::test]
548    async fn mutation_without_a_session_uses_the_efficient_tier() {
549        let messages = vec![tool_call(
550            "exec_command",
551            json!({"cmd": "printf 'done\\n' > task.txt"}),
552        )];
553
554        let (selected, routed) = route_and_capture(
555            algorithm(PlanExecuteConfig::default()),
556            request(messages.clone(), None),
557        )
558        .await;
559
560        assert_eq!(selected, EFFICIENT);
561        assert_eq!(routed.llm_request.messages, messages);
562        assert!(routed.llm_request.instructions.is_empty());
563    }
564
565    #[tokio::test]
566    async fn editor_view_keeps_planning() {
567        let messages = vec![tool_call(
568            "str_replace_based_edit_tool",
569            json!({"command": "view", "path": "/app/main.py"}),
570        )];
571
572        let (selected, _) = route_and_capture(
573            algorithm(PlanExecuteConfig::default()),
574            request(messages, None),
575        )
576        .await;
577
578        assert_eq!(selected, CAPABLE);
579    }
580
581    #[test]
582    fn rejects_blank_prompts() {
583        for config in [
584            PlanExecuteConfig {
585                planning_prompt: "  ".to_string(),
586                ..PlanExecuteConfig::default()
587            },
588            PlanExecuteConfig {
589                handoff_prompt: Some("  ".to_string()),
590                ..PlanExecuteConfig::default()
591            },
592        ] {
593            assert!(matches!(
594                PlanExecute::new(config),
595                Err(LibsyError::AlgorithmError { .. })
596            ));
597        }
598    }
599}