1use std::collections::{HashMap, HashSet};
29use std::hash::{DefaultHasher, Hash, Hasher};
30use std::sync::Arc;
31use std::time::Instant;
32
33use parking_lot::Mutex;
34use switchyard_protocol::{
35 ContentBlock, InstructionBlock, LlmRequest, Message, ModelId, OutputParams, Request, Role,
36 SamplingParams,
37};
38
39use crate::algorithms::util::tool_signals::ToolSignals;
40use crate::core::algorithm::{Algorithm, Driver, RoutingOutcome};
41use crate::{LibsyError, Result};
42
43mod telemetry;
44#[cfg(test)]
45mod tests;
46mod transcript;
47mod turn;
48
49use telemetry::{
50 ReviewAudit, emit_discarded_audit, emit_review_audit, record_consult_failure, record_discarded,
51 record_review,
52};
53use transcript::{VERDICT_PATTERN, Verdict, advisor_reply_text, parse_verdict, review_transcript};
54use turn::{GatedTurn, buffer_turn, has_tool_use, reasoning_text, visible_text};
55
56pub const REVIEWER_SYSTEM_PROMPT: &str =
58 include_str!("../prompts/advisor-gate/reviewer-system-prompt.md");
59
60pub const REDO_FEEDBACK_PREFIX: &str = concat!(
63 include_str!("../prompts/advisor-gate/redo-feedback-prefix.md"),
64 "\n"
65);
66
67const REASONING_TAIL_LABEL: &str =
71 "(the executor produced no visible text this turn; its internal reasoning follows)\n";
72const EMPTY_ECHO_PLACEHOLDER: &str = "(the executor produced no output this turn)";
75const MAX_FAILED_CONSULTS: u32 = 3;
80const MAX_TRACKED_SCOPES: usize = 1_024;
83const BENCH_SESSION_HEADER: &str = "proxy_x_session_id";
87
88#[derive(Clone, Debug, PartialEq)]
90pub enum GateTrigger {
91 NoToolCall,
93 Pattern(String),
97}
98
99#[derive(Clone, Debug)]
101pub struct AdvisorGateConfig {
102 pub reviewer_system_prompt: String,
104 pub redo_feedback_prefix: String,
106 pub gate_trigger: GateTrigger,
108 pub max_reviews: u32,
112 pub gate_stall_turns: u32,
117 pub gate_min_tool_results: u32,
121 pub advisor_max_tokens: u64,
123 pub advisor_temperature: Option<f64>,
125 pub transcript_max_chars: usize,
128 pub fail_open: bool,
131}
132
133impl Default for AdvisorGateConfig {
134 fn default() -> Self {
135 Self {
136 reviewer_system_prompt: REVIEWER_SYSTEM_PROMPT.to_string(),
137 redo_feedback_prefix: REDO_FEEDBACK_PREFIX.to_string(),
138 gate_trigger: GateTrigger::NoToolCall,
139 max_reviews: 1,
140 gate_stall_turns: 0,
141 gate_min_tool_results: 0,
142 advisor_max_tokens: 2048,
143 advisor_temperature: None,
144 transcript_max_chars: 200_000,
145 fail_open: true,
146 }
147 }
148}
149
150enum CompiledTrigger {
152 NoToolCall,
153 Pattern(regex::Regex),
154}
155
156#[derive(Clone, Debug, Hash, PartialEq, Eq)]
160enum ScopeKey {
161 Instance,
162 Client(String),
163 Session(String),
164}
165
166#[derive(Default)]
168struct ScopeState {
169 reviews: u32,
170 failed_consults: u32,
171 exhaustion_logged: bool,
172}
173
174#[derive(Default)]
177struct GateState {
178 scopes: HashMap<ScopeKey, ScopeState>,
179 stall_fired: HashSet<u64>,
180}
181
182pub struct AdvisorGate {
186 executor: ModelId,
187 advisor: ModelId,
188 config: AdvisorGateConfig,
189 trigger: CompiledTrigger,
190 verdict_re: regex::Regex,
191 state: Mutex<GateState>,
192}
193
194impl AdvisorGate {
195 pub fn new(executor: ModelId, advisor: ModelId, config: AdvisorGateConfig) -> Result<Self> {
197 if config.max_reviews < 1 {
198 return Err(algorithm_error("max_reviews must be at least 1"));
199 }
200 if config.advisor_max_tokens < 1 {
201 return Err(algorithm_error("advisor_max_tokens must be at least 1"));
202 }
203 if config.transcript_max_chars < 256 {
204 return Err(algorithm_error("transcript_max_chars must be at least 256"));
205 }
206 let trigger = match &config.gate_trigger {
207 GateTrigger::NoToolCall => CompiledTrigger::NoToolCall,
208 GateTrigger::Pattern(pattern) => {
209 if pattern.is_empty() {
210 return Err(algorithm_error(
211 "gate_trigger 'pattern' requires a non-empty gate_trigger_pattern",
212 ));
213 }
214 CompiledTrigger::Pattern(regex::Regex::new(pattern).map_err(|error| {
215 algorithm_error(format!(
216 "gate_trigger_pattern is not a valid regex: {error}"
217 ))
218 })?)
219 }
220 };
221 let verdict_re = regex::Regex::new(VERDICT_PATTERN).map_err(|error| {
222 algorithm_error(format!("verdict pattern failed to compile: {error}"))
223 })?;
224 Ok(Self {
225 executor,
226 advisor,
227 config,
228 trigger,
229 verdict_re,
230 state: Mutex::new(GateState::default()),
231 })
232 }
233
234 fn check_exhausted(&self, scope: &ScopeKey) -> bool {
238 let mut state = self.state.lock();
239 let Some(entry) = state.scopes.get_mut(scope) else {
240 return false;
241 };
242 let exhausted = entry.reviews >= self.config.max_reviews
243 || entry.failed_consults >= MAX_FAILED_CONSULTS;
244 if exhausted && !entry.exhaustion_logged {
245 entry.exhaustion_logged = true;
246 tracing::info!(
247 target: "libsy",
248 scope = ?scope,
249 "advisor gate: review budget spent; passing through"
250 );
251 }
252 exhausted
253 }
254
255 fn try_reserve(&self, scope: &ScopeKey) -> bool {
259 let mut state = self.state.lock();
260 if state.scopes.len() >= MAX_TRACKED_SCOPES && !state.scopes.contains_key(scope) {
261 let evict = state
262 .scopes
263 .keys()
264 .find(|key| **key != ScopeKey::Instance)
265 .cloned();
266 if let Some(key) = evict {
267 state.scopes.remove(&key);
268 }
269 }
270 let max_reviews = self.config.max_reviews;
271 let entry = state.scopes.entry(scope.clone()).or_default();
272 if entry.reviews >= max_reviews || entry.failed_consults >= MAX_FAILED_CONSULTS {
273 return false;
274 }
275 entry.reviews += 1;
276 true
277 }
278
279 fn refund_failure(&self, scope: &ScopeKey) {
283 let mut state = self.state.lock();
284 let entry = state.scopes.entry(scope.clone()).or_default();
285 entry.reviews = entry.reviews.saturating_sub(1);
286 entry.failed_consults += 1;
287 }
288
289 fn evict_scope(&self, scope: &ScopeKey) {
291 if *scope == ScopeKey::Instance {
292 return;
293 }
294 self.state.lock().scopes.remove(scope);
295 }
296
297 fn stall_already_fired(&self, key: u64) -> bool {
298 self.state.lock().stall_fired.contains(&key)
299 }
300
301 fn mark_stall_fired(&self, key: u64) {
302 let mut state = self.state.lock();
303 if state.stall_fired.len() >= MAX_TRACKED_SCOPES {
304 let drop = state.stall_fired.iter().next().copied();
305 if let Some(key) = drop {
306 state.stall_fired.remove(&key);
307 }
308 }
309 state.stall_fired.insert(key);
310 }
311
312 async fn route_inner(
315 &self,
316 driver: &Driver,
317 request: Request,
318 scope: &ScopeKey,
319 ) -> Result<RoutingOutcome> {
320 if self.check_exhausted(scope) {
325 return Ok(RoutingOutcome::route_to(
326 self.executor.clone(),
327 Vec::new(),
328 request,
329 ));
330 }
331
332 let response = driver
335 .call_model(request.clone(), vec![self.executor.clone()])
336 .await?;
337 let turn = buffer_turn(self.executor.as_str(), response).await?;
338
339 let signals = ToolSignals::from_request(&request, None);
343 let stall_key = stall_key(&request);
347 let stall = self.config.gate_stall_turns > 0
348 && !self.stall_already_fired(stall_key)
349 && signals.assistant_turn_count >= self.config.gate_stall_turns;
350 let triggered = match &self.trigger {
351 CompiledTrigger::Pattern(pattern) => {
352 pattern.is_match(visible_text(&turn.agg).as_deref().unwrap_or(""))
353 }
354 CompiledTrigger::NoToolCall => {
355 !has_tool_use(&turn.agg)
356 && signals.tool_result_count >= self.config.gate_min_tool_results
357 }
358 };
359 if !(triggered || stall) {
360 return Ok(RoutingOutcome::answered(
361 self.executor.clone(),
362 request,
363 turn.into_response(),
364 ));
365 }
366 if stall && !triggered {
369 self.mark_stall_fired(stall_key);
370 }
371 if !self.try_reserve(scope) {
372 return Ok(RoutingOutcome::answered(
373 self.executor.clone(),
374 request,
375 turn.into_response(),
376 ));
377 }
378
379 let trigger_label = match (&self.trigger, triggered) {
380 (CompiledTrigger::Pattern(_), true) => "pattern",
381 (CompiledTrigger::NoToolCall, true) => "no_tool_call",
382 _ => "stall",
383 };
384 let review_tail = visible_text(&turn.agg).or_else(|| {
385 reasoning_text(&turn.agg).map(|reasoning| format!("{REASONING_TAIL_LABEL}{reasoning}"))
386 });
387 match self
388 .consult(driver, &request, review_tail.as_deref(), trigger_label)
389 .await
390 {
391 Ok(ConsultOutcome::Approve) => Ok(RoutingOutcome::answered(
392 self.executor.clone(),
393 request,
394 turn.into_response(),
395 )),
396 Ok(ConsultOutcome::Redo { plan }) => Ok(self.redo(request, turn, &plan)),
397 Ok(ConsultOutcome::Failed) => {
398 self.refund_failure(scope);
399 Ok(RoutingOutcome::answered(
400 self.executor.clone(),
401 request,
402 turn.into_response(),
403 ))
404 }
405 Err(error) => {
406 self.refund_failure(scope);
407 Err(error)
408 }
409 }
410 }
411
412 fn redo(&self, request: Request, turn: GatedTurn, plan: &str) -> RoutingOutcome {
416 record_discarded(&turn.agg.usage);
417 emit_discarded_audit(self.executor.as_str(), &turn.agg.usage);
418 let echo = visible_text(&turn.agg)
419 .or_else(|| reasoning_text(&turn.agg))
420 .unwrap_or_else(|| EMPTY_ECHO_PLACEHOLDER.to_string());
421 let mut redo = request;
422 redo.llm_request
423 .messages
424 .push(Message::text(Role::Assistant, echo));
425 redo.llm_request.messages.push(Message::text(
426 Role::User,
427 format!("{}{}", self.config.redo_feedback_prefix, plan),
428 ));
429 crate::algorithms::util::prompts::drop_exact_replay(&mut redo);
433 RoutingOutcome::route_to(self.executor.clone(), Vec::new(), redo)
434 }
435
436 async fn consult(
440 &self,
441 driver: &Driver,
442 base: &Request,
443 review_tail: Option<&str>,
444 trigger: &'static str,
445 ) -> Result<ConsultOutcome> {
446 let transcript_messages: Vec<Message> = base
451 .llm_request
452 .instructions
453 .iter()
454 .map(|block| Message {
455 role: block.role,
456 content: block.content.clone(),
457 })
458 .chain(base.llm_request.messages.iter().cloned())
459 .collect();
460 let transcript = review_transcript(
461 &transcript_messages,
462 review_tail,
463 self.config.transcript_max_chars,
464 );
465 let consult_request = self.build_consult_request(base, transcript);
466 let started = Instant::now();
467 let reply = match driver
468 .call_model(consult_request, vec![self.advisor.clone()])
469 .await
470 {
471 Ok(response) => response
472 .llm_response
473 .into_agg()
474 .await
475 .map_err(|source| LibsyError::client_call(self.advisor.clone(), source)),
476 Err(error) => Err(error),
477 };
478 let latency_ms = started.elapsed().as_secs_f64() * 1000.0;
479 let agg = match reply {
480 Ok(agg) => agg,
481 Err(error) => {
482 record_consult_failure(crate::algorithms::util::llm_judge::libsy_error_reason(
483 &error,
484 ));
485 if !self.config.fail_open {
486 return Err(algorithm_error(format!(
492 "advisor consult failed (fail_open = false): {error}"
493 )));
494 }
495 tracing::warn!(
496 target: "libsy",
497 error = %error,
498 "advisor gate: consult failed; passing the turn through (fail open)"
499 );
500 emit_review_audit(ReviewAudit {
501 verdict: "APPROVE",
502 error: Some(error.to_string()),
503 latency_ms,
504 reply_head: None,
505 usage: None,
506 });
507 return Ok(ConsultOutcome::Failed);
508 }
509 };
510 let reply_text = advisor_reply_text(&agg);
511 let reply_head: String = reply_text.chars().take(160).collect();
512 match parse_verdict(&self.verdict_re, &reply_text) {
513 Some(Verdict::Approve) => {
514 record_review("approve", trigger);
515 emit_review_audit(ReviewAudit {
516 verdict: "APPROVE",
517 error: None,
518 latency_ms,
519 reply_head: Some(reply_head),
520 usage: Some(&agg.usage),
521 });
522 Ok(ConsultOutcome::Approve)
523 }
524 Some(Verdict::Redo { plan }) => {
525 record_review("redo", trigger);
526 emit_review_audit(ReviewAudit {
527 verdict: "REDO",
528 error: None,
529 latency_ms,
530 reply_head: Some(reply_head),
531 usage: Some(&agg.usage),
532 });
533 Ok(ConsultOutcome::Redo { plan })
534 }
535 None => {
536 record_review("unparseable", trigger);
540 emit_review_audit(ReviewAudit {
541 verdict: "UNPARSEABLE",
542 error: None,
543 latency_ms,
544 reply_head: Some(reply_head),
545 usage: Some(&agg.usage),
546 });
547 Ok(ConsultOutcome::Failed)
548 }
549 }
550 }
551
552 fn build_consult_request(&self, base: &Request, transcript: String) -> Request {
555 Request {
556 llm_request: LlmRequest {
557 model: base.llm_request.model.clone(),
558 instructions: vec![InstructionBlock {
559 role: Role::System,
560 content: vec![ContentBlock::Text {
561 text: self.config.reviewer_system_prompt.clone(),
562 }],
563 }],
564 messages: vec![Message::text(Role::User, transcript)],
565 sampling: SamplingParams {
566 temperature: self.config.advisor_temperature,
567 ..SamplingParams::default()
568 },
569 output: OutputParams {
570 max_output_tokens: Some(self.config.advisor_max_tokens),
571 response_format: None,
572 },
573 ..LlmRequest::default()
574 },
575 raw_request: None,
576 metadata: base.metadata.clone(),
577 }
578 }
579}
580
581#[async_trait::async_trait]
582impl Algorithm for AdvisorGate {
583 fn name(&self) -> &str {
584 "advisor_gate"
585 }
586
587 async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<RoutingOutcome> {
588 let scope = budget_scope(&request);
589 let session_final = request
590 .metadata
591 .as_ref()
592 .and_then(|metadata| metadata.session_final)
593 == Some(true);
594 let result = self.route_inner(&driver, request, &scope).await;
595 if session_final {
596 self.evict_scope(&scope);
597 }
598 result
599 }
600}
601
602enum ConsultOutcome {
604 Approve,
605 Redo { plan: String },
606 Failed,
607}
608
609fn budget_scope(request: &Request) -> ScopeKey {
615 let metadata = request.metadata.as_ref();
616 if let Some(value) = metadata
617 .and_then(|metadata| metadata.http_headers.as_ref())
618 .and_then(|headers| headers.get(BENCH_SESSION_HEADER))
619 .and_then(|value| value.to_str().ok())
620 && !value.is_empty()
621 {
622 return ScopeKey::Client(value.to_string());
623 }
624 if let Some(id) = metadata.and_then(|metadata| metadata.session_id.as_deref())
625 && !id.is_empty()
626 {
627 return ScopeKey::Session(id.to_string());
628 }
629 ScopeKey::Instance
630}
631
632fn stall_key(request: &Request) -> u64 {
635 let text = request
636 .llm_request
637 .messages
638 .iter()
639 .find(|message| message.role == Role::User)
640 .and_then(|message| message.text_content("\n"))
641 .unwrap_or_default();
642 let mut hasher = DefaultHasher::new();
643 text.hash(&mut hasher);
644 hasher.finish()
645}
646
647fn algorithm_error(message: impl Into<String>) -> LibsyError {
648 LibsyError::AlgorithmError {
649 message: message.into(),
650 }
651}