Skip to main content

vv_agent/runner/
event_stream.rs

1use serde_json::Value;
2
3use crate::events::{RunEvent, ToolStatus};
4use crate::result::RunResult;
5
6pub struct RunEventStream {
7    events: std::vec::IntoIter<Result<RunEvent, String>>,
8    result: Option<RunResult>,
9}
10
11impl RunEventStream {
12    pub(super) fn from_events(result: RunResult, events: Vec<RunEvent>) -> Self {
13        let events = events.into_iter().map(Ok).collect::<Vec<_>>().into_iter();
14        Self {
15            events,
16            result: Some(result),
17        }
18    }
19
20    pub async fn next(&mut self) -> Option<Result<RunEvent, String>> {
21        self.events.next()
22    }
23
24    pub async fn into_result(mut self) -> Result<RunResult, String> {
25        self.result
26            .take()
27            .ok_or_else(|| "stream result already taken".to_string())
28    }
29}
30
31pub(super) fn map_runtime_event(
32    event: &str,
33    payload: &std::collections::BTreeMap<String, Value>,
34) -> Option<RunEvent> {
35    let run_id = payload
36        .get("task_id")
37        .and_then(Value::as_str)
38        .or_else(|| payload.get("run_id").and_then(Value::as_str))
39        .unwrap_or("run")
40        .to_string();
41    match event {
42        "run_started" => Some(RunEvent::RunStarted {
43            run_id,
44            agent_name: payload
45                .get("agent_name")
46                .and_then(Value::as_str)
47                .unwrap_or("agent")
48                .to_string(),
49        }),
50        "cycle_started" => Some(RunEvent::AgentStarted {
51            run_id,
52            agent_name: payload
53                .get("agent_name")
54                .and_then(Value::as_str)
55                .unwrap_or("agent")
56                .to_string(),
57            cycle_index: payload
58                .get("cycle")
59                .and_then(Value::as_u64)
60                .unwrap_or_default() as u32,
61        }),
62        "cycle_llm_response" => {
63            let cycle_index = payload
64                .get("cycle")
65                .and_then(Value::as_u64)
66                .unwrap_or_default() as u32;
67            payload
68                .get("assistant_message")
69                .and_then(Value::as_str)
70                .filter(|delta| !delta.is_empty())
71                .map(|delta| RunEvent::AssistantDelta {
72                    run_id,
73                    delta: delta.to_string(),
74                    cycle_index,
75                })
76        }
77        "tool_result" => {
78            let metadata = payload.get("metadata").and_then(Value::as_object);
79            if let Some(interruption_id) = metadata
80                .and_then(|metadata| metadata.get("approval_interruption_id"))
81                .and_then(Value::as_str)
82            {
83                return Some(RunEvent::ToolApprovalRequested {
84                    run_id,
85                    interruption_id: interruption_id.to_string(),
86                    tool_name: payload
87                        .get("tool_name")
88                        .and_then(Value::as_str)
89                        .unwrap_or_default()
90                        .to_string(),
91                });
92            }
93            let status = match payload
94                .get("status")
95                .and_then(Value::as_str)
96                .unwrap_or_default()
97            {
98                "error" | "ERROR" => ToolStatus::Error,
99                "wait_response" | "WAIT_RESPONSE" => ToolStatus::WaitResponse,
100                _ => ToolStatus::Success,
101            };
102            Some(RunEvent::ToolFinished {
103                run_id,
104                tool_call_id: payload
105                    .get("tool_call_id")
106                    .and_then(Value::as_str)
107                    .unwrap_or_default()
108                    .to_string(),
109                tool_name: payload
110                    .get("tool_name")
111                    .and_then(Value::as_str)
112                    .unwrap_or_default()
113                    .to_string(),
114                status,
115            })
116        }
117        "run_completed" => Some(RunEvent::RunCompleted {
118            run_id,
119            status: crate::types::AgentStatus::Completed,
120        }),
121        "run_wait_user" => Some(RunEvent::RunCompleted {
122            run_id,
123            status: crate::types::AgentStatus::WaitUser,
124        }),
125        "run_max_cycles" => Some(RunEvent::RunCompleted {
126            run_id,
127            status: crate::types::AgentStatus::MaxCycles,
128        }),
129        "cycle_failed" => Some(RunEvent::RunFailed {
130            run_id,
131            error: crate::events::AgentErrorPayload::new(
132                payload
133                    .get("error")
134                    .and_then(Value::as_str)
135                    .unwrap_or("cycle failed"),
136            ),
137        }),
138        _ => None,
139    }
140}