1use std::future::Future;
9use std::pin::Pin;
10use std::sync::Arc;
11
12use async_trait::async_trait;
13use deepstrike_core::context::renderer::RenderedContext;
14use deepstrike_core::harness::eval::SkillCandidate;
15use deepstrike_core::orchestration::workflow::WorkflowNode;
16use deepstrike_core::types::message::{Message, Role};
17use futures::{Stream, StreamExt};
18
19use crate::harness::{Criterion, Verdict};
20use crate::providers::{LLMProvider, StreamEvent};
21use crate::run_event::RunEvent;
22use crate::runtime::RuntimeRunner;
23use crate::tools::ToolChunk;
24use crate::{Error, Result};
25
26pub type AttemptBodyStream<'a> = Pin<Box<dyn Stream<Item = Result<AttemptBodyEvent>> + 'a>>;
27pub type AttemptLoopStream<'a> = Pin<Box<dyn Stream<Item = Result<AttemptLoopEvent>> + 'a>>;
28
29#[derive(Debug, Clone)]
30pub struct AttemptRequest {
31 pub session_id: String,
32 pub goal: String,
33 pub criteria: Vec<Criterion>,
34 pub extensions: Option<serde_json::Value>,
35}
36
37impl AttemptRequest {
38 pub fn new(session_id: impl Into<String>, goal: impl Into<String>) -> Self {
39 Self {
40 session_id: session_id.into(),
41 goal: goal.into(),
42 criteria: Vec::new(),
43 extensions: None,
44 }
45 }
46
47 pub fn generated(goal: impl Into<String>) -> Self {
48 Self::new(uuid::Uuid::new_v4().to_string(), goal)
49 }
50}
51
52#[derive(Debug, Clone)]
53pub struct AttemptBodyContext {
54 pub session_id: String,
55 pub goal: String,
56 pub criteria: Vec<Criterion>,
57 pub extensions: Option<serde_json::Value>,
58 pub attempt: u32,
59 pub context_input: Option<String>,
61}
62
63#[derive(Debug, Clone)]
64pub enum AttemptBodyEvent {
65 Token(String),
66 ToolCall {
67 id: String,
68 name: String,
69 },
70 ToolDelta {
71 call_id: String,
72 name: String,
73 chunk: ToolChunk,
74 },
75 ToolSuspend {
76 call_id: String,
77 name: String,
78 suspension_id: String,
79 payload: Option<serde_json::Value>,
80 },
81 ToolResult {
82 call_id: String,
83 content: String,
84 is_error: bool,
85 },
86 WorkflowNodesSubmitted(Vec<WorkflowNode>),
87 BodyError(String),
88 BodyDone {
89 run_status: String,
90 result: String,
91 turns: u32,
92 total_tokens: u64,
93 },
94}
95
96pub trait AttemptBody: Send + Sync {
97 fn run<'a>(&'a self, context: AttemptBodyContext) -> AttemptBodyStream<'a>;
98}
99
100pub struct RuntimeAttemptBody<'a> {
102 runner: &'a RuntimeRunner,
103}
104
105impl<'a> RuntimeAttemptBody<'a> {
106 pub fn new(runner: &'a RuntimeRunner) -> Self {
107 Self { runner }
108 }
109}
110
111impl AttemptBody for RuntimeAttemptBody<'_> {
112 fn run<'a>(&'a self, context: AttemptBodyContext) -> AttemptBodyStream<'a> {
113 Box::pin(async_stream::try_stream! {
114 let mut criteria: Vec<String> = context
115 .criteria
116 .iter()
117 .map(|criterion| criterion.text.clone())
118 .collect();
119 if let Some(note) = &context.context_input {
123 criteria.push(format!("[Attempt feedback] {note}"));
124 }
125
126 let mut result = String::new();
127 let mut terminal: Option<(u32, u64, String)> = None;
128 let stream = self.runner
129 .run_streaming(
130 &context.goal,
131 &criteria,
132 context.extensions.as_ref(),
133 Some(&context.session_id),
134 )
135 .await;
136 let mut stream = match stream {
137 Ok(stream) => stream,
138 Err(error) => {
139 yield AttemptBodyEvent::BodyError(error.to_string());
140 yield AttemptBodyEvent::BodyDone {
141 run_status: "error".to_string(),
142 result,
143 turns: 0,
144 total_tokens: 0,
145 };
146 return;
147 }
148 };
149
150 while let Some(event) = stream.next().await {
151 match event {
152 Err(error) => {
153 yield AttemptBodyEvent::BodyError(error.to_string());
154 terminal = Some((0, 0, "error".to_string()));
155 break;
156 }
157 Ok(RunEvent::TextDelta(delta)) => {
158 result.push_str(&delta);
159 yield AttemptBodyEvent::Token(delta);
160 }
161 Ok(RunEvent::ToolCall { id, name }) => {
162 yield AttemptBodyEvent::ToolCall { id, name };
163 }
164 Ok(RunEvent::ToolDelta { call_id, name, chunk }) => {
165 yield AttemptBodyEvent::ToolDelta { call_id, name, chunk };
166 }
167 Ok(RunEvent::ToolSuspend { call_id, name, suspension_id, payload }) => {
168 yield AttemptBodyEvent::ToolSuspend {
169 call_id,
170 name,
171 suspension_id,
172 payload,
173 };
174 }
175 Ok(RunEvent::ToolResult { call_id, content, is_error, .. }) => {
176 yield AttemptBodyEvent::ToolResult { call_id, content, is_error };
177 }
178 Ok(RunEvent::Error(message)) => {
179 yield AttemptBodyEvent::BodyError(message);
180 terminal = Some((0, 0, "error".to_string()));
181 break;
182 }
183 Ok(RunEvent::Done { iterations, total_tokens, status }) => {
184 terminal = Some((iterations, total_tokens, status));
185 }
186 Ok(_) => {}
187 }
188 }
189
190 let (turns, total_tokens, run_status) =
191 terminal.unwrap_or_else(|| (0, 0, "error".to_string()));
192 yield AttemptBodyEvent::BodyDone {
193 run_status,
194 result,
195 turns,
196 total_tokens,
197 };
198 })
199 }
200}
201
202#[derive(Debug, Clone)]
203pub struct JudgeContext {
204 pub goal: String,
205 pub criteria: Vec<Criterion>,
206 pub attempt: u32,
207 pub result: String,
208}
209
210#[derive(Debug, Clone)]
211pub struct JudgeResult {
212 pub verdict: Verdict,
213 pub skill_candidate: Option<SkillCandidate>,
214}
215
216impl JudgeResult {
217 pub fn new(verdict: Verdict) -> Self {
218 Self {
219 verdict,
220 skill_candidate: None,
221 }
222 }
223}
224
225#[async_trait]
226pub trait AttemptJudge: Send + Sync {
227 async fn judge(&self, context: &JudgeContext) -> Result<Option<JudgeResult>>;
229}
230
231pub type VerdictFn = Arc<dyn Fn(&JudgeContext) -> Option<Verdict> + Send + Sync>;
232
233pub struct VerdictFnJudge {
234 verdict_fn: VerdictFn,
235}
236
237impl VerdictFnJudge {
238 pub fn new(verdict_fn: VerdictFn) -> Self {
239 Self { verdict_fn }
240 }
241}
242
243#[async_trait]
244impl AttemptJudge for VerdictFnJudge {
245 async fn judge(&self, context: &JudgeContext) -> Result<Option<JudgeResult>> {
246 Ok((self.verdict_fn)(context).map(JudgeResult::new))
247 }
248}
249
250pub struct LlmEvalJudge {
251 eval_provider: Box<dyn LLMProvider>,
252 extract_skill_on_pass: bool,
253}
254
255impl LlmEvalJudge {
256 pub fn new(eval_provider: impl LLMProvider + 'static) -> Self {
257 Self {
258 eval_provider: Box::new(eval_provider),
259 extract_skill_on_pass: false,
260 }
261 }
262
263 pub fn extract_skill_on_pass(mut self, enabled: bool) -> Self {
264 self.extract_skill_on_pass = enabled;
265 self
266 }
267}
268
269#[async_trait]
270impl AttemptJudge for LlmEvalJudge {
271 async fn judge(&self, context: &JudgeContext) -> Result<Option<JudgeResult>> {
272 use deepstrike_core::harness::eval::{
273 build_eval_messages, parse_verdict, Criterion as EvalCriterion,
274 };
275
276 let criteria: Vec<EvalCriterion> = context
277 .criteria
278 .iter()
279 .map(|criterion| EvalCriterion {
280 text: criterion.text.clone(),
281 required: criterion.required,
282 weight: criterion.weight,
283 })
284 .collect();
285 let messages = build_eval_messages(
286 &context.goal,
287 &criteria,
288 &context.result,
289 context.attempt,
290 self.extract_skill_on_pass,
291 );
292 let rendered = rendered_context_from_messages(messages);
293 let state = self.eval_provider.create_run_state();
294 let mut stream = self
295 .eval_provider
296 .stream(&rendered, &[], None, state.as_ref())
297 .await?;
298 let mut text = String::new();
299 while let Some(event) = stream.next().await {
300 if let StreamEvent::TextDelta { delta } = event? {
301 text.push_str(&delta);
302 }
303 }
304 if text.is_empty() {
305 return Err(Error::Other("attempt judge produced no text".to_string()));
306 }
307
308 let parsed = parse_verdict(&text);
309 Ok(Some(JudgeResult {
310 verdict: Verdict {
311 passed: parsed.passed,
312 overall_score: parsed.overall_score,
313 feedback: parsed.feedback,
314 details: parsed
315 .details
316 .into_iter()
317 .map(|detail| crate::harness::CriterionResult {
318 criterion: detail.criterion,
319 passed: detail.passed,
320 score: detail.score,
321 feedback: detail.feedback,
322 })
323 .collect(),
324 },
325 skill_candidate: parsed.skill_candidate,
326 }))
327 }
328}
329
330pub struct HybridJudge<P, F> {
331 primary: P,
332 fallback: F,
333}
334
335impl<P, F> HybridJudge<P, F> {
336 pub fn new(primary: P, fallback: F) -> Self {
337 Self { primary, fallback }
338 }
339}
340
341#[async_trait]
342impl<P, F> AttemptJudge for HybridJudge<P, F>
343where
344 P: AttemptJudge,
345 F: AttemptJudge,
346{
347 async fn judge(&self, context: &JudgeContext) -> Result<Option<JudgeResult>> {
348 match self.primary.judge(context).await? {
349 Some(result) => Ok(Some(result)),
350 None => self.fallback.judge(context).await,
351 }
352 }
353}
354
355#[derive(Debug, Clone)]
356pub struct CarryContext {
357 pub root_session_id: String,
358 pub goal: String,
359 pub attempt: u32,
360 pub previous_verdict: Option<Verdict>,
361}
362
363#[derive(Debug, Clone)]
364pub struct PreparedAttempt {
365 pub session_id: String,
366 pub goal: String,
367 pub context_input: Option<String>,
368}
369
370#[async_trait]
371pub trait CarryPolicy: Send + Sync {
372 async fn prepare(&self, context: CarryContext) -> Result<PreparedAttempt>;
373}
374
375#[derive(Debug, Clone, Copy, Default)]
377pub struct ContinueSession;
378
379#[async_trait]
380impl CarryPolicy for ContinueSession {
381 async fn prepare(&self, context: CarryContext) -> Result<PreparedAttempt> {
382 Ok(PreparedAttempt {
383 session_id: context.root_session_id,
384 goal: context.goal,
385 context_input: context
386 .previous_verdict
387 .map(|verdict| verdict.feedback)
388 .filter(|feedback| !feedback.is_empty()),
389 })
390 }
391}
392
393#[derive(Debug, Clone, Copy, Default)]
395pub struct FreshWithFeedback;
396
397#[async_trait]
398impl CarryPolicy for FreshWithFeedback {
399 async fn prepare(&self, context: CarryContext) -> Result<PreparedAttempt> {
400 let goal = match context.previous_verdict {
401 Some(verdict) if !verdict.feedback.is_empty() => format!(
402 "{}\n\n[Attempt {} feedback: {}]",
403 context.goal,
404 context.attempt.saturating_sub(1),
405 verdict.feedback
406 ),
407 _ => context.goal,
408 };
409 Ok(PreparedAttempt {
410 session_id: if context.attempt == 1 {
411 context.root_session_id
412 } else {
413 uuid::Uuid::new_v4().to_string()
414 },
415 goal,
416 context_input: None,
417 })
418 }
419}
420
421pub type DigestFuture = Pin<Box<dyn Future<Output = Result<String>> + Send>>;
422pub type DigestFn = Arc<dyn Fn(Verdict, u32) -> DigestFuture + Send + Sync>;
423pub type PassHookFuture<'a> = Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
424pub type PassHook =
425 Arc<dyn for<'a> Fn(&'a AttemptOutcome, &'a JudgeResult) -> PassHookFuture<'a> + Send + Sync>;
426
427pub struct FreshWithDigest {
428 digest: DigestFn,
429}
430
431impl FreshWithDigest {
432 pub fn new(digest: DigestFn) -> Self {
433 Self { digest }
434 }
435}
436
437#[async_trait]
438impl CarryPolicy for FreshWithDigest {
439 async fn prepare(&self, context: CarryContext) -> Result<PreparedAttempt> {
440 let goal = if let Some(verdict) = context.previous_verdict {
441 let digest = (self.digest)(verdict, context.attempt.saturating_sub(1)).await?;
442 format!("{}\n\n[Prior attempt digest: {digest}]", context.goal)
443 } else {
444 context.goal
445 };
446 Ok(PreparedAttempt {
447 session_id: if context.attempt == 1 {
448 context.root_session_id
449 } else {
450 uuid::Uuid::new_v4().to_string()
451 },
452 goal,
453 context_input: None,
454 })
455 }
456}
457
458#[derive(Debug, Clone)]
459pub struct StopPolicy {
460 pub max_attempts: u32,
461 pub max_total_tokens: Option<u64>,
462 pub stop_on_failed_verdict: bool,
463}
464
465impl StopPolicy {
466 pub fn new(max_attempts: u32) -> Self {
467 Self {
468 max_attempts,
469 max_total_tokens: None,
470 stop_on_failed_verdict: false,
471 }
472 }
473
474 pub fn max_total_tokens(mut self, limit: u64) -> Self {
475 self.max_total_tokens = Some(limit);
476 self
477 }
478
479 pub fn stop_on_failed_verdict(mut self, enabled: bool) -> Self {
480 self.stop_on_failed_verdict = enabled;
481 self
482 }
483}
484
485#[derive(Debug, Clone, Copy, PartialEq, Eq)]
486pub enum AttemptOutcomeKind {
487 Passed,
488 FailedJudge,
489 Exhausted,
490 RunError,
491}
492
493#[derive(Debug, Clone)]
494pub struct AttemptOutcome {
495 pub outcome: AttemptOutcomeKind,
496 pub run_status: String,
497 pub verdict: Option<Verdict>,
498 pub result: String,
499 pub attempts: u32,
500 pub turns: u32,
501 pub total_tokens: u64,
502 pub submitted_nodes: Option<Vec<WorkflowNode>>,
503}
504
505#[derive(Debug, Clone)]
506pub enum AttemptLoopEvent {
507 Token(String),
508 ToolCall {
509 id: String,
510 name: String,
511 },
512 ToolDelta {
513 call_id: String,
514 name: String,
515 chunk: ToolChunk,
516 },
517 ToolSuspend {
518 call_id: String,
519 name: String,
520 suspension_id: String,
521 payload: Option<serde_json::Value>,
522 },
523 ToolResult {
524 call_id: String,
525 content: String,
526 is_error: bool,
527 },
528 WorkflowNodesSubmitted(Vec<WorkflowNode>),
529 BodyError(String),
530 Judging {
531 attempt: u32,
532 },
533 Retrying {
534 attempt: u32,
535 verdict: Verdict,
536 },
537 Completed(AttemptOutcome),
538}
539
540pub struct AttemptLoop<B, J, C = ContinueSession> {
541 body: B,
542 judge: J,
543 carry: C,
544 stop: StopPolicy,
545 on_pass: Option<PassHook>,
546}
547
548impl<B, J> AttemptLoop<B, J, ContinueSession> {
549 pub fn new(body: B, judge: J, stop: StopPolicy) -> Result<Self> {
550 validate_stop_policy(&stop)?;
551 Ok(Self {
552 body,
553 judge,
554 carry: ContinueSession,
555 stop,
556 on_pass: None,
557 })
558 }
559}
560
561impl<B, J, C> AttemptLoop<B, J, C> {
562 pub fn with_carry<Next>(self, carry: Next) -> AttemptLoop<B, J, Next> {
563 AttemptLoop {
564 body: self.body,
565 judge: self.judge,
566 carry,
567 stop: self.stop,
568 on_pass: self.on_pass,
569 }
570 }
571
572 pub fn with_on_pass(mut self, hook: PassHook) -> Self {
574 self.on_pass = Some(hook);
575 self
576 }
577}
578
579impl<B, J, C> AttemptLoop<B, J, C>
580where
581 B: AttemptBody,
582 J: AttemptJudge,
583 C: CarryPolicy,
584{
585 pub async fn run(&self, request: AttemptRequest) -> Result<AttemptOutcome> {
586 let mut outcome = None;
587 let mut stream = self.stream(request);
588 while let Some(event) = stream.next().await {
589 if let AttemptLoopEvent::Completed(completed) = event? {
590 outcome = Some(completed);
591 }
592 }
593 outcome.ok_or_else(|| Error::Other("AttemptLoop ended without an outcome".to_string()))
594 }
595
596 pub fn stream<'a>(&'a self, request: AttemptRequest) -> AttemptLoopStream<'a> {
597 Box::pin(async_stream::try_stream! {
598 let root_session_id = request.session_id.clone();
599 let mut previous_verdict: Option<Verdict> = None;
600 let mut total_tokens = 0u64;
601 let mut total_turns = 0u32;
602 let mut submitted_nodes = Vec::new();
603
604 for attempt in 1..=self.stop.max_attempts {
605 let prepared = self.carry
606 .prepare(CarryContext {
607 root_session_id: root_session_id.clone(),
608 goal: request.goal.clone(),
609 attempt,
610 previous_verdict: previous_verdict.clone(),
611 })
612 .await?;
613 let mut body = self.body.run(AttemptBodyContext {
614 session_id: prepared.session_id,
615 goal: prepared.goal,
616 criteria: request.criteria.clone(),
617 extensions: request.extensions.clone(),
618 attempt,
619 context_input: prepared.context_input,
620 });
621 let mut terminal = None;
622 while let Some(event) = body.next().await {
623 match event? {
624 AttemptBodyEvent::Token(text) => yield AttemptLoopEvent::Token(text),
625 AttemptBodyEvent::ToolCall { id, name } => {
626 yield AttemptLoopEvent::ToolCall { id, name };
627 }
628 AttemptBodyEvent::ToolDelta { call_id, name, chunk } => {
629 yield AttemptLoopEvent::ToolDelta { call_id, name, chunk };
630 }
631 AttemptBodyEvent::ToolSuspend {
632 call_id,
633 name,
634 suspension_id,
635 payload,
636 } => {
637 yield AttemptLoopEvent::ToolSuspend {
638 call_id,
639 name,
640 suspension_id,
641 payload,
642 };
643 }
644 AttemptBodyEvent::ToolResult { call_id, content, is_error } => {
645 yield AttemptLoopEvent::ToolResult { call_id, content, is_error };
646 }
647 AttemptBodyEvent::WorkflowNodesSubmitted(nodes) => {
648 submitted_nodes.extend(nodes.iter().cloned());
649 yield AttemptLoopEvent::WorkflowNodesSubmitted(nodes);
650 }
651 AttemptBodyEvent::BodyError(message) => {
652 yield AttemptLoopEvent::BodyError(message);
653 }
654 AttemptBodyEvent::BodyDone {
655 run_status,
656 result,
657 turns,
658 total_tokens,
659 } => terminal = Some((run_status, result, turns, total_tokens)),
660 }
661 }
662 let (run_status, result, turns, attempt_tokens) = terminal.ok_or_else(|| {
663 Error::Other("AttemptBody ended without body_done".to_string())
664 })?;
665 total_tokens = total_tokens.saturating_add(attempt_tokens);
666 total_turns = total_turns.saturating_add(turns);
667
668 if is_run_error(&run_status) {
669 yield AttemptLoopEvent::Completed(AttemptOutcome {
670 outcome: AttemptOutcomeKind::RunError,
671 run_status,
672 verdict: None,
673 result,
674 attempts: attempt,
675 turns: total_turns,
676 total_tokens,
677 submitted_nodes: (!submitted_nodes.is_empty())
678 .then(|| submitted_nodes.clone()),
679 });
680 return;
681 }
682
683 yield AttemptLoopEvent::Judging { attempt };
684 let judged = self
685 .judge
686 .judge(&JudgeContext {
687 goal: request.goal.clone(),
688 criteria: request.criteria.clone(),
689 attempt,
690 result: result.clone(),
691 })
692 .await?
693 .ok_or_else(|| Error::Other("AttemptLoop judge produced no verdict".to_string()))?;
694 let verdict = judged.verdict.clone();
695
696 if verdict.passed {
697 let outcome = AttemptOutcome {
698 outcome: AttemptOutcomeKind::Passed,
699 run_status,
700 verdict: Some(verdict),
701 result,
702 attempts: attempt,
703 turns: total_turns,
704 total_tokens,
705 submitted_nodes: (!submitted_nodes.is_empty())
706 .then(|| submitted_nodes.clone()),
707 };
708 if let Some(on_pass) = &self.on_pass {
709 on_pass(&outcome, &judged).await?;
710 }
711 yield AttemptLoopEvent::Completed(outcome);
712 return;
713 }
714
715 previous_verdict = Some(verdict.clone());
716 let token_limit_reached = self
717 .stop
718 .max_total_tokens
719 .is_some_and(|limit| total_tokens >= limit);
720 if self.stop.stop_on_failed_verdict
721 || attempt == self.stop.max_attempts
722 || token_limit_reached
723 {
724 yield AttemptLoopEvent::Completed(AttemptOutcome {
725 outcome: if self.stop.stop_on_failed_verdict {
726 AttemptOutcomeKind::FailedJudge
727 } else {
728 AttemptOutcomeKind::Exhausted
729 },
730 run_status,
731 verdict: Some(verdict),
732 result,
733 attempts: attempt,
734 turns: total_turns,
735 total_tokens,
736 submitted_nodes: (!submitted_nodes.is_empty())
737 .then(|| submitted_nodes.clone()),
738 });
739 return;
740 }
741
742 yield AttemptLoopEvent::Retrying { attempt, verdict };
743 }
744 })
745 }
746}
747
748fn validate_stop_policy(stop: &StopPolicy) -> Result<()> {
749 if stop.max_attempts == 0 {
750 return Err(Error::Other(
751 "AttemptLoop stop.max_attempts must be positive".to_string(),
752 ));
753 }
754 Ok(())
755}
756
757fn is_run_error(status: &str) -> bool {
758 matches!(
759 status.to_ascii_lowercase().as_str(),
760 "error" | "invalid_arg" | "user_abort"
761 )
762}
763
764fn rendered_context_from_messages(messages: Vec<Message>) -> RenderedContext {
765 let mut system_parts = Vec::new();
766 let mut turns = Vec::new();
767 for message in messages {
768 if message.role == Role::System {
769 if let Some(text) = message.content.as_text() {
770 system_parts.push(text.to_owned());
771 }
772 } else {
773 turns.push(message);
774 }
775 }
776 let system_text = system_parts.join("\n\n");
777 RenderedContext {
778 system_text: system_text.clone(),
779 system_stable: system_text,
780 system_knowledge: String::new(),
781 turns,
782 state_turn: None,
783 frozen_prefix_len: None,
784 budget_overflow: None,
785 }
786}