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