1use std::sync::Arc;
2
3use bamboo_agent_core::tools::{FunctionSchema, ToolCall, ToolSchema};
4use bamboo_agent_core::{AgentError, AgentEvent, GoldCheckpoint, GoldConfidence, GoldDecision};
5use bamboo_agent_core::{Message, Role, Session};
6use bamboo_compression::{TiktokenTokenCounter, TokenCounter};
7use bamboo_domain::ReasoningEffort;
8use bamboo_llm::{LLMProvider, LLMRequestOptions};
9use chrono::Utc;
10use serde_json::json;
11use tokio::sync::mpsc;
12use tokio_util::sync::CancellationToken;
13
14use crate::runtime::config::GoldConfig;
15use crate::runtime::stream::handler::{
16 await_stream_bootstrap, consume_llm_stream_silent_with_context, StreamTimeoutContext,
17};
18use crate::runtime::task_context::TaskLoopContext;
19use bamboo_metrics::TokenUsage as MetricsTokenUsage;
20
21pub struct GoldEvalFrame<'a> {
25 pub event_tx: &'a mpsc::Sender<AgentEvent>,
26 pub session_id: &'a str,
27 pub model: &'a str,
28 pub timeout_context: StreamTimeoutContext,
29 pub reasoning_effort: Option<ReasoningEffort>,
30 pub checkpoint: GoldCheckpoint,
31 pub iteration: u32,
32}
33
34#[derive(Debug, Clone)]
35pub struct GoldEvaluationResult {
36 pub checkpoint: GoldCheckpoint,
37 pub iteration: u32,
38 pub decision: GoldDecision,
39 pub confidence: GoldConfidence,
40 pub reasoning: String,
41 pub missing_information: Vec<String>,
43 pub next_action: Option<String>,
46 pub prompt_tokens: u64,
47 pub completion_tokens: u64,
48}
49
50#[derive(Debug, Clone)]
51pub(crate) struct AsyncGoldEvaluationRequest {
52 pub(crate) session_id: String,
53 pub(crate) round_number: usize,
54 pub(crate) model_name: String,
55 pub(crate) timeout_context: StreamTimeoutContext,
56 pub(crate) reasoning_effort: Option<ReasoningEffort>,
57 pub(crate) checkpoint: GoldCheckpoint,
58 pub(crate) session_snapshot: Session,
59 pub(crate) task_context_snapshot: Option<TaskLoopContext>,
60 pub(crate) gold_config: GoldConfig,
61}
62
63#[derive(Debug, Clone)]
64pub(crate) struct AsyncGoldEvaluationResult {
65 pub(crate) round_number: usize,
66 pub(crate) model_name: String,
67 pub(crate) evaluation_result: GoldEvaluationResult,
68}
69
70fn normalize_lightweight_reasoning_effort(
71 reasoning_effort: Option<ReasoningEffort>,
72) -> Option<ReasoningEffort> {
73 reasoning_effort.map(|effort| match effort {
74 ReasoningEffort::Xhigh | ReasoningEffort::Max => ReasoningEffort::High,
75 other => other,
76 })
77}
78
79fn estimate_prompt_tokens(messages: &[Message]) -> u64 {
80 let counter = TiktokenTokenCounter::default();
81 u64::from(counter.count_messages(messages))
82}
83
84fn estimate_completion_tokens(content: &str, tool_calls: &[ToolCall]) -> u64 {
85 let counter = TiktokenTokenCounter::default();
86 let mut completion_surface = content.to_string();
87
88 for call in tool_calls {
89 if !completion_surface.is_empty() {
90 completion_surface.push('\n');
91 }
92 completion_surface.push_str(&call.function.name);
93 completion_surface.push('\n');
94 completion_surface.push_str(&call.function.arguments);
95 }
96
97 u64::from(counter.count_text(&completion_surface))
98}
99
100#[allow(clippy::too_many_arguments)]
101pub(crate) fn build_async_gold_evaluation_request(
102 task_context: &Option<TaskLoopContext>,
103 session: &Session,
104 session_id: &str,
105 round_number: usize,
106 model_name: Option<&str>,
107 reasoning_effort: Option<ReasoningEffort>,
108 checkpoint: GoldCheckpoint,
109 gold_config: &GoldConfig,
110 timeout_context: StreamTimeoutContext,
111) -> Result<Option<AsyncGoldEvaluationRequest>, AgentError> {
112 if !gold_config.enabled {
113 return Ok(None);
114 }
115
116 let model_name = gold_config
117 .model_name
118 .as_deref()
119 .or(model_name)
120 .ok_or_else(|| AgentError::LLM("gold evaluation model_name is required".to_string()))?;
121
122 Ok(Some(AsyncGoldEvaluationRequest {
123 session_id: session_id.to_string(),
124 round_number,
125 model_name: model_name.to_string(),
126 timeout_context,
127 reasoning_effort,
128 checkpoint,
129 session_snapshot: session.clone(),
130 task_context_snapshot: task_context.clone(),
131 gold_config: gold_config.clone(),
132 }))
133}
134
135pub(crate) async fn execute_async_gold_evaluation(
136 request: AsyncGoldEvaluationRequest,
137 llm: Arc<dyn LLMProvider>,
138 event_tx: mpsc::Sender<AgentEvent>,
139 configured_limit: usize,
140) -> AsyncGoldEvaluationResult {
141 let budget_provider = llm.clone();
142 let budget_model = request.model_name.clone();
143 let acquire_dispatch_guard = async move {
144 crate::runtime::runner::auxiliary_budget::acquire(
145 &budget_provider,
146 &budget_model,
147 configured_limit,
148 )
149 .await
150 };
151 let evaluation_result = match evaluate_gold_with_dispatch(
152 &request.session_snapshot,
153 request.task_context_snapshot.as_ref(),
154 &request.gold_config,
155 llm,
156 &GoldEvalFrame {
157 event_tx: &event_tx,
158 session_id: &request.session_id,
159 model: &request.model_name,
160 timeout_context: request.timeout_context.clone(),
161 reasoning_effort: request.reasoning_effort,
162 checkpoint: request.checkpoint,
163 iteration: request.round_number as u32,
164 },
165 acquire_dispatch_guard,
166 )
167 .await
168 {
169 Ok(result) => result,
170 Err(error) => GoldEvaluationResult {
171 checkpoint: request.checkpoint,
172 iteration: request.round_number as u32,
173 decision: GoldDecision::Continue,
174 confidence: GoldConfidence::Low,
175 reasoning: format!("Gold evaluation failed: {error}"),
176 missing_information: Vec::new(),
177 next_action: None,
178 prompt_tokens: 0,
179 completion_tokens: 0,
180 },
181 };
182
183 AsyncGoldEvaluationResult {
184 round_number: request.round_number,
185 model_name: request.model_name,
186 evaluation_result,
187 }
188}
189
190pub async fn evaluate_gold(
191 session: &Session,
192 task_context: Option<&TaskLoopContext>,
193 config: &GoldConfig,
194 llm: Arc<dyn LLMProvider>,
195 frame: &GoldEvalFrame<'_>,
196) -> Result<GoldEvaluationResult, AgentError> {
197 evaluate_gold_with_dispatch(
198 session,
199 task_context,
200 config,
201 llm,
202 frame,
203 std::future::ready(()),
204 )
205 .await
206}
207
208pub(crate) async fn evaluate_gold_with_dispatch<G, Fut>(
209 session: &Session,
210 task_context: Option<&TaskLoopContext>,
211 config: &GoldConfig,
212 llm: Arc<dyn LLMProvider>,
213 frame: &GoldEvalFrame<'_>,
214 acquire_dispatch_guard: Fut,
215) -> Result<GoldEvaluationResult, AgentError>
216where
217 Fut: std::future::Future<Output = G>,
218{
219 let event_tx = frame.event_tx;
221 let session_id = frame.session_id;
222 let model = frame.model;
223 let reasoning_effort = frame.reasoning_effort;
224 let checkpoint = frame.checkpoint;
225 let iteration = frame.iteration;
226
227 let _ = event_tx
228 .send(AgentEvent::GoldEvaluationStarted {
229 session_id: session_id.to_string(),
230 checkpoint,
231 iteration,
232 })
233 .await;
234
235 let messages = build_gold_messages(session, task_context, config, checkpoint);
236 let prompt_tokens = estimate_prompt_tokens(&messages);
237 let tools = get_gold_evaluation_tools();
238
239 let request_reasoning_effort = normalize_lightweight_reasoning_effort(reasoning_effort);
240 let request_options = LLMRequestOptions {
241 session_id: Some(session_id.to_string()),
242 reasoning_effort: request_reasoning_effort,
243 parallel_tool_calls: None,
244 required_tool: None,
245 responses: None,
246 request_purpose: Some("gold_evaluation".to_string()),
247 cache: None,
248 };
249
250 let cancel_token = CancellationToken::new();
251 let _dispatch_guard = acquire_dispatch_guard.await;
252 let timeout_context = frame.timeout_context.clone().begin_request();
253 let stream = await_stream_bootstrap(
254 llm.chat_stream_with_options(
255 &messages,
256 &tools,
257 Some(config.max_output_tokens),
258 model,
259 Some(&request_options),
260 ),
261 &cancel_token,
262 session_id,
263 &timeout_context,
264 )
265 .await?
266 .map_err(|error| AgentError::LLM(error.to_string()))?;
267 let stream_output =
268 consume_llm_stream_silent_with_context(stream, &cancel_token, session_id, &timeout_context)
269 .await?;
270
271 let result = parse_gold_evaluation(
272 &stream_output.content,
273 &stream_output.tool_calls,
274 checkpoint,
275 iteration,
276 prompt_tokens,
277 );
278
279 let _ = event_tx
280 .send(AgentEvent::GoldEvaluationCompleted {
281 session_id: session_id.to_string(),
282 checkpoint: result.checkpoint,
283 iteration: result.iteration,
284 decision: result.decision,
285 confidence: result.confidence,
286 reasoning: result.reasoning.clone(),
287 })
288 .await;
289
290 Ok(result)
291}
292
293pub fn build_gold_messages(
294 session: &Session,
295 task_context: Option<&TaskLoopContext>,
296 config: &GoldConfig,
297 checkpoint: GoldCheckpoint,
298) -> Vec<Message> {
299 let mut messages = Vec::new();
300
301 let mut system_prompt = String::from(
302 "You are a gold progress evaluator. Judge whether the agent has already achieved the user's goal, should continue execution, needs user input, is blocked, or is exhausted.\n\nRules:\n1. This phase is observe-only: do not mutate state or invent actions.\n2. You must call report_gold_evaluation exactly once.\n3. Use achieved only when the user's actual goal is satisfied.\n4. Use continue when more agent work is still appropriate.\n5. Use need_input only when missing user input is the true next blocker.\n6. Use blocked only for a concrete blocking condition.\n7. Use exhausted for loops, budget exhaustion, or clear inability to make progress.\n8. Keep reasoning short, concrete, and evidence-based."
303 );
304
305 if let Some(extra) = config
306 .evaluation_prompt
307 .as_deref()
308 .map(str::trim)
309 .filter(|value| !value.is_empty())
310 {
311 system_prompt.push_str("\n\nAdditional instructions:\n");
312 system_prompt.push_str(extra);
313 }
314
315 messages.push(Message::system(system_prompt));
316
317 let task_summary = task_context
318 .map(TaskLoopContext::format_for_prompt)
319 .filter(|value| !value.trim().is_empty())
320 .unwrap_or_else(|| "## Current Task List\nNo task list available.".to_string());
321
322 let pending_question_summary = session
323 .pending_question
324 .as_ref()
325 .map(|question| {
326 let options = if question.options.is_empty() {
327 "none".to_string()
328 } else {
329 question.options.join(" | ")
330 };
331 let tool_name = if question.tool_name.trim().is_empty() {
332 "unknown".to_string()
333 } else {
334 question.tool_name.clone()
335 };
336 format!(
337 "question={} | options={} | tool={} | source={:?}",
338 question.question, options, tool_name, question.source
339 )
340 })
341 .unwrap_or_else(|| "none".to_string());
342
343 let runtime_summary = session
344 .agent_runtime_state
345 .as_ref()
346 .map(|state| {
347 format!(
348 "status={:?} | current_round={} | max_rounds={} | suspend_reason={} | waiting_for_children={}",
349 state.status,
350 state.round.current_round,
351 state.round.max_rounds,
352 state
353 .suspension
354 .as_ref()
355 .map(|s| s.reason.clone())
356 .unwrap_or_else(|| "none".to_string()),
357 state.waiting_for_children.is_some()
358 )
359 })
360 .unwrap_or_else(|| "runtime_state=none".to_string());
361
362 let recent_messages = format_recent_messages(session, 6);
363
364 let goal_section = config
365 .effective_goal()
366 .map(|goal| format!("## Goal\n{goal}"))
367 .unwrap_or_else(|| {
368 "## Goal\nNo explicit goal set. Judge against the user's request inferred from the conversation.".to_string()
369 });
370
371 let user_prompt = format!(
372 "## Gold Checkpoint\ncheckpoint={}\n\n{}\n\n## Runtime\n{}\n\n## Pending Question\n{}\n\n{}\n\n## Recent Conversation\n{}\n\n## Instruction\nReport the best current Gold judgment for this checkpoint by measuring progress against the goal above. Remember: Phase 1 is observe-only, so only report decision/confidence/reasoning.",
373 checkpoint.as_str(),
374 goal_section,
375 runtime_summary,
376 pending_question_summary,
377 task_summary,
378 recent_messages,
379 );
380
381 messages.push(Message::user(user_prompt));
382 messages
383}
384
385fn format_recent_messages(session: &Session, limit: usize) -> String {
386 let start = session.messages.len().saturating_sub(limit);
387 let mut lines = Vec::new();
388
389 for message in session.messages.iter().skip(start) {
390 let role = match message.role {
391 Role::System => "system",
392 Role::User => "user",
393 Role::Assistant => "assistant",
394 Role::Tool => "tool",
395 };
396
397 let mut content = message.content.trim().replace('\n', " ");
398 if content.chars().count() > 240 {
399 content = format!("{}…", content.chars().take(240).collect::<String>());
400 }
401 if content.is_empty() {
402 content = "<empty>".to_string();
403 }
404
405 lines.push(format!("- [{}] {}", role, content));
406 }
407
408 if lines.is_empty() {
409 "- <no messages>".to_string()
410 } else {
411 lines.join("\n")
412 }
413}
414
415pub fn get_gold_evaluation_tools() -> Vec<ToolSchema> {
416 vec![ToolSchema {
417 schema_type: "function".to_string(),
418 function: FunctionSchema {
419 name: "report_gold_evaluation".to_string(),
420 description: "Report the current Gold evaluation decision for the session".to_string(),
421 parameters: json!({
422 "type": "object",
423 "properties": {
424 "decision": {
425 "type": "string",
426 "enum": ["continue", "achieved", "blocked", "need_input", "exhausted"]
427 },
428 "confidence": {
429 "type": "string",
430 "enum": ["low", "medium", "high"]
431 },
432 "reasoning": {
433 "type": "string",
434 "description": "Short concrete reasoning for the decision"
435 },
436 "missing_information": {
437 "type": "array",
438 "items": { "type": "string" },
439 "description": "Concrete pieces of information still missing to achieve the goal. Empty when nothing is missing."
440 },
441 "next_action": {
442 "type": "string",
443 "description": "The single most useful next action the agent should take. Provide when decision is continue."
444 }
445 },
446 "required": ["decision", "confidence", "reasoning"],
447 "additionalProperties": false
448 }),
449 },
450 }]
451}
452
453pub fn parse_gold_evaluation(
454 content: &str,
455 tool_calls: &[ToolCall],
456 checkpoint: GoldCheckpoint,
457 iteration: u32,
458 prompt_tokens: u64,
459) -> GoldEvaluationResult {
460 let completion_tokens = estimate_completion_tokens(content, tool_calls);
461 let parsed = parse_gold_result_from_tool_calls(tool_calls);
462
463 let parsed = parsed.unwrap_or_else(|| {
464 let fallback_reasoning = content.trim().to_string();
465 ParsedGoldResult {
466 decision: GoldDecision::Continue,
467 confidence: GoldConfidence::Low,
468 reasoning: if fallback_reasoning.is_empty() {
469 "Gold evaluation returned no structured result; defaulting to continue.".to_string()
470 } else {
471 fallback_reasoning
472 },
473 missing_information: Vec::new(),
474 next_action: None,
475 }
476 });
477
478 GoldEvaluationResult {
479 checkpoint,
480 iteration,
481 decision: parsed.decision,
482 confidence: parsed.confidence,
483 reasoning: parsed.reasoning,
484 missing_information: parsed.missing_information,
485 next_action: parsed.next_action,
486 prompt_tokens,
487 completion_tokens,
488 }
489}
490
491struct ParsedGoldResult {
492 decision: GoldDecision,
493 confidence: GoldConfidence,
494 reasoning: String,
495 missing_information: Vec<String>,
496 next_action: Option<String>,
497}
498
499fn parse_gold_result_from_tool_calls(tool_calls: &[ToolCall]) -> Option<ParsedGoldResult> {
500 for tool_call in tool_calls {
501 if tool_call.function.name != "report_gold_evaluation" {
502 continue;
503 }
504
505 let Ok(args) = serde_json::from_str::<serde_json::Value>(&tool_call.function.arguments)
506 else {
507 continue;
508 };
509
510 let decision = match args.get("decision").and_then(|value| value.as_str()) {
511 Some("continue") => GoldDecision::Continue,
512 Some("achieved") => GoldDecision::Achieved,
513 Some("blocked") => GoldDecision::Blocked,
514 Some("need_input") => GoldDecision::NeedInput,
515 Some("exhausted") => GoldDecision::Exhausted,
516 _ => continue,
517 };
518
519 let confidence = match args.get("confidence").and_then(|value| value.as_str()) {
520 Some("low") => GoldConfidence::Low,
521 Some("medium") => GoldConfidence::Medium,
522 Some("high") => GoldConfidence::High,
523 _ => GoldConfidence::Low,
524 };
525
526 let reasoning = args
527 .get("reasoning")
528 .and_then(|value| value.as_str())
529 .map(str::trim)
530 .filter(|value| !value.is_empty())
531 .unwrap_or("Gold evaluation produced no reasoning")
532 .to_string();
533
534 let missing_information = args
535 .get("missing_information")
536 .and_then(|value| value.as_array())
537 .map(|items| {
538 items
539 .iter()
540 .filter_map(|item| item.as_str())
541 .map(str::trim)
542 .filter(|value| !value.is_empty())
543 .map(str::to_string)
544 .collect::<Vec<_>>()
545 })
546 .unwrap_or_default();
547
548 let next_action = args
549 .get("next_action")
550 .and_then(|value| value.as_str())
551 .map(str::trim)
552 .filter(|value| !value.is_empty())
553 .map(str::to_string);
554
555 return Some(ParsedGoldResult {
556 decision,
557 confidence,
558 reasoning,
559 missing_information,
560 next_action,
561 });
562 }
563
564 None
565}
566
567pub(crate) fn apply_gold_evaluation_result(
568 session: &mut Session,
569 result: &GoldEvaluationResult,
570) -> MetricsTokenUsage {
571 let evaluation_count = session
572 .metadata
573 .get("gold.evaluation_count")
574 .and_then(|value| value.parse::<u64>().ok())
575 .unwrap_or(0)
576 .saturating_add(1);
577
578 let summary = json!({
579 "checkpoint": result.checkpoint.as_str(),
580 "iteration": result.iteration,
581 "decision": result.decision.as_str(),
582 "confidence": result.confidence.as_str(),
583 "reasoning": result.reasoning,
584 "recorded_at": Utc::now().to_rfc3339(),
585 });
586
587 session
588 .metadata
589 .insert("gold.last_evaluation".to_string(), summary.to_string());
590 session.metadata.insert(
591 "gold.last_decision".to_string(),
592 result.decision.as_str().to_string(),
593 );
594 session.metadata.insert(
595 "gold.last_confidence".to_string(),
596 result.confidence.as_str().to_string(),
597 );
598 session
599 .metadata
600 .insert("gold.last_reasoning".to_string(), result.reasoning.clone());
601 session.metadata.insert(
602 "gold.last_checkpoint".to_string(),
603 result.checkpoint.as_str().to_string(),
604 );
605 session.metadata.insert(
606 "gold.last_iteration".to_string(),
607 result.iteration.to_string(),
608 );
609 session.metadata.insert(
610 "gold.evaluation_count".to_string(),
611 evaluation_count.to_string(),
612 );
613 session.updated_at = Utc::now();
614
615 let mut usage = MetricsTokenUsage {
616 prompt_tokens: result.prompt_tokens,
617 completion_tokens: result.completion_tokens,
618 ..Default::default()
619 };
620 usage.recompute_total();
621 usage
622}
623
624#[cfg(test)]
625mod tests {
626 use super::*;
627 use bamboo_agent_core::tools::FunctionCall;
628
629 fn report_call(arguments: serde_json::Value) -> ToolCall {
630 ToolCall {
631 id: "call-1".to_string(),
632 tool_type: "function".to_string(),
633 function: FunctionCall {
634 name: "report_gold_evaluation".to_string(),
635 arguments: arguments.to_string(),
636 },
637 }
638 }
639
640 #[test]
641 fn parse_gold_result_from_tool_calls_reads_structured_fields() {
642 let parsed = parse_gold_result_from_tool_calls(&[report_call(json!({
643 "decision": "blocked",
644 "confidence": "high",
645 "reasoning": "Missing credentials",
646 "missing_information": ["API key", " ", "Database URL"],
647 "next_action": " Ask the user for the API key "
648 }))])
649 .expect("gold result should parse");
650
651 assert_eq!(parsed.decision, GoldDecision::Blocked);
652 assert_eq!(parsed.confidence, GoldConfidence::High);
653 assert_eq!(parsed.reasoning, "Missing credentials");
654 assert_eq!(
655 parsed.missing_information,
656 vec!["API key".to_string(), "Database URL".to_string()]
657 );
658 assert_eq!(
659 parsed.next_action.as_deref(),
660 Some("Ask the user for the API key")
661 );
662 }
663
664 #[test]
665 fn parse_gold_result_from_tool_calls_ignores_other_tools() {
666 let parsed = parse_gold_result_from_tool_calls(&[ToolCall {
667 id: "call-1".to_string(),
668 tool_type: "function".to_string(),
669 function: FunctionCall {
670 name: "other_tool".to_string(),
671 arguments: "{}".to_string(),
672 },
673 }]);
674
675 assert!(parsed.is_none());
676 }
677
678 #[test]
679 fn apply_gold_evaluation_result_updates_metadata_keys() {
680 let mut session = Session::new("session-1", "model");
681 let result = GoldEvaluationResult {
682 checkpoint: GoldCheckpoint::PostRound,
683 iteration: 2,
684 decision: GoldDecision::Achieved,
685 confidence: GoldConfidence::Medium,
686 reasoning: "Goal satisfied".to_string(),
687 missing_information: Vec::new(),
688 next_action: None,
689 prompt_tokens: 10,
690 completion_tokens: 5,
691 };
692
693 let usage = apply_gold_evaluation_result(&mut session, &result);
694
695 assert_eq!(
696 session
697 .metadata
698 .get("gold.last_decision")
699 .map(String::as_str),
700 Some("achieved")
701 );
702 assert_eq!(
703 session
704 .metadata
705 .get("gold.last_confidence")
706 .map(String::as_str),
707 Some("medium")
708 );
709 assert_eq!(
710 session
711 .metadata
712 .get("gold.last_checkpoint")
713 .map(String::as_str),
714 Some("post_round")
715 );
716 assert_eq!(
717 session
718 .metadata
719 .get("gold.evaluation_count")
720 .map(String::as_str),
721 Some("1")
722 );
723 assert_eq!(usage.prompt_tokens, 10);
724 assert_eq!(usage.completion_tokens, 5);
725 assert_eq!(usage.total_tokens, 15);
726 }
727}