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, Decision, InstructionBlock, LlmRequest, Message, ModelId, OutputParams, Request,
36 Response, Role, SamplingParams,
37};
38
39use crate::core::algorithm::{Algorithm, Driver};
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 executor_decision(&self) -> Decision {
239 Decision::new(self.executor.clone(), true)
240 }
241
242 fn check_exhausted(&self, scope: &ScopeKey) -> bool {
246 let mut state = self.state.lock();
247 let Some(entry) = state.scopes.get_mut(scope) else {
248 return false;
249 };
250 let exhausted = entry.reviews >= self.config.max_reviews
251 || entry.failed_consults >= MAX_FAILED_CONSULTS;
252 if exhausted && !entry.exhaustion_logged {
253 entry.exhaustion_logged = true;
254 tracing::info!(
255 target: "libsy",
256 scope = ?scope,
257 "advisor gate: review budget spent; passing through"
258 );
259 }
260 exhausted
261 }
262
263 fn try_reserve(&self, scope: &ScopeKey) -> bool {
267 let mut state = self.state.lock();
268 if state.scopes.len() >= MAX_TRACKED_SCOPES && !state.scopes.contains_key(scope) {
269 let evict = state
270 .scopes
271 .keys()
272 .find(|key| **key != ScopeKey::Instance)
273 .cloned();
274 if let Some(key) = evict {
275 state.scopes.remove(&key);
276 }
277 }
278 let max_reviews = self.config.max_reviews;
279 let entry = state.scopes.entry(scope.clone()).or_default();
280 if entry.reviews >= max_reviews || entry.failed_consults >= MAX_FAILED_CONSULTS {
281 return false;
282 }
283 entry.reviews += 1;
284 true
285 }
286
287 fn refund_failure(&self, scope: &ScopeKey) {
291 let mut state = self.state.lock();
292 let entry = state.scopes.entry(scope.clone()).or_default();
293 entry.reviews = entry.reviews.saturating_sub(1);
294 entry.failed_consults += 1;
295 }
296
297 fn evict_scope(&self, scope: &ScopeKey) {
299 if *scope == ScopeKey::Instance {
300 return;
301 }
302 self.state.lock().scopes.remove(scope);
303 }
304
305 fn stall_already_fired(&self, key: u64) -> bool {
306 self.state.lock().stall_fired.contains(&key)
307 }
308
309 fn mark_stall_fired(&self, key: u64) {
310 let mut state = self.state.lock();
311 if state.stall_fired.len() >= MAX_TRACKED_SCOPES {
312 let drop = state.stall_fired.iter().next().copied();
313 if let Some(key) = drop {
314 state.stall_fired.remove(&key);
315 }
316 }
317 state.stall_fired.insert(key);
318 }
319
320 async fn route_inner(
323 &self,
324 driver: &Driver,
325 request: Request,
326 scope: &ScopeKey,
327 ) -> Result<Response> {
328 if self.check_exhausted(scope) {
333 driver.decide(self.executor_decision()).await?;
334 return driver
335 .call_model(request, vec![self.executor.clone()], true)
336 .await;
337 }
338
339 driver.decide(self.executor_decision()).await?;
342 let response = driver
343 .call_model(request.clone(), vec![self.executor.clone()], true)
344 .await?;
345 let turn = buffer_turn(self.executor.as_str(), response).await?;
346
347 let stall_key = stall_key(&request);
351 let stall = self.config.gate_stall_turns > 0
352 && !self.stall_already_fired(stall_key)
353 && assistant_turns(&request.llm_request.messages) >= self.config.gate_stall_turns;
354 let triggered = match &self.trigger {
355 CompiledTrigger::Pattern(pattern) => {
356 pattern.is_match(visible_text(&turn.agg).as_deref().unwrap_or(""))
357 }
358 CompiledTrigger::NoToolCall => {
359 !has_tool_use(&turn.agg)
360 && count_tool_results(&request.llm_request.messages)
361 >= self.config.gate_min_tool_results
362 }
363 };
364 if !(triggered || stall) {
365 return Ok(turn.into_response());
366 }
367 if stall && !triggered {
370 self.mark_stall_fired(stall_key);
371 }
372 if !self.try_reserve(scope) {
373 return Ok(turn.into_response());
374 }
375
376 let trigger_label = match (&self.trigger, triggered) {
377 (CompiledTrigger::Pattern(_), true) => "pattern",
378 (CompiledTrigger::NoToolCall, true) => "no_tool_call",
379 _ => "stall",
380 };
381 let review_tail = visible_text(&turn.agg).or_else(|| {
382 reasoning_text(&turn.agg).map(|reasoning| format!("{REASONING_TAIL_LABEL}{reasoning}"))
383 });
384 match self
385 .consult(driver, &request, review_tail.as_deref(), trigger_label)
386 .await
387 {
388 Ok(ConsultOutcome::Approve) => Ok(turn.into_response()),
389 Ok(ConsultOutcome::Redo { plan }) => self.redo(driver, request, turn, &plan).await,
390 Ok(ConsultOutcome::Failed) => {
391 self.refund_failure(scope);
392 Ok(turn.into_response())
393 }
394 Err(error) => {
395 self.refund_failure(scope);
396 Err(error)
397 }
398 }
399 }
400
401 async fn redo(
405 &self,
406 driver: &Driver,
407 request: Request,
408 turn: GatedTurn,
409 plan: &str,
410 ) -> Result<Response> {
411 record_discarded(&turn.agg.usage);
412 emit_discarded_audit(self.executor.as_str(), &turn.agg.usage);
413 let echo = visible_text(&turn.agg)
414 .or_else(|| reasoning_text(&turn.agg))
415 .unwrap_or_else(|| EMPTY_ECHO_PLACEHOLDER.to_string());
416 let mut redo = request;
417 redo.llm_request
418 .messages
419 .push(Message::text(Role::Assistant, echo));
420 redo.llm_request.messages.push(Message::text(
421 Role::User,
422 format!("{}{}", self.config.redo_feedback_prefix, plan),
423 ));
424 crate::algorithms::util::prompts::drop_exact_replay(&mut redo);
428 driver.decide(self.executor_decision()).await?;
429 driver
430 .call_model(redo, vec![self.executor.clone()], true)
431 .await
432 }
433
434 async fn consult(
438 &self,
439 driver: &Driver,
440 base: &Request,
441 review_tail: Option<&str>,
442 trigger: &'static str,
443 ) -> Result<ConsultOutcome> {
444 let transcript_messages: Vec<Message> = base
449 .llm_request
450 .instructions
451 .iter()
452 .map(|block| Message {
453 role: block.role,
454 content: block.content.clone(),
455 })
456 .chain(base.llm_request.messages.iter().cloned())
457 .collect();
458 let transcript = review_transcript(
459 &transcript_messages,
460 review_tail,
461 self.config.transcript_max_chars,
462 );
463 let consult_request = self.build_consult_request(base, transcript);
464 let started = Instant::now();
465 let reply = match driver
468 .call_model(consult_request, vec![self.advisor.clone()], false)
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<Response> {
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}