Skip to main content

switchyard_libsy/algorithms/util/
escalation.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Trajectory-judge components for the escalation router — the judges, their shared verdict
5//! policy, and the transcript condenser they read.
6//!
7//! [`build_judge`] (the trouble judge) and [`build_recovery_judge`] (the hand-back judge
8//! consulted on latched turns) are the whole surface; the confirmation policy that consumes
9//! their verdicts lives with the assembled algorithm in [`crate::algorithms::escalation`].
10
11use serde::Deserialize;
12use switchyard_protocol::{ContentBlock, Message, Role};
13
14use super::classifier_contract::{ClassifierContract, ClassifierContractConfig};
15use super::llm_judge::{
16    ClassifierInput, JudgeClassifier, JudgePolicy, JudgeRuntimeConfig, SerdeDecoder,
17    StructuredJudge,
18};
19use crate::core::algorithm::LlmTarget;
20use crate::core::classifier::{Classification, Score};
21use crate::core::state::State;
22use crate::{LibsyError, Result};
23use switchyard_protocol::Request;
24
25const PROMPT_TEMPLATE: &str = include_str!("../../prompts/escalation/prompt.md");
26const RECOVERY_PROMPT_TEMPLATE: &str = include_str!("../../prompts/escalation/recovery_prompt.md");
27const SCHEMA_TEMPLATE: &str = include_str!("../../prompts/escalation/schema.json");
28
29/// Separator marking where [`truncate_middle`] dropped a message's interior.
30const TRIM_MARKER: &str = " ...[trimmed] ";
31
32/// Suffix marking a transcript cut off by [`MAX_REQUEST_CHARS`].
33const TRUNCATION_SUFFIX: &str = "...<truncated>";
34
35/// Per-message cap for system and developer anchors, which carry no trajectory signal but
36/// which coding-agent harnesses make very large.
37const SYSTEM_CHARS: usize = 1_000;
38
39/// Cap for the first user message — the task statement, so it gets the widest anchor budget.
40const FIRST_USER_CHARS: usize = 2_000;
41
42/// Backstop on the assembled transcript; the per-message caps normally bind first.
43const MAX_REQUEST_CHARS: usize = 18_000;
44
45/// The tuning surface for the trajectory judge.
46///
47/// The routing settings retain their benchmarked defaults. Everything else is a fixed invariant
48/// (the constants above).
49#[derive(Clone, Debug, Deserialize)]
50#[serde(default, deny_unknown_fields)]
51pub struct EscalationJudgeConfig {
52    /// Consecutive escalate verdicts required before a turn moves to the capable tier, which
53    /// is also the turn that latches the session. Any decline clears the streak.
54    /// `1` escalates on the first verdict; the router's main cost dial.
55    /// `2` or higher needs a session id, since the streak is retained per session.
56    pub confirmations: u32,
57    /// Consecutive clear verdicts required, while latched, before the session de-latches back
58    /// to the efficient tier. `0` (the default) disables recovery: a latch is permanent for
59    /// the session's remainder and latched turns skip the judge entirely. Any escalate verdict
60    /// clears the recovery streak, and a de-latched session can re-escalate as usual.
61    pub recovery_confirmations: u32,
62    /// Trailing messages shown on top of the anchors. A loop longer than this is invisible.
63    pub recent_turn_window: usize,
64    /// Per-message cap inside the trailing window.
65    pub window_message_chars: usize,
66}
67
68impl EscalationJudgeConfig {
69    /// Rejects settings that would leave the judge with nothing useful to read.
70    fn validate(&self) -> Result<()> {
71        let reject = |message: String| Err(LibsyError::AlgorithmError { message });
72        if self.confirmations == 0 {
73            return reject("confirmations must be at least 1".to_string());
74        }
75        if self.recent_turn_window == 0 {
76            return reject("recent_turn_window must be at least 1".to_string());
77        }
78        if self.window_message_chars < 50 {
79            return reject(format!(
80                "window_message_chars must be at least 50, got {}",
81                self.window_message_chars
82            ));
83        }
84        Ok(())
85    }
86}
87
88impl Default for EscalationJudgeConfig {
89    fn default() -> Self {
90        Self {
91            confirmations: 2,
92            recovery_confirmations: 0,
93            recent_turn_window: 28,
94            window_message_chars: 500,
95        }
96    }
97}
98
99/// The judge's verdict. The schema also requires a `reason`, which makes the judge state its
100/// case and measurably sharpens the verdict — routing reads only the boolean, so it is
101/// deserialized away rather than carried.
102#[derive(Deserialize)]
103pub(crate) struct EscalationVerdict {
104    escalate: bool,
105}
106
107/// Builds the condensed trajectory presented to the escalation judge.
108pub(crate) struct EscalationInput {
109    config: EscalationJudgeConfig,
110}
111
112impl ClassifierInput for EscalationInput {
113    fn build_messages(&self, _state: &State, request: &Request) -> Vec<Message> {
114        let messages = &request.llm_request.messages;
115        let summary = summarize_for_judge(messages, conversation_turn(request), &self.config);
116        vec![Message::text(Role::User, summary)]
117    }
118}
119
120/// Structured trajectory judge with a typed escalation verdict.
121pub(crate) type EscalationJudge = StructuredJudge<EscalationInput, SerdeDecoder<EscalationVerdict>>;
122
123/// Maps the judge's verdict to a classification. A verdict names the tier to serve — capable
124/// on escalate, efficient on decline — so the caller reads it straight off the winning score.
125/// [`Classification::Ambiguous`] carries the unavailable case, which names no tier: both a
126/// decline and an outage stay efficient, but only a decline is evidence, so only a decline
127/// clears the streak.
128pub(crate) struct EscalationPolicy {
129    capable: String,
130    efficient: String,
131}
132
133impl JudgePolicy for EscalationPolicy {
134    type Verdict = EscalationVerdict;
135
136    fn to_classification(&self, verdict: Option<&EscalationVerdict>) -> Classification {
137        match verdict {
138            Some(verdict) if verdict.escalate => Classification::Scores(vec![Score {
139                target: self.capable.clone(),
140                confidence: 1.0,
141            }]),
142            Some(_) => Classification::Scores(vec![Score {
143                target: self.efficient.clone(),
144                confidence: 1.0,
145            }]),
146            None => Classification::Ambiguous(Vec::new()),
147        }
148    }
149}
150
151/// Builds the trajectory judge over `judge_target`, scoring `capable` when it escalates.
152///
153/// Loads the packaged prompt and schema, so an unusable asset or an unusable `config` value
154/// fails here rather than on the first request.
155pub(crate) fn build_judge(
156    judge_target: LlmTarget,
157    capable: String,
158    efficient: String,
159    contract_config: &ClassifierContractConfig,
160    config: EscalationJudgeConfig,
161    max_output_tokens: u64,
162) -> Result<JudgeClassifier<EscalationJudge, EscalationPolicy>> {
163    let contract =
164        ClassifierContract::from_config(contract_config, PROMPT_TEMPLATE, SCHEMA_TEMPLATE)?;
165    assemble_judge(
166        judge_target,
167        capable,
168        efficient,
169        contract,
170        config,
171        max_output_tokens,
172    )
173}
174
175/// Builds the hand-back judge consulted on latched turns when recovery is enabled.
176///
177/// The recovery question — could the efficient tier carry the *remaining* work — is
178/// different from the trouble question the trajectory judge answers, and the latched
179/// transcript it reads is the capable tier's own (usually healthy-looking) work. It
180/// therefore always uses the packaged recovery prompt: the route-level `prompt` override
181/// replaces the trouble rubric only. Same verdict schema, so `escalate: true` keeps the
182/// session on the capable tier.
183pub(crate) fn build_recovery_judge(
184    judge_target: LlmTarget,
185    capable: String,
186    efficient: String,
187    config: EscalationJudgeConfig,
188    max_output_tokens: u64,
189) -> Result<JudgeClassifier<EscalationJudge, EscalationPolicy>> {
190    let contract = ClassifierContract::from_config(
191        &ClassifierContractConfig::default(),
192        RECOVERY_PROMPT_TEMPLATE,
193        SCHEMA_TEMPLATE,
194    )?;
195    assemble_judge(
196        judge_target,
197        capable,
198        efficient,
199        contract,
200        config,
201        max_output_tokens,
202    )
203}
204
205/// Validates `config` and assembles a structured judge over the rendered `contract`.
206fn assemble_judge(
207    judge_target: LlmTarget,
208    capable: String,
209    efficient: String,
210    contract: ClassifierContract,
211    config: EscalationJudgeConfig,
212    max_output_tokens: u64,
213) -> Result<JudgeClassifier<EscalationJudge, EscalationPolicy>> {
214    config.validate()?;
215    Ok(JudgeClassifier::new(
216        StructuredJudge::new(
217            EscalationInput { config },
218            contract,
219            SerdeDecoder::new(),
220            JudgeRuntimeConfig::new(max_output_tokens)?,
221        ),
222        judge_target,
223        EscalationPolicy { capable, efficient },
224    ))
225}
226
227/// The 1-indexed model invocation the transcript ends on: one per assistant reply.
228///
229/// The judge reads the turn *including* the reply it is judging, so the newest assistant
230/// message is this turn — no `+ 1`. Counting the caller's request instead would report the
231/// turn after the one under judgement.
232///
233/// Messages are already normalized by `switchyard-protocol`, so this needs no
234/// per-format branching.
235pub(crate) fn conversation_turn(request: &Request) -> usize {
236    request
237        .llm_request
238        .messages
239        .iter()
240        .filter(|message| message.role == Role::Assistant)
241        .count()
242}
243
244/// Flattens a message to plain text, tool calls and tool results included.
245///
246/// [`Message::text_content`] is deliberately not used here: it keeps only text and refusal
247/// blocks, which would erase exactly the repeated-command signal the judge's loop detection
248/// relies on.
249fn message_text(message: &Message) -> String {
250    let mut parts = Vec::new();
251    collect_text(&message.content, &mut parts);
252    parts.join(" ")
253}
254
255/// Appends the judge-relevant text of each block, descending into tool results.
256fn collect_text(content: &[ContentBlock], parts: &mut Vec<String>) {
257    for block in content {
258        match block {
259            ContentBlock::Text { text } | ContentBlock::Refusal { text } => {
260                parts.push(text.clone());
261            }
262            ContentBlock::ToolCall(call) => {
263                parts.push(format!("tool_call {}({})", call.name, call.arguments));
264            }
265            ContentBlock::ToolResult(result) => collect_text(&result.content, parts),
266            _ => {}
267        }
268    }
269}
270
271/// Keeps the head and tail of `text` within `limit` characters.
272///
273/// The head gets two thirds of the surviving budget: for a trajectory judge the command or
274/// error signature that opens a message carries more signal than its trailing output.
275fn truncate_middle(text: &str, limit: usize) -> String {
276    let chars: Vec<char> = text.chars().collect();
277    if chars.len() <= limit {
278        return text.to_string();
279    }
280    let keep = limit
281        .saturating_sub(TRIM_MARKER.chars().count())
282        .max(20)
283        .min(chars.len());
284    let head = keep * 2 / 3;
285    let tail = keep - head;
286    let mut out: String = chars[..head].iter().collect();
287    out.push_str(TRIM_MARKER);
288    out.extend(chars[chars.len() - tail..].iter());
289    out
290}
291
292/// Renders a compact role-labelled transcript for the judge.
293///
294/// The framing anchors — system/developer messages and the first user message, where agent
295/// harnesses put the task statement — are kept unconditionally and capped individually. The
296/// trailing window carries recent activity. A coverage header states how much history is not
297/// shown, so the judge can reason about pace rather than assuming it sees everything.
298///
299/// When the assembled text still exceeds `max_request_chars`, the oldest window lines go
300/// first: for a trajectory judge the newest evidence is strictly the most valuable.
301fn summarize_for_judge(
302    messages: &[Message],
303    turn: usize,
304    config: &EscalationJudgeConfig,
305) -> String {
306    let mut anchors: Vec<String> = Vec::new();
307    let mut window: Vec<String> = Vec::new();
308    let mut first_user_seen = false;
309
310    for message in messages {
311        let text = message_text(message);
312        match message.role {
313            Role::System | Role::Developer => anchors.push(format!(
314                "[{}] {}",
315                role_label(message.role),
316                truncate_middle(&text, SYSTEM_CHARS)
317            )),
318            Role::User if !first_user_seen => {
319                first_user_seen = true;
320                anchors.push(format!(
321                    "[user (task)] {}",
322                    truncate_middle(&text, FIRST_USER_CHARS)
323                ));
324            }
325            role => window.push(format!(
326                "[{}] {}",
327                role_label(role),
328                truncate_middle(&text, config.window_message_chars)
329            )),
330        }
331    }
332
333    if window.len() > config.recent_turn_window {
334        window.drain(..window.len() - config.recent_turn_window);
335    }
336
337    let assemble = |window: &[String]| {
338        let header = format!(
339            "Conversation turn {turn}; showing the last {} of {} messages after the task framing.",
340            window.len(),
341            messages.len(),
342        );
343        std::iter::once(header)
344            .chain(anchors.iter().cloned())
345            .chain(window.iter().cloned())
346            .collect::<Vec<_>>()
347            .join("\n")
348    };
349
350    let mut text = assemble(&window);
351    while text.chars().count() > MAX_REQUEST_CHARS && !window.is_empty() {
352        window.remove(0);
353        text = assemble(&window);
354    }
355    if text.chars().count() > MAX_REQUEST_CHARS {
356        let keep = MAX_REQUEST_CHARS.saturating_sub(TRUNCATION_SUFFIX.chars().count() + 1);
357        text = text.chars().take(keep).collect::<String>() + TRUNCATION_SUFFIX;
358    }
359    text
360}
361
362/// The transcript label for a role.
363fn role_label(role: Role) -> &'static str {
364    match role {
365        Role::System => "system",
366        Role::Developer => "developer",
367        Role::User => "user",
368        Role::Assistant => "assistant",
369        Role::Tool => "tool",
370    }
371}
372
373/// A request whose conversation sits at `turn`: `turn - 1` prior assistant replies, each
374/// answered by a further user message.
375///
376/// Shared with the assembled router's tests, which drive the same conversation shape.
377#[cfg(test)]
378pub(crate) fn request_at_turn(session_id: Option<&str>, turn: usize) -> Request {
379    use switchyard_protocol::{LlmRequest, Metadata};
380
381    let mut messages = vec![Message::text(Role::User, "What is 2+2?")];
382    for attempt in 1..turn {
383        messages.push(Message::text(Role::Assistant, format!("attempt {attempt}")));
384        messages.push(Message::text(Role::User, format!("still wrong {attempt}")));
385    }
386    Request {
387        llm_request: LlmRequest {
388            model: Some("auto".to_string()),
389            messages,
390            ..LlmRequest::default()
391        },
392        raw_request: None,
393        metadata: session_id.map(|id| Metadata {
394            session_id: Some(id.to_string()),
395            ..Metadata::default()
396        }),
397    }
398}
399
400#[cfg(test)]
401mod tests {
402    use serde_json::json;
403    use switchyard_protocol::{ContentBlock, Message, Role, ToolCall, ToolResult};
404
405    use super::*;
406    use crate::algorithms::util::llm_judge::Judge;
407
408    fn escalation_judge(max_output_tokens: u64) -> Result<EscalationJudge> {
409        Ok(StructuredJudge::new(
410            EscalationInput {
411                config: EscalationJudgeConfig::default(),
412            },
413            ClassifierContract::from_config(
414                &ClassifierContractConfig::default(),
415                PROMPT_TEMPLATE,
416                SCHEMA_TEMPLATE,
417            )?,
418            SerdeDecoder::new(),
419            JudgeRuntimeConfig::new(max_output_tokens)?,
420        ))
421    }
422
423    #[test]
424    fn judge_request_is_rubric_plus_summary_under_a_completion_cap() -> Result<()> {
425        let judge = escalation_judge(super::super::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)?;
426
427        // As the classifier calls it: the turn's reply is already on the transcript.
428        let mut judged = request_at_turn(None, 4);
429        judged
430            .llm_request
431            .messages
432            .push(Message::text(Role::Assistant, "this turn's reply"));
433        let built = judge.build_request(&State::default(), &judged);
434
435        // Rubric in instructions, condensed trajectory as the sole user message.
436        assert_eq!(built.llm_request.instructions.len(), 1);
437        assert_eq!(built.llm_request.instructions[0].role, Role::System);
438        assert_eq!(built.llm_request.messages.len(), 1);
439        assert_eq!(built.llm_request.messages[0].role, Role::User);
440        assert!(
441            built.llm_request.messages[0]
442                .text_content("")
443                .is_some_and(|text| text.contains("Conversation turn 4"))
444        );
445        // Bounded output, so a reasoning judge cannot run away mid-verdict.
446        assert_eq!(
447            built.llm_request.output.max_output_tokens,
448            Some(super::super::DEFAULT_JUDGE_MAX_OUTPUT_TOKENS)
449        );
450        assert!(built.llm_request.output.response_format.is_some());
451        Ok(())
452    }
453
454    #[test]
455    fn judge_request_uses_the_configured_completion_cap() -> Result<()> {
456        let judge = escalation_judge(512)?;
457
458        let built = judge.build_request(&State::default(), &request_at_turn(None, 1));
459
460        assert_eq!(built.llm_request.output.max_output_tokens, Some(512));
461        Ok(())
462    }
463
464    #[test]
465    fn conversation_turn_counts_assistant_replies() {
466        // The judge is handed the transcript with this turn's reply already appended, which
467        // is the shape asserted here: a request entering turn N plus its reply *is* turn N.
468        for turn in [1, 5] {
469            let mut judged = request_at_turn(None, turn);
470            judged
471                .llm_request
472                .messages
473                .push(Message::text(Role::Assistant, "this turn's reply"));
474            assert_eq!(conversation_turn(&judged), turn);
475        }
476    }
477
478    #[test]
479    fn message_text_keeps_tool_calls_and_results() {
480        let call = Message {
481            role: Role::Assistant,
482            content: vec![
483                ContentBlock::Text {
484                    text: "running it".to_string(),
485                },
486                ContentBlock::ToolCall(ToolCall {
487                    id: "call-1".to_string(),
488                    name: "bash".to_string(),
489                    arguments: json!({"cmd": "ls"}),
490                }),
491            ],
492        };
493        let text = message_text(&call);
494        assert!(text.contains("running it"), "{text}");
495        assert!(text.contains(r#"tool_call bash({"cmd":"ls"})"#), "{text}");
496
497        let result = Message {
498            role: Role::Tool,
499            content: vec![ContentBlock::ToolResult(ToolResult {
500                tool_call_id: "call-1".to_string(),
501                content: vec![ContentBlock::Text {
502                    text: "no such file".to_string(),
503                }],
504                is_error: Some(true),
505            })],
506        };
507        assert_eq!(message_text(&result), "no such file");
508    }
509
510    #[test]
511    fn truncate_middle_keeps_head_and_tail() {
512        let text = "a".repeat(40) + &"z".repeat(40);
513        let trimmed = truncate_middle(&text, 50);
514        assert!(trimmed.chars().count() <= 50, "{trimmed}");
515        assert!(trimmed.starts_with('a'));
516        assert!(trimmed.ends_with('z'));
517        assert!(trimmed.contains("[trimmed]"));
518
519        // Under the limit the text is returned untouched.
520        assert_eq!(truncate_middle("short", 50), "short");
521    }
522
523    #[test]
524    fn summary_keeps_anchors_and_the_recent_window() {
525        let mut messages = vec![
526            Message::text(Role::System, "you are a coding agent"),
527            Message::text(Role::User, "fix the failing test"),
528        ];
529        for i in 0..10 {
530            messages.push(Message::text(Role::Assistant, format!("step {i}")));
531        }
532        let config = EscalationJudgeConfig {
533            recent_turn_window: 3,
534            ..EscalationJudgeConfig::default()
535        };
536
537        let summary = summarize_for_judge(&messages, 11, &config);
538
539        assert!(
540            summary.contains("[system] you are a coding agent"),
541            "{summary}"
542        );
543        assert!(
544            summary.contains("[user (task)] fix the failing test"),
545            "{summary}"
546        );
547        assert!(summary.contains("Conversation turn 11; showing the last 3 of 12 messages"));
548        // Only the newest window entries survive.
549        assert!(summary.contains("step 9"), "{summary}");
550        assert!(summary.contains("step 7"), "{summary}");
551        assert!(!summary.contains("step 6"), "{summary}");
552    }
553
554    #[test]
555    fn summary_drops_oldest_window_lines_under_the_char_cap() {
556        // MAX_REQUEST_CHARS is a backstop, not a dial: at default settings the window caps
557        // bind first (28 x 500 plus anchors sits under it), so reaching it takes an unusually
558        // wide per-message cap. That is the point — it only fires on pathological input.
559        let mut messages = vec![
560            Message::text(Role::System, "framing"),
561            Message::text(Role::User, "task"),
562        ];
563        for i in 0..20 {
564            messages.push(Message::text(
565                Role::Assistant,
566                format!("{i} {}", "x".repeat(2_000)),
567            ));
568        }
569        let config = EscalationJudgeConfig {
570            window_message_chars: 2_000,
571            ..EscalationJudgeConfig::default()
572        };
573
574        let summary = summarize_for_judge(&messages, 21, &config);
575
576        assert!(
577            summary.chars().count() <= MAX_REQUEST_CHARS,
578            "{}",
579            summary.chars().count()
580        );
581        // Anchors are never dropped, and the newest activity outlives the oldest.
582        assert!(summary.contains("[system] framing"), "{summary}");
583        assert!(summary.contains("[user (task)] task"), "{summary}");
584        assert!(summary.contains("19 xxx"), "{summary}");
585        assert!(!summary.contains("0 xxx"), "{summary}");
586    }
587}