1use 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
21pub 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#[derive(Clone, Debug)]
29pub struct PlanExecuteConfig {
30 pub judge: Option<TaskClassifierConfig>,
32 pub tool_semantics: ToolSemantics,
34 pub planning_prompt: String,
36 pub handoff_prompt: Option<String>,
38 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
61pub struct PlanExecute {
64 config: PlanExecuteConfig,
65 judge: Option<LlmTaskClassifier>,
66 executing_sessions: Mutex<HashSet<RoutingIdentity>>,
67}
68
69impl PlanExecute {
70 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}