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 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) -> AsyncGoldEvaluationResult {
140 let evaluation_result = match evaluate_gold(
141 &request.session_snapshot,
142 request.task_context_snapshot.as_ref(),
143 &request.gold_config,
144 llm,
145 &GoldEvalFrame {
146 event_tx: &event_tx,
147 session_id: &request.session_id,
148 model: &request.model_name,
149 timeout_context: request.timeout_context.clone(),
150 reasoning_effort: request.reasoning_effort,
151 checkpoint: request.checkpoint,
152 iteration: request.round_number as u32,
153 },
154 )
155 .await
156 {
157 Ok(result) => result,
158 Err(error) => GoldEvaluationResult {
159 checkpoint: request.checkpoint,
160 iteration: request.round_number as u32,
161 decision: GoldDecision::Continue,
162 confidence: GoldConfidence::Low,
163 reasoning: format!("Gold evaluation failed: {error}"),
164 missing_information: Vec::new(),
165 next_action: None,
166 prompt_tokens: 0,
167 completion_tokens: 0,
168 },
169 };
170
171 AsyncGoldEvaluationResult {
172 round_number: request.round_number,
173 model_name: request.model_name,
174 evaluation_result,
175 }
176}
177
178pub async fn evaluate_gold(
179 session: &Session,
180 task_context: Option<&TaskLoopContext>,
181 config: &GoldConfig,
182 llm: Arc<dyn LLMProvider>,
183 frame: &GoldEvalFrame<'_>,
184) -> Result<GoldEvaluationResult, AgentError> {
185 let event_tx = frame.event_tx;
187 let session_id = frame.session_id;
188 let model = frame.model;
189 let reasoning_effort = frame.reasoning_effort;
190 let checkpoint = frame.checkpoint;
191 let iteration = frame.iteration;
192
193 let _ = event_tx
194 .send(AgentEvent::GoldEvaluationStarted {
195 session_id: session_id.to_string(),
196 checkpoint,
197 iteration,
198 })
199 .await;
200
201 let messages = build_gold_messages(session, task_context, config, checkpoint);
202 let prompt_tokens = estimate_prompt_tokens(&messages);
203 let tools = get_gold_evaluation_tools();
204
205 let request_reasoning_effort = normalize_lightweight_reasoning_effort(reasoning_effort);
206 let request_options = LLMRequestOptions {
207 session_id: Some(session_id.to_string()),
208 reasoning_effort: request_reasoning_effort,
209 parallel_tool_calls: None,
210 required_tool: None,
211 responses: None,
212 request_purpose: Some("gold_evaluation".to_string()),
213 cache: None,
214 };
215
216 match llm
217 .chat_stream_with_options(
218 &messages,
219 &tools,
220 Some(config.max_output_tokens),
221 model,
222 Some(&request_options),
223 )
224 .await
225 {
226 Ok(stream) => {
227 let stream_output = consume_llm_stream_silent_with_context(
228 stream,
229 &CancellationToken::new(),
230 session_id,
231 &frame.timeout_context,
232 )
233 .await?;
234
235 let result = parse_gold_evaluation(
236 &stream_output.content,
237 &stream_output.tool_calls,
238 checkpoint,
239 iteration,
240 prompt_tokens,
241 );
242
243 let _ = event_tx
244 .send(AgentEvent::GoldEvaluationCompleted {
245 session_id: session_id.to_string(),
246 checkpoint: result.checkpoint,
247 iteration: result.iteration,
248 decision: result.decision,
249 confidence: result.confidence,
250 reasoning: result.reasoning.clone(),
251 })
252 .await;
253
254 Ok(result)
255 }
256 Err(error) => Err(AgentError::LLM(error.to_string())),
257 }
258}
259
260pub fn build_gold_messages(
261 session: &Session,
262 task_context: Option<&TaskLoopContext>,
263 config: &GoldConfig,
264 checkpoint: GoldCheckpoint,
265) -> Vec<Message> {
266 let mut messages = Vec::new();
267
268 let mut system_prompt = String::from(
269 "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."
270 );
271
272 if let Some(extra) = config
273 .evaluation_prompt
274 .as_deref()
275 .map(str::trim)
276 .filter(|value| !value.is_empty())
277 {
278 system_prompt.push_str("\n\nAdditional instructions:\n");
279 system_prompt.push_str(extra);
280 }
281
282 messages.push(Message::system(system_prompt));
283
284 let task_summary = task_context
285 .map(TaskLoopContext::format_for_prompt)
286 .filter(|value| !value.trim().is_empty())
287 .unwrap_or_else(|| "## Current Task List\nNo task list available.".to_string());
288
289 let pending_question_summary = session
290 .pending_question
291 .as_ref()
292 .map(|question| {
293 let options = if question.options.is_empty() {
294 "none".to_string()
295 } else {
296 question.options.join(" | ")
297 };
298 let tool_name = if question.tool_name.trim().is_empty() {
299 "unknown".to_string()
300 } else {
301 question.tool_name.clone()
302 };
303 format!(
304 "question={} | options={} | tool={} | source={:?}",
305 question.question, options, tool_name, question.source
306 )
307 })
308 .unwrap_or_else(|| "none".to_string());
309
310 let runtime_summary = session
311 .agent_runtime_state
312 .as_ref()
313 .map(|state| {
314 format!(
315 "status={:?} | current_round={} | max_rounds={} | suspend_reason={} | waiting_for_children={}",
316 state.status,
317 state.round.current_round,
318 state.round.max_rounds,
319 state
320 .suspension
321 .as_ref()
322 .map(|s| s.reason.clone())
323 .unwrap_or_else(|| "none".to_string()),
324 state.waiting_for_children.is_some()
325 )
326 })
327 .unwrap_or_else(|| "runtime_state=none".to_string());
328
329 let recent_messages = format_recent_messages(session, 6);
330
331 let goal_section = config
332 .effective_goal()
333 .map(|goal| format!("## Goal\n{goal}"))
334 .unwrap_or_else(|| {
335 "## Goal\nNo explicit goal set. Judge against the user's request inferred from the conversation.".to_string()
336 });
337
338 let user_prompt = format!(
339 "## 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.",
340 checkpoint.as_str(),
341 goal_section,
342 runtime_summary,
343 pending_question_summary,
344 task_summary,
345 recent_messages,
346 );
347
348 messages.push(Message::user(user_prompt));
349 messages
350}
351
352fn format_recent_messages(session: &Session, limit: usize) -> String {
353 let start = session.messages.len().saturating_sub(limit);
354 let mut lines = Vec::new();
355
356 for message in session.messages.iter().skip(start) {
357 let role = match message.role {
358 Role::System => "system",
359 Role::User => "user",
360 Role::Assistant => "assistant",
361 Role::Tool => "tool",
362 };
363
364 let mut content = message.content.trim().replace('\n', " ");
365 if content.chars().count() > 240 {
366 content = format!("{}…", content.chars().take(240).collect::<String>());
367 }
368 if content.is_empty() {
369 content = "<empty>".to_string();
370 }
371
372 lines.push(format!("- [{}] {}", role, content));
373 }
374
375 if lines.is_empty() {
376 "- <no messages>".to_string()
377 } else {
378 lines.join("\n")
379 }
380}
381
382pub fn get_gold_evaluation_tools() -> Vec<ToolSchema> {
383 vec![ToolSchema {
384 schema_type: "function".to_string(),
385 function: FunctionSchema {
386 name: "report_gold_evaluation".to_string(),
387 description: "Report the current Gold evaluation decision for the session".to_string(),
388 parameters: json!({
389 "type": "object",
390 "properties": {
391 "decision": {
392 "type": "string",
393 "enum": ["continue", "achieved", "blocked", "need_input", "exhausted"]
394 },
395 "confidence": {
396 "type": "string",
397 "enum": ["low", "medium", "high"]
398 },
399 "reasoning": {
400 "type": "string",
401 "description": "Short concrete reasoning for the decision"
402 },
403 "missing_information": {
404 "type": "array",
405 "items": { "type": "string" },
406 "description": "Concrete pieces of information still missing to achieve the goal. Empty when nothing is missing."
407 },
408 "next_action": {
409 "type": "string",
410 "description": "The single most useful next action the agent should take. Provide when decision is continue."
411 }
412 },
413 "required": ["decision", "confidence", "reasoning"],
414 "additionalProperties": false
415 }),
416 },
417 }]
418}
419
420pub fn parse_gold_evaluation(
421 content: &str,
422 tool_calls: &[ToolCall],
423 checkpoint: GoldCheckpoint,
424 iteration: u32,
425 prompt_tokens: u64,
426) -> GoldEvaluationResult {
427 let completion_tokens = estimate_completion_tokens(content, tool_calls);
428 let parsed = parse_gold_result_from_tool_calls(tool_calls);
429
430 let parsed = parsed.unwrap_or_else(|| {
431 let fallback_reasoning = content.trim().to_string();
432 ParsedGoldResult {
433 decision: GoldDecision::Continue,
434 confidence: GoldConfidence::Low,
435 reasoning: if fallback_reasoning.is_empty() {
436 "Gold evaluation returned no structured result; defaulting to continue.".to_string()
437 } else {
438 fallback_reasoning
439 },
440 missing_information: Vec::new(),
441 next_action: None,
442 }
443 });
444
445 GoldEvaluationResult {
446 checkpoint,
447 iteration,
448 decision: parsed.decision,
449 confidence: parsed.confidence,
450 reasoning: parsed.reasoning,
451 missing_information: parsed.missing_information,
452 next_action: parsed.next_action,
453 prompt_tokens,
454 completion_tokens,
455 }
456}
457
458struct ParsedGoldResult {
459 decision: GoldDecision,
460 confidence: GoldConfidence,
461 reasoning: String,
462 missing_information: Vec<String>,
463 next_action: Option<String>,
464}
465
466fn parse_gold_result_from_tool_calls(tool_calls: &[ToolCall]) -> Option<ParsedGoldResult> {
467 for tool_call in tool_calls {
468 if tool_call.function.name != "report_gold_evaluation" {
469 continue;
470 }
471
472 let Ok(args) = serde_json::from_str::<serde_json::Value>(&tool_call.function.arguments)
473 else {
474 continue;
475 };
476
477 let decision = match args.get("decision").and_then(|value| value.as_str()) {
478 Some("continue") => GoldDecision::Continue,
479 Some("achieved") => GoldDecision::Achieved,
480 Some("blocked") => GoldDecision::Blocked,
481 Some("need_input") => GoldDecision::NeedInput,
482 Some("exhausted") => GoldDecision::Exhausted,
483 _ => continue,
484 };
485
486 let confidence = match args.get("confidence").and_then(|value| value.as_str()) {
487 Some("low") => GoldConfidence::Low,
488 Some("medium") => GoldConfidence::Medium,
489 Some("high") => GoldConfidence::High,
490 _ => GoldConfidence::Low,
491 };
492
493 let reasoning = args
494 .get("reasoning")
495 .and_then(|value| value.as_str())
496 .map(str::trim)
497 .filter(|value| !value.is_empty())
498 .unwrap_or("Gold evaluation produced no reasoning")
499 .to_string();
500
501 let missing_information = args
502 .get("missing_information")
503 .and_then(|value| value.as_array())
504 .map(|items| {
505 items
506 .iter()
507 .filter_map(|item| item.as_str())
508 .map(str::trim)
509 .filter(|value| !value.is_empty())
510 .map(str::to_string)
511 .collect::<Vec<_>>()
512 })
513 .unwrap_or_default();
514
515 let next_action = args
516 .get("next_action")
517 .and_then(|value| value.as_str())
518 .map(str::trim)
519 .filter(|value| !value.is_empty())
520 .map(str::to_string);
521
522 return Some(ParsedGoldResult {
523 decision,
524 confidence,
525 reasoning,
526 missing_information,
527 next_action,
528 });
529 }
530
531 None
532}
533
534pub(crate) fn apply_gold_evaluation_result(
535 session: &mut Session,
536 result: &GoldEvaluationResult,
537) -> MetricsTokenUsage {
538 let evaluation_count = session
539 .metadata
540 .get("gold.evaluation_count")
541 .and_then(|value| value.parse::<u64>().ok())
542 .unwrap_or(0)
543 .saturating_add(1);
544
545 let summary = json!({
546 "checkpoint": result.checkpoint.as_str(),
547 "iteration": result.iteration,
548 "decision": result.decision.as_str(),
549 "confidence": result.confidence.as_str(),
550 "reasoning": result.reasoning,
551 "recorded_at": Utc::now().to_rfc3339(),
552 });
553
554 session
555 .metadata
556 .insert("gold.last_evaluation".to_string(), summary.to_string());
557 session.metadata.insert(
558 "gold.last_decision".to_string(),
559 result.decision.as_str().to_string(),
560 );
561 session.metadata.insert(
562 "gold.last_confidence".to_string(),
563 result.confidence.as_str().to_string(),
564 );
565 session
566 .metadata
567 .insert("gold.last_reasoning".to_string(), result.reasoning.clone());
568 session.metadata.insert(
569 "gold.last_checkpoint".to_string(),
570 result.checkpoint.as_str().to_string(),
571 );
572 session.metadata.insert(
573 "gold.last_iteration".to_string(),
574 result.iteration.to_string(),
575 );
576 session.metadata.insert(
577 "gold.evaluation_count".to_string(),
578 evaluation_count.to_string(),
579 );
580 session.updated_at = Utc::now();
581
582 let mut usage = MetricsTokenUsage {
583 prompt_tokens: result.prompt_tokens,
584 completion_tokens: result.completion_tokens,
585 ..Default::default()
586 };
587 usage.recompute_total();
588 usage
589}
590
591#[cfg(test)]
592mod tests {
593 use super::*;
594 use bamboo_agent_core::tools::FunctionCall;
595
596 fn report_call(arguments: serde_json::Value) -> ToolCall {
597 ToolCall {
598 id: "call-1".to_string(),
599 tool_type: "function".to_string(),
600 function: FunctionCall {
601 name: "report_gold_evaluation".to_string(),
602 arguments: arguments.to_string(),
603 },
604 }
605 }
606
607 #[test]
608 fn parse_gold_result_from_tool_calls_reads_structured_fields() {
609 let parsed = parse_gold_result_from_tool_calls(&[report_call(json!({
610 "decision": "blocked",
611 "confidence": "high",
612 "reasoning": "Missing credentials",
613 "missing_information": ["API key", " ", "Database URL"],
614 "next_action": " Ask the user for the API key "
615 }))])
616 .expect("gold result should parse");
617
618 assert_eq!(parsed.decision, GoldDecision::Blocked);
619 assert_eq!(parsed.confidence, GoldConfidence::High);
620 assert_eq!(parsed.reasoning, "Missing credentials");
621 assert_eq!(
622 parsed.missing_information,
623 vec!["API key".to_string(), "Database URL".to_string()]
624 );
625 assert_eq!(
626 parsed.next_action.as_deref(),
627 Some("Ask the user for the API key")
628 );
629 }
630
631 #[test]
632 fn parse_gold_result_from_tool_calls_ignores_other_tools() {
633 let parsed = parse_gold_result_from_tool_calls(&[ToolCall {
634 id: "call-1".to_string(),
635 tool_type: "function".to_string(),
636 function: FunctionCall {
637 name: "other_tool".to_string(),
638 arguments: "{}".to_string(),
639 },
640 }]);
641
642 assert!(parsed.is_none());
643 }
644
645 #[test]
646 fn apply_gold_evaluation_result_updates_metadata_keys() {
647 let mut session = Session::new("session-1", "model");
648 let result = GoldEvaluationResult {
649 checkpoint: GoldCheckpoint::PostRound,
650 iteration: 2,
651 decision: GoldDecision::Achieved,
652 confidence: GoldConfidence::Medium,
653 reasoning: "Goal satisfied".to_string(),
654 missing_information: Vec::new(),
655 next_action: None,
656 prompt_tokens: 10,
657 completion_tokens: 5,
658 };
659
660 let usage = apply_gold_evaluation_result(&mut session, &result);
661
662 assert_eq!(
663 session
664 .metadata
665 .get("gold.last_decision")
666 .map(String::as_str),
667 Some("achieved")
668 );
669 assert_eq!(
670 session
671 .metadata
672 .get("gold.last_confidence")
673 .map(String::as_str),
674 Some("medium")
675 );
676 assert_eq!(
677 session
678 .metadata
679 .get("gold.last_checkpoint")
680 .map(String::as_str),
681 Some("post_round")
682 );
683 assert_eq!(
684 session
685 .metadata
686 .get("gold.evaluation_count")
687 .map(String::as_str),
688 Some("1")
689 );
690 assert_eq!(usage.prompt_tokens, 10);
691 assert_eq!(usage.completion_tokens, 5);
692 assert_eq!(usage.total_tokens, 15);
693 }
694}