rskit_agent/agent/
stream.rs

1use std::pin::Pin;
2
3use async_stream::stream;
4use futures::Stream;
5use rskit_hook::CancellationToken;
6use rskit_llm::types::{Message, Usage};
7
8use super::Agent;
9use crate::runtime::state::RunState;
10use crate::runtime::stop;
11use crate::runtime::turn::{self, TurnOutcome};
12use crate::types::{AgentEvent, StopReason};
13
14impl Agent {
15    /// Stream the agent loop, yielding [`AgentEvent`]s for each lifecycle point.
16    pub fn stream(
17        &self,
18        messages: Vec<Message>,
19    ) -> Pin<Box<dyn Stream<Item = AgentEvent> + Send + '_>> {
20        Box::pin(stream! {
21            let mut state = RunState::new(&self.config.system_prompt, messages);
22
23            if let Some(stop_reason) = stop::initial_stop(&self.config) {
24                yield AgentEvent::Complete {
25                    result: state.finish(0, stop_reason),
26                };
27                return;
28            }
29
30            let hook_token = CancellationToken::new();
31
32            for turn in 0..self.config.max_turns {
33                if let Some(stop_reason) = stop::wall_clock_stop(&state, &self.config) {
34                    yield AgentEvent::Complete {
35                        result: state.finish(turn, stop_reason),
36                    };
37                    return;
38                }
39
40                let usage_before = state.total_usage;
41                yield AgentEvent::TurnStart { turn };
42
43                let outcome = match turn::run_turn(
44                    self.provider.as_ref(),
45                    &self.config,
46                    &mut state,
47                    turn,
48                    &hook_token,
49                )
50                .await
51                {
52                    Ok(outcome) => outcome,
53                    Err(error) => {
54                        tracing::error!(error = %error, "agent.run.failed");
55                        yield AgentEvent::Complete {
56                            result: state.finish(turn, StopReason::Aborted),
57                        };
58                        return;
59                    }
60                };
61
62                let usage = usage_delta(usage_before, state.total_usage);
63                let completed_turn = match &outcome {
64                    TurnOutcome::Continue => true,
65                    TurnOutcome::Stop { turn_count, .. } => *turn_count > turn,
66                };
67
68                if completed_turn {
69                    yield AgentEvent::TurnComplete {
70                        turn,
71                        message: state.last_assistant.clone(),
72                        usage,
73                    };
74                }
75
76                match outcome {
77                    TurnOutcome::Continue => {}
78                    TurnOutcome::Stop { turn_count, reason } => {
79                        yield AgentEvent::Complete {
80                            result: state.finish(turn_count, reason),
81                        };
82                        return;
83                    }
84                }
85            }
86
87            yield AgentEvent::Complete {
88                result: state.finish(self.config.max_turns, StopReason::MaxTurns),
89            };
90        })
91    }
92}
93
94const fn usage_delta(before: Usage, after: Usage) -> Usage {
95    Usage {
96        input_tokens: after.input_tokens.saturating_sub(before.input_tokens),
97        output_tokens: after.output_tokens.saturating_sub(before.output_tokens),
98        cached_tokens: after.cached_tokens.saturating_sub(before.cached_tokens),
99        reasoning_tokens: after
100            .reasoning_tokens
101            .saturating_sub(before.reasoning_tokens),
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use rskit_llm::types::Usage;
108
109    use super::usage_delta;
110
111    #[test]
112    fn usage_delta_reports_per_turn_usage() {
113        let before = Usage {
114            input_tokens: 10,
115            output_tokens: 5,
116            cached_tokens: 2,
117            reasoning_tokens: 1,
118        };
119        let after = Usage {
120            input_tokens: 17,
121            output_tokens: 9,
122            cached_tokens: 3,
123            reasoning_tokens: 6,
124        };
125
126        assert_eq!(
127            usage_delta(before, after),
128            Usage {
129                input_tokens: 7,
130                output_tokens: 4,
131                cached_tokens: 1,
132                reasoning_tokens: 5,
133            }
134        );
135    }
136
137    #[test]
138    fn usage_delta_saturates_if_counters_reset() {
139        let before = Usage {
140            input_tokens: 10,
141            output_tokens: 5,
142            cached_tokens: 2,
143            reasoning_tokens: 1,
144        };
145        let after = Usage {
146            input_tokens: 7,
147            output_tokens: 9,
148            cached_tokens: 1,
149            reasoning_tokens: 6,
150        };
151
152        assert_eq!(
153            usage_delta(before, after),
154            Usage {
155                input_tokens: 0,
156                output_tokens: 4,
157                cached_tokens: 0,
158                reasoning_tokens: 5,
159            }
160        );
161    }
162}
rskit_agent/agent/stream.rs

rskit_agent/agent/
stream.rs
