switchyard_libsy/algorithms/util/
turn_pin.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Holds one classifier verdict across the tool-call turns that follow a user message.
5
6use std::sync::Arc;
7
8use async_trait::async_trait;
9use serde::Deserialize;
10use switchyard_protocol::{ContentBlock, Message, ModelId, Request, Response, Role};
11
12use super::decisive;
13use crate::Result;
14use crate::core::algorithm::Driver;
15use crate::core::classifier::{Classification, Classifier};
16use crate::core::state::{State, StateValue};
17
18/// `State.extra` key holding the pinned target.
19const PINNED_TARGET_KEY: &str = "classifier_pinned_target";
20
21/// How often the classifier re-decides a session's target.
22#[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq)]
23#[serde(rename_all = "snake_case")]
24pub enum ClassifyTrigger {
25    /// Judge every request, tool continuations included.
26    #[default]
27    EveryRequest,
28    /// Judge each new user message and hold that target across the tool calls between.
29    UserTurn,
30    /// Judge once and reuse that target for the session.
31    NewSession,
32}
33
34/// True when the user spoke. Anthropic carries tool results as a `Role::User` message, so
35/// role alone cannot tell a human turn from a tool continuation. The same invariant is
36/// encoded in the Anthropic codec's `message_is_tool_result_only`.
37fn is_user_turn(message: &Message) -> bool {
38    message.role == Role::User
39        && !message
40            .content
41            .iter()
42            .all(|block| matches!(block, ContentBlock::ToolResult(_)))
43}
44
45/// True when the conversation ends with the user speaking, so the agent is not mid-task.
46fn starts_new_turn(messages: &[Message]) -> bool {
47    messages.last().is_some_and(is_user_turn)
48}
49
50fn pinned_target(state: &State) -> Option<ModelId> {
51    match state.extra.get(PINNED_TARGET_KEY) {
52        Some(StateValue::String(target)) => Some(ModelId::new(target.clone())),
53        _ => None,
54    }
55}
56
57/// Pins the inner classifier's verdict until the user speaks again.
58///
59/// Without a session id there is no retained state, and every turn is classified.
60pub(crate) struct TurnPin {
61    inner: Arc<dyn Classifier<State>>,
62}
63
64impl TurnPin {
65    pub(crate) fn new(inner: Arc<dyn Classifier<State>>) -> Self {
66        Self { inner }
67    }
68}
69
70#[async_trait]
71impl Classifier<State> for TurnPin {
72    fn routing_tier(&self, selected_model_id: &ModelId) -> Option<&'static str> {
73        self.inner.routing_tier(selected_model_id)
74    }
75
76    async fn score(
77        &self,
78        state: &mut State,
79        request: &mut Request,
80        driver: Option<&Driver>,
81    ) -> Result<(Classification, Option<Response>)> {
82        if let Some(target) = pinned_target(state)
83            && !starts_new_turn(&request.llm_request.messages)
84        {
85            return Ok((decisive(&target), None));
86        }
87
88        let (classification, response) = self.inner.score(state, request, driver).await?;
89        // An abstention clears the pin. Holding the old target would serve this turn from the
90        // fall-through default while its tool continuations reused the previous turn's target.
91        match classification.argmax(false)? {
92            Some(score) => {
93                state.extra.insert(
94                    PINNED_TARGET_KEY.to_string(),
95                    StateValue::String(score.target.as_str().to_string()),
96                );
97            }
98            None => {
99                state.extra.remove(PINNED_TARGET_KEY);
100            }
101        }
102        Ok((classification, response))
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    use parking_lot::Mutex;
111    use serde_json::Value;
112    use switchyard_protocol::{LlmRequest, ToolCall, ToolResult};
113
114    /// Returns queued verdicts in order, counting each consultation.
115    struct RecordingClassifier {
116        verdicts: Mutex<Vec<Option<&'static str>>>,
117        consultations: Mutex<u32>,
118    }
119
120    impl RecordingClassifier {
121        fn new(verdicts: Vec<Option<&'static str>>) -> Arc<Self> {
122            Arc::new(Self {
123                verdicts: Mutex::new(verdicts.into_iter().rev().collect()),
124                consultations: Mutex::new(0),
125            })
126        }
127
128        fn consultations(&self) -> u32 {
129            *self.consultations.lock()
130        }
131    }
132
133    #[async_trait]
134    impl Classifier<State> for RecordingClassifier {
135        async fn score(
136            &self,
137            _state: &mut State,
138            _request: &mut Request,
139            _driver: Option<&Driver>,
140        ) -> Result<(Classification, Option<Response>)> {
141            *self.consultations.lock() += 1;
142            let verdict = self.verdicts.lock().pop().flatten();
143            Ok((
144                match verdict {
145                    Some(target) => decisive(&ModelId::new(target)),
146                    None => Classification::Scores(Vec::new()),
147                },
148                None,
149            ))
150        }
151    }
152
153    fn user(text: &str) -> Message {
154        Message::text(Role::User, text)
155    }
156
157    fn tool_result(id: &str) -> Message {
158        Message {
159            role: Role::User,
160            content: vec![ContentBlock::ToolResult(ToolResult {
161                tool_call_id: id.to_string(),
162                content: Vec::new(),
163                is_error: None,
164            })],
165        }
166    }
167
168    fn tool_call(id: &str) -> Message {
169        Message {
170            role: Role::Assistant,
171            content: vec![ContentBlock::ToolCall(ToolCall {
172                id: id.to_string(),
173                name: "read".to_string(),
174                arguments: Value::Null,
175            })],
176        }
177    }
178
179    async fn selected(
180        pin: &TurnPin,
181        state: &mut State,
182        messages: Vec<Message>,
183    ) -> Result<Option<String>> {
184        let mut request = Request {
185            llm_request: LlmRequest {
186                messages,
187                ..LlmRequest::default()
188            },
189            raw_request: None,
190            metadata: None,
191        };
192        let (classification, _) = pin.score(state, &mut request, None).await?;
193        Ok(classification
194            .argmax(false)?
195            .map(|score| score.target.as_str().to_string()))
196    }
197
198    #[tokio::test]
199    async fn tool_continuation_turns_hold_the_pinned_target() -> Result<()> {
200        let inner = RecordingClassifier::new(vec![Some("capable"), Some("efficient")]);
201        let pin = TurnPin::new(inner.clone());
202        let mut state = State::default();
203        let opening = vec![user("debug this")];
204
205        assert_eq!(
206            selected(&pin, &mut state, opening.clone()).await?,
207            Some("capable".to_string())
208        );
209        let mut continued = opening;
210        continued.push(tool_call("call-1"));
211        continued.push(tool_result("call-1"));
212        assert_eq!(
213            selected(&pin, &mut state, continued).await?,
214            Some("capable".to_string())
215        );
216        assert_eq!(inner.consultations(), 1);
217        Ok(())
218    }
219
220    #[tokio::test]
221    async fn a_fresh_user_message_re_classifies() -> Result<()> {
222        let inner = RecordingClassifier::new(vec![Some("efficient"), Some("capable")]);
223        let pin = TurnPin::new(inner.clone());
224        let mut state = State::default();
225        let opening = vec![user("summarise this file")];
226
227        assert_eq!(
228            selected(&pin, &mut state, opening.clone()).await?,
229            Some("efficient".to_string())
230        );
231        let mut followed_up = opening;
232        followed_up.push(tool_call("call-1"));
233        followed_up.push(tool_result("call-1"));
234        followed_up.push(user("now find the race condition"));
235        assert_eq!(
236            selected(&pin, &mut state, followed_up).await?,
237            Some("capable".to_string())
238        );
239        assert_eq!(inner.consultations(), 2);
240        Ok(())
241    }
242
243    #[tokio::test]
244    async fn an_abstention_clears_an_earlier_pin() -> Result<()> {
245        let inner = RecordingClassifier::new(vec![Some("efficient"), None, Some("capable")]);
246        let pin = TurnPin::new(inner.clone());
247        let mut state = State::default();
248        let opening = vec![user("summarise this file")];
249
250        assert_eq!(
251            selected(&pin, &mut state, opening.clone()).await?,
252            Some("efficient".to_string())
253        );
254        let mut followed_up = opening;
255        followed_up.push(user("now find the race condition"));
256        assert_eq!(selected(&pin, &mut state, followed_up.clone()).await?, None);
257
258        // The abstaining turn is served by the fall-through default, so its tool
259        // continuations must not reuse the target pinned by the previous turn.
260        followed_up.push(tool_call("call-1"));
261        followed_up.push(tool_result("call-1"));
262        assert_eq!(
263            selected(&pin, &mut state, followed_up).await?,
264            Some("capable".to_string())
265        );
266        assert_eq!(inner.consultations(), 3);
267        Ok(())
268    }
269
270    #[test]
271    fn a_turn_starts_only_when_the_user_spoke_last() {
272        let mut messages = vec![user("debug this")];
273        assert!(starts_new_turn(&messages));
274        messages.push(tool_call("call-1"));
275        assert!(!starts_new_turn(&messages));
276        // Anthropic sends this as a user-role message, so it must not count as a turn.
277        messages.push(tool_result("call-1"));
278        assert!(!starts_new_turn(&messages));
279        messages.push(user("still broken"));
280        assert!(starts_new_turn(&messages));
281        assert!(!starts_new_turn(&[]));
282
283        // A tool result the user appended to counts, because the user did speak.
284        let mut mixed = tool_result("call-2");
285        mixed.content.extend(user("and also rename foo").content);
286        assert!(starts_new_turn(&[mixed]));
287    }
288}
switchyard_libsy/algorithms/util/turn_pin.rs

switchyard_libsy/algorithms/util/
turn_pin.rs
