vtcode_exec_events/
atif.rs

1//! Agent Trajectory Interchange Format (ATIF) types and builder.
2//!
3//! Implements the [ATIF specification](https://github.com/laude-institute/harbor/blob/main/docs/rfcs/0001-trajectory-format.md)
4//! v1.4 for logging complete agent interaction histories in a standardized,
5//! JSON-based format usable across debugging, visualization, SFT, and RL
6//! pipelines.
7//!
8//! # Overview
9//!
10//! ATIF provides a complete session trajectory: user messages, agent responses,
11//! tool executions, observations, and per-step/aggregate LLM metrics. The
12//! [`AtifTrajectoryBuilder`] converts live [`ThreadEvent`]
13//! streams into a finished [`Trajectory`].
14//!
15//! # Example
16//!
17//! ```rust
18//! use vtcode_exec_events::atif::*;
19//!
20//! let agent = AtifAgent::new("vtcode", env!("CARGO_PKG_VERSION"));
21//! let mut builder = AtifTrajectoryBuilder::new(agent);
22//!
23//! // Feed ThreadEvents as they arrive …
24//! // builder.process_event(&event);
25//!
26//! let trajectory = builder.finish(None);
27//! let json = serde_json::to_string_pretty(&trajectory).unwrap();
28//! ```
29
30use chrono::{DateTime, Utc};
31use serde::{Deserialize, Serialize};
32use serde_json::Value;
33
34use crate::{ThreadEvent, ThreadItemDetails, ToolCallStatus};
35
36/// Current ATIF schema version supported by this implementation.
37pub const ATIF_SCHEMA_VERSION: &str = "ATIF-v1.4";
38
39// ============================================================================
40// Core ATIF Types
41// ============================================================================
42
43/// Root-level ATIF trajectory object.
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct Trajectory {
46    /// ATIF schema version (e.g., "ATIF-v1.4").
47    pub schema_version: String,
48    /// Unique identifier for the entire agent run.
49    pub session_id: String,
50    /// Agent configuration for this trajectory.
51    pub agent: AtifAgent,
52    /// Ordered interaction steps.
53    pub steps: Vec<Step>,
54    /// Optional developer notes.
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub notes: Option<String>,
57    /// Aggregate metrics for the full trajectory.
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub final_metrics: Option<FinalMetrics>,
60    /// Optional custom root-level metadata.
61    #[serde(skip_serializing_if = "Option::is_none")]
62    pub extra: Option<Value>,
63}
64
65/// Agent configuration metadata.
66#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct AtifAgent {
68    /// Agent system name (e.g., "vtcode").
69    pub name: String,
70    /// Agent system version.
71    pub version: String,
72    /// Default LLM model used. Step-level `model_name` overrides this.
73    #[serde(skip_serializing_if = "Option::is_none")]
74    pub model_name: Option<String>,
75    /// Optional custom agent metadata.
76    #[serde(skip_serializing_if = "Option::is_none")]
77    pub extra: Option<Value>,
78}
79
80impl AtifAgent {
81    /// Create a new agent descriptor.
82    pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
83        Self {
84            name: name.into(),
85            version: version.into(),
86            model_name: None,
87            extra: None,
88        }
89    }
90
91    /// Create a vtcode agent descriptor using the crate version.
92    pub fn vtcode() -> Self {
93        Self::new("vtcode", env!("CARGO_PKG_VERSION"))
94    }
95
96    /// Set the default model name.
97    pub fn with_model(mut self, model: impl Into<String>) -> Self {
98        self.model_name = Some(model.into());
99        self
100    }
101}
102
103/// The originator of a step.
104#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
105#[serde(rename_all = "lowercase")]
106pub enum StepSource {
107    /// System prompt or system-initiated operation.
108    System,
109    /// User message.
110    User,
111    /// Agent response.
112    Agent,
113}
114
115/// Individual interaction step.
116#[derive(Debug, Clone, Serialize, Deserialize)]
117pub struct Step {
118    /// Ordinal index (starting from 1).
119    pub step_id: u64,
120    /// ISO 8601 timestamp.
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub timestamp: Option<String>,
123    /// Originator of this step.
124    pub source: StepSource,
125    /// LLM model used for this step (agent steps only).
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub model_name: Option<String>,
128    /// Step content — text message or array.
129    #[serde(skip_serializing_if = "Option::is_none")]
130    pub message: Option<String>,
131    /// Agent internal reasoning content.
132    #[serde(skip_serializing_if = "Option::is_none")]
133    pub reasoning_content: Option<String>,
134    /// Tool/function invocations (agent steps only).
135    #[serde(skip_serializing_if = "Option::is_none")]
136    pub tool_calls: Option<Vec<AtifToolCall>>,
137    /// Environment feedback after actions.
138    #[serde(skip_serializing_if = "Option::is_none")]
139    pub observation: Option<Observation>,
140    /// LLM operational metrics (agent steps only).
141    #[serde(skip_serializing_if = "Option::is_none")]
142    pub metrics: Option<StepMetrics>,
143    /// Custom step-level metadata.
144    #[serde(skip_serializing_if = "Option::is_none")]
145    pub extra: Option<Value>,
146}
147
148impl Step {
149    /// Create a user step.
150    pub fn user(step_id: u64, message: impl Into<String>) -> Self {
151        Self {
152            step_id,
153            timestamp: Some(Utc::now().to_rfc3339()),
154            source: StepSource::User,
155            model_name: None,
156            message: Some(message.into()),
157            reasoning_content: None,
158            tool_calls: None,
159            observation: None,
160            metrics: None,
161            extra: None,
162        }
163    }
164
165    /// Create an agent step.
166    pub fn agent(step_id: u64, message: impl Into<String>) -> Self {
167        Self {
168            step_id,
169            timestamp: Some(Utc::now().to_rfc3339()),
170            source: StepSource::Agent,
171            model_name: None,
172            message: Some(message.into()),
173            reasoning_content: None,
174            tool_calls: None,
175            observation: None,
176            metrics: None,
177            extra: None,
178        }
179    }
180
181    /// Create a system step.
182    pub fn system(step_id: u64, message: impl Into<String>) -> Self {
183        Self {
184            step_id,
185            timestamp: Some(Utc::now().to_rfc3339()),
186            source: StepSource::System,
187            model_name: None,
188            message: Some(message.into()),
189            reasoning_content: None,
190            tool_calls: None,
191            observation: None,
192            metrics: None,
193            extra: None,
194        }
195    }
196}
197
198/// Structured tool/function invocation.
199#[derive(Debug, Clone, Serialize, Deserialize)]
200pub struct AtifToolCall {
201    /// Unique identifier for the tool call.
202    pub tool_call_id: String,
203    /// Function/tool name.
204    pub function_name: String,
205    /// Arguments passed to the tool.
206    #[serde(skip_serializing_if = "Option::is_none")]
207    pub arguments: Option<Value>,
208}
209
210/// Environment feedback container.
211#[derive(Debug, Clone, Serialize, Deserialize)]
212pub struct Observation {
213    /// Results from tool calls or system operations.
214    pub results: Vec<ObservationResult>,
215}
216
217/// Individual observation result tied to a tool call.
218#[derive(Debug, Clone, Serialize, Deserialize)]
219pub struct ObservationResult {
220    /// Identifier of the originating tool call.
221    pub source_call_id: String,
222    /// Content/output of the observation.
223    pub content: String,
224}
225
226/// Per-step LLM operational metrics.
227#[derive(Debug, Clone, Serialize, Deserialize)]
228pub struct StepMetrics {
229    /// Total input tokens for this step (cached + non-cached).
230    #[serde(skip_serializing_if = "Option::is_none")]
231    pub prompt_tokens: Option<u64>,
232    /// Completion tokens generated.
233    #[serde(skip_serializing_if = "Option::is_none")]
234    pub completion_tokens: Option<u64>,
235    /// Subset of prompt_tokens that were cache hits.
236    #[serde(skip_serializing_if = "Option::is_none")]
237    pub cached_tokens: Option<u64>,
238    /// Estimated cost in USD for this step.
239    #[serde(skip_serializing_if = "Option::is_none")]
240    pub cost_usd: Option<f64>,
241    /// Log probabilities for completion tokens.
242    #[serde(skip_serializing_if = "Option::is_none")]
243    pub logprobs: Option<Vec<f64>>,
244    /// Completion token IDs for RL training.
245    #[serde(skip_serializing_if = "Option::is_none")]
246    pub completion_token_ids: Option<Vec<u64>>,
247    /// Prompt token IDs.
248    #[serde(skip_serializing_if = "Option::is_none")]
249    pub prompt_token_ids: Option<Vec<u64>>,
250    /// Custom metrics.
251    #[serde(skip_serializing_if = "Option::is_none")]
252    pub extra: Option<Value>,
253}
254
255impl StepMetrics {
256    /// Create metrics from vtcode Usage.
257    pub fn from_usage(usage: &crate::Usage) -> Self {
258        Self {
259            prompt_tokens: Some(usage.input_tokens),
260            completion_tokens: Some(usage.output_tokens),
261            cached_tokens: if usage.cached_input_tokens > 0 {
262                Some(usage.cached_input_tokens)
263            } else {
264                None
265            },
266            cost_usd: None,
267            logprobs: None,
268            completion_token_ids: None,
269            prompt_token_ids: None,
270            extra: if usage.cache_creation_tokens > 0 {
271                Some(serde_json::json!({
272                    "cache_creation_tokens": usage.cache_creation_tokens
273                }))
274            } else {
275                None
276            },
277        }
278    }
279}
280
281/// Trajectory-level aggregate metrics.
282#[derive(Debug, Clone, Default, Serialize, Deserialize)]
283pub struct FinalMetrics {
284    /// Sum of all prompt tokens across steps.
285    #[serde(skip_serializing_if = "Option::is_none")]
286    pub total_prompt_tokens: Option<u64>,
287    /// Sum of all completion tokens across steps.
288    #[serde(skip_serializing_if = "Option::is_none")]
289    pub total_completion_tokens: Option<u64>,
290    /// Sum of all cached tokens across steps.
291    #[serde(skip_serializing_if = "Option::is_none")]
292    pub total_cached_tokens: Option<u64>,
293    /// Total estimated cost in USD.
294    #[serde(skip_serializing_if = "Option::is_none")]
295    pub total_cost_usd: Option<f64>,
296    /// Total number of steps.
297    #[serde(skip_serializing_if = "Option::is_none")]
298    pub total_steps: Option<u64>,
299    /// Custom aggregate metrics.
300    #[serde(skip_serializing_if = "Option::is_none")]
301    pub extra: Option<Value>,
302}
303
304// ============================================================================
305// Builder — converts live ThreadEvent streams into ATIF Trajectory
306// ============================================================================
307
308/// Stateful collector that converts a live [`ThreadEvent`] stream into an
309/// ATIF-compliant [`Trajectory`].
310///
311/// Feed events via [`process_event`](Self::process_event) (timestamps at
312/// observation time) or [`process_event_at`](Self::process_event_at)
313/// (deterministic timestamps for tests). Call [`finish`](Self::finish) to
314/// produce the final trajectory.
315pub struct AtifTrajectoryBuilder {
316    agent: AtifAgent,
317    session_id: Option<String>,
318    steps: Vec<Step>,
319    next_step_id: u64,
320    // Running token accumulators for final metrics
321    total_input_tokens: u64,
322    total_output_tokens: u64,
323    total_cached_tokens: u64,
324    num_turns: usize,
325    /// Pending tool invocations awaiting matching ToolOutput.
326    pending_tool_calls: Vec<PendingToolCall>,
327}
328
329struct PendingToolCall {
330    call_id: String,
331    tool_call_id: Option<String>,
332    tool_name: String,
333    arguments: Option<Value>,
334    timestamp: String,
335}
336
337impl AtifTrajectoryBuilder {
338    /// Create a new builder for the given agent.
339    pub fn new(agent: AtifAgent) -> Self {
340        Self {
341            agent,
342            session_id: None,
343            steps: Vec::new(),
344            next_step_id: 1,
345            total_input_tokens: 0,
346            total_output_tokens: 0,
347            total_cached_tokens: 0,
348            num_turns: 0,
349            pending_tool_calls: Vec::new(),
350        }
351    }
352
353    /// Set the session ID explicitly. If not set, it will be derived from
354    /// `ThreadStarted` or `ThreadCompleted` events.
355    pub fn set_session_id(&mut self, id: impl Into<String>) {
356        self.session_id = Some(id.into());
357    }
358
359    /// Process a thread event using the current wall-clock time.
360    pub fn process_event(&mut self, event: &ThreadEvent) {
361        self.process_event_at(event, Utc::now());
362    }
363
364    /// Process a thread event with an explicit timestamp (for deterministic tests).
365    pub fn process_event_at(&mut self, event: &ThreadEvent, ts: DateTime<Utc>) {
366        let ts_str = ts.to_rfc3339();
367        match event {
368            ThreadEvent::ThreadStarted(e) => {
369                if self.session_id.is_none() {
370                    self.session_id = Some(e.thread_id.clone());
371                }
372            }
373            ThreadEvent::ThreadCompleted(e) => {
374                if self.session_id.is_none() {
375                    self.session_id = Some(e.session_id.clone());
376                }
377                // Accumulate aggregate usage
378                self.total_input_tokens =
379                    self.total_input_tokens.saturating_add(e.usage.input_tokens);
380                self.total_output_tokens = self
381                    .total_output_tokens
382                    .saturating_add(e.usage.output_tokens);
383                self.total_cached_tokens = self
384                    .total_cached_tokens
385                    .saturating_add(e.usage.cached_input_tokens);
386                self.num_turns = e.num_turns;
387            }
388            ThreadEvent::TurnCompleted(e) => {
389                self.total_input_tokens =
390                    self.total_input_tokens.saturating_add(e.usage.input_tokens);
391                self.total_output_tokens = self
392                    .total_output_tokens
393                    .saturating_add(e.usage.output_tokens);
394                self.total_cached_tokens = self
395                    .total_cached_tokens
396                    .saturating_add(e.usage.cached_input_tokens);
397                self.num_turns += 1;
398
399                let mut step = Step::system(self.next_step_id, "turn_completed");
400                step.timestamp = Some(ts_str);
401                step.metrics = Some(StepMetrics::from_usage(&e.usage));
402                self.push_step(step);
403            }
404            ThreadEvent::TurnFailed(e) => {
405                if let Some(usage) = &e.usage {
406                    self.total_input_tokens =
407                        self.total_input_tokens.saturating_add(usage.input_tokens);
408                    self.total_output_tokens =
409                        self.total_output_tokens.saturating_add(usage.output_tokens);
410                }
411                let mut step = Step::system(self.next_step_id, &e.message);
412                step.timestamp = Some(ts_str);
413                step.metrics = e.usage.as_ref().map(StepMetrics::from_usage);
414                self.push_step(step);
415            }
416            ThreadEvent::ItemCompleted(e) => {
417                self.process_item_completed(&e.item.id, &e.item.details, &ts_str);
418            }
419            ThreadEvent::ThreadCompactBoundary(e) => {
420                let msg = format!(
421                    "context_compaction: {} messages -> {} messages ({})",
422                    e.original_message_count,
423                    e.compacted_message_count,
424                    e.trigger.as_str()
425                );
426                let mut step = Step::system(self.next_step_id, msg);
427                step.timestamp = Some(ts_str);
428                self.push_step(step);
429            }
430            ThreadEvent::Error(e) => {
431                let mut step = Step::system(self.next_step_id, &e.message);
432                step.timestamp = Some(ts_str);
433                self.push_step(step);
434            }
435            // Skip streaming/lifecycle events that don't map to ATIF steps
436            ThreadEvent::TurnStarted(_)
437            | ThreadEvent::ItemStarted(_)
438            | ThreadEvent::ItemUpdated(_)
439            | ThreadEvent::PlanDelta(_)
440            | ThreadEvent::Unknown => {}
441        }
442    }
443
444    fn process_item_completed(&mut self, item_id: &str, details: &ThreadItemDetails, ts: &str) {
445        match details {
446            ThreadItemDetails::AgentMessage(msg) => {
447                let mut step = Step::agent(self.next_step_id, &msg.text);
448                step.timestamp = Some(ts.to_string());
449                self.push_step(step);
450            }
451            ThreadItemDetails::Plan(plan) => {
452                let mut step = Step::agent(self.next_step_id, &plan.text);
453                step.timestamp = Some(ts.to_string());
454                step.extra = Some(serde_json::json!({ "vtcode_item_type": "plan" }));
455                self.push_step(step);
456            }
457            ThreadItemDetails::Reasoning(r) => {
458                let mut step = Step::agent(self.next_step_id, "");
459                step.timestamp = Some(ts.to_string());
460                step.reasoning_content = Some(r.text.clone());
461                step.message = None;
462                self.push_step(step);
463            }
464            ThreadItemDetails::ToolInvocation(inv) => {
465                // Buffer the invocation; we'll pair it with the ToolOutput
466                self.pending_tool_calls.push(PendingToolCall {
467                    call_id: item_id.to_string(),
468                    tool_call_id: inv.tool_call_id.clone(),
469                    tool_name: inv.tool_name.clone(),
470                    arguments: inv.arguments.clone(),
471                    timestamp: ts.to_string(),
472                });
473            }
474            ThreadItemDetails::ToolOutput(output) => {
475                // Find the matching pending invocation
476                let pending_idx = self
477                    .pending_tool_calls
478                    .iter()
479                    .position(|p| p.call_id == output.call_id);
480
481                let (tool_name, arguments, tool_call_id, inv_ts) = if let Some(idx) = pending_idx {
482                    let p = self.pending_tool_calls.remove(idx);
483                    (p.tool_name, p.arguments, p.tool_call_id, p.timestamp)
484                } else {
485                    (
486                        "unknown".to_string(),
487                        None,
488                        output.tool_call_id.clone(),
489                        ts.to_string(),
490                    )
491                };
492
493                let call_id = tool_call_id
494                    .clone()
495                    .unwrap_or_else(|| output.call_id.clone());
496
497                let mut step = Step::agent(self.next_step_id, "");
498                step.timestamp = Some(inv_ts);
499                step.message = None;
500                step.tool_calls = Some(vec![AtifToolCall {
501                    tool_call_id: call_id.clone(),
502                    function_name: tool_name,
503                    arguments,
504                }]);
505
506                let status_suffix = match output.status {
507                    ToolCallStatus::Failed => " [FAILED]",
508                    ToolCallStatus::InProgress => " [IN_PROGRESS]",
509                    ToolCallStatus::Completed => "",
510                };
511                let content = format!("{}{}", output.output, status_suffix);
512                step.observation = Some(Observation {
513                    results: vec![ObservationResult {
514                        source_call_id: call_id,
515                        content,
516                    }],
517                });
518                self.push_step(step);
519            }
520            ThreadItemDetails::CommandExecution(cmd) => {
521                let call_id = item_id.to_string();
522                let mut step = Step::agent(self.next_step_id, "");
523                step.timestamp = Some(ts.to_string());
524                step.message = None;
525                step.tool_calls = Some(vec![AtifToolCall {
526                    tool_call_id: call_id.clone(),
527                    function_name: "command_execution".to_string(),
528                    arguments: Some(serde_json::json!({
529                        "command": cmd.command,
530                        "arguments": cmd.arguments,
531                    })),
532                }]);
533                step.observation = Some(Observation {
534                    results: vec![ObservationResult {
535                        source_call_id: call_id,
536                        content: cmd.aggregated_output.clone(),
537                    }],
538                });
539                if let Some(exit_code) = cmd.exit_code {
540                    step.extra = Some(serde_json::json!({ "exit_code": exit_code }));
541                }
542                self.push_step(step);
543            }
544            ThreadItemDetails::McpToolCall(mcp) => {
545                let call_id = item_id.to_string();
546                let mut step = Step::agent(self.next_step_id, "");
547                step.timestamp = Some(ts.to_string());
548                step.message = None;
549                step.tool_calls = Some(vec![AtifToolCall {
550                    tool_call_id: call_id.clone(),
551                    function_name: mcp.tool_name.clone(),
552                    arguments: mcp.arguments.clone(),
553                }]);
554                if let Some(result) = &mcp.result {
555                    step.observation = Some(Observation {
556                        results: vec![ObservationResult {
557                            source_call_id: call_id,
558                            content: result.clone(),
559                        }],
560                    });
561                }
562                self.push_step(step);
563            }
564            ThreadItemDetails::FileChange(fc) => {
565                let changes: Vec<String> = fc
566                    .changes
567                    .iter()
568                    .map(|c| format!("{}: {:?}", c.path, c.kind))
569                    .collect();
570                let msg = format!("file_changes: {}", changes.join(", "));
571                let mut step = Step::system(self.next_step_id, msg);
572                step.timestamp = Some(ts.to_string());
573                self.push_step(step);
574            }
575            ThreadItemDetails::WebSearch(ws) => {
576                let mut step = Step::system(self.next_step_id, format!("web_search: {}", ws.query));
577                step.timestamp = Some(ts.to_string());
578                if let Some(results) = &ws.results {
579                    step.observation = Some(Observation {
580                        results: results
581                            .iter()
582                            .enumerate()
583                            .map(|(i, r)| ObservationResult {
584                                source_call_id: format!("search_{i}"),
585                                content: r.clone(),
586                            })
587                            .collect(),
588                    });
589                }
590                self.push_step(step);
591            }
592            ThreadItemDetails::Harness(h) => {
593                let msg = format!("harness: {:?}", h.event);
594                let mut step = Step::system(self.next_step_id, msg);
595                step.timestamp = Some(ts.to_string());
596                if let Some(m) = &h.message {
597                    step.extra = Some(serde_json::json!({ "harness_message": m }));
598                }
599                self.push_step(step);
600            }
601            ThreadItemDetails::Error(e) => {
602                let mut step = Step::system(self.next_step_id, &e.message);
603                step.timestamp = Some(ts.to_string());
604                self.push_step(step);
605            }
606        }
607    }
608
609    fn push_step(&mut self, step: Step) {
610        self.next_step_id = step.step_id + 1;
611        self.steps.push(step);
612    }
613
614    /// Consume the builder and produce the final ATIF trajectory.
615    ///
616    /// Pass optional `FinalMetrics` to override the accumulated values.
617    /// If `None`, final metrics are derived from observed events.
618    pub fn finish(self, override_metrics: Option<FinalMetrics>) -> Trajectory {
619        let final_metrics = override_metrics.unwrap_or_else(|| FinalMetrics {
620            total_prompt_tokens: Some(self.total_input_tokens),
621            total_completion_tokens: Some(self.total_output_tokens),
622            total_cached_tokens: if self.total_cached_tokens > 0 {
623                Some(self.total_cached_tokens)
624            } else {
625                None
626            },
627            total_cost_usd: None,
628            total_steps: Some(self.steps.len() as u64),
629            extra: Some(serde_json::json!({ "num_turns": self.num_turns })),
630        });
631
632        Trajectory {
633            schema_version: ATIF_SCHEMA_VERSION.to_string(),
634            session_id: self
635                .session_id
636                .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
637            agent: self.agent,
638            steps: self.steps,
639            notes: None,
640            final_metrics: Some(final_metrics),
641            extra: None,
642        }
643    }
644
645    /// Returns the number of steps collected so far.
646    pub fn step_count(&self) -> usize {
647        self.steps.len()
648    }
649}
650
651impl crate::EventEmitter for AtifTrajectoryBuilder {
652    fn emit(&mut self, event: &ThreadEvent) {
653        self.process_event(event);
654    }
655}
656
657#[cfg(test)]
658mod tests {
659    use super::*;
660    use crate::{
661        AgentMessageItem, ItemCompletedEvent, ThreadItem, ThreadStartedEvent, ToolInvocationItem,
662        ToolOutputItem, TurnCompletedEvent, Usage,
663    };
664
665    fn fixed_ts() -> DateTime<Utc> {
666        "2025-01-15T10:30:00Z".parse().unwrap()
667    }
668
669    #[test]
670    fn trajectory_round_trip() {
671        let trajectory = Trajectory {
672            schema_version: ATIF_SCHEMA_VERSION.to_string(),
673            session_id: "test-session".to_string(),
674            agent: AtifAgent::vtcode(),
675            steps: vec![Step::user(1, "hello")],
676            notes: None,
677            final_metrics: None,
678            extra: None,
679        };
680
681        let json = serde_json::to_string_pretty(&trajectory).unwrap();
682        let restored: Trajectory = serde_json::from_str(&json).unwrap();
683        assert_eq!(restored.schema_version, ATIF_SCHEMA_VERSION);
684        assert_eq!(restored.session_id, "test-session");
685        assert_eq!(restored.steps.len(), 1);
686    }
687
688    #[test]
689    fn builder_thread_started_sets_session_id() {
690        let mut builder = AtifTrajectoryBuilder::new(AtifAgent::vtcode());
691        let event = ThreadEvent::ThreadStarted(ThreadStartedEvent {
692            thread_id: "thread-abc".to_string(),
693        });
694        builder.process_event_at(&event, fixed_ts());
695        let trajectory = builder.finish(None);
696        assert_eq!(trajectory.session_id, "thread-abc");
697    }
698
699    #[test]
700    fn builder_agent_message_step() {
701        let mut builder = AtifTrajectoryBuilder::new(AtifAgent::vtcode());
702        let event = ThreadEvent::ItemCompleted(ItemCompletedEvent {
703            item: ThreadItem {
704                id: "msg-1".to_string(),
705                details: ThreadItemDetails::AgentMessage(AgentMessageItem {
706                    text: "Hello, world!".to_string(),
707                }),
708            },
709        });
710        builder.process_event_at(&event, fixed_ts());
711        let trajectory = builder.finish(None);
712
713        assert_eq!(trajectory.steps.len(), 1);
714        let step = &trajectory.steps[0];
715        assert_eq!(step.step_id, 1);
716        assert_eq!(step.source, StepSource::Agent);
717        assert_eq!(step.message.as_deref(), Some("Hello, world!"));
718    }
719
720    #[test]
721    fn builder_tool_invocation_with_output() {
722        let mut builder = AtifTrajectoryBuilder::new(AtifAgent::vtcode());
723        let ts = fixed_ts();
724
725        // Tool invocation
726        let inv_event = ThreadEvent::ItemCompleted(ItemCompletedEvent {
727            item: ThreadItem {
728                id: "tool_1".to_string(),
729                details: ThreadItemDetails::ToolInvocation(ToolInvocationItem {
730                    tool_name: "read_file".to_string(),
731                    arguments: Some(serde_json::json!({"path": "README.md"})),
732                    tool_call_id: Some("tc_0".to_string()),
733                    status: ToolCallStatus::Completed,
734                }),
735            },
736        });
737        builder.process_event_at(&inv_event, ts);
738
739        // Tool output
740        let out_event = ThreadEvent::ItemCompleted(ItemCompletedEvent {
741            item: ThreadItem {
742                id: "tool_1:output".to_string(),
743                details: ThreadItemDetails::ToolOutput(ToolOutputItem {
744                    call_id: "tool_1".to_string(),
745                    tool_call_id: Some("tc_0".to_string()),
746                    spool_path: None,
747                    output: "file contents here".to_string(),
748                    exit_code: Some(0),
749                    status: ToolCallStatus::Completed,
750                }),
751            },
752        });
753        builder.process_event_at(&out_event, ts);
754
755        let trajectory = builder.finish(None);
756        // Only one step: the invocation is buffered until output arrives
757        assert_eq!(trajectory.steps.len(), 1);
758        let step = &trajectory.steps[0];
759        assert_eq!(step.source, StepSource::Agent);
760
761        let calls = step.tool_calls.as_ref().unwrap();
762        assert_eq!(calls.len(), 1);
763        assert_eq!(calls[0].function_name, "read_file");
764        assert_eq!(calls[0].tool_call_id, "tc_0");
765
766        let obs = step.observation.as_ref().unwrap();
767        assert_eq!(obs.results.len(), 1);
768        assert_eq!(obs.results[0].content, "file contents here");
769    }
770
771    #[test]
772    fn builder_turn_completed_accumulates_metrics() {
773        let mut builder = AtifTrajectoryBuilder::new(AtifAgent::vtcode());
774        let event = ThreadEvent::TurnCompleted(TurnCompletedEvent {
775            usage: Usage {
776                input_tokens: 500,
777                cached_input_tokens: 100,
778                cache_creation_tokens: 0,
779                output_tokens: 200,
780            },
781        });
782        builder.process_event_at(&event, fixed_ts());
783
784        let trajectory = builder.finish(None);
785        let fm = trajectory.final_metrics.as_ref().unwrap();
786        assert_eq!(fm.total_prompt_tokens, Some(500));
787        assert_eq!(fm.total_completion_tokens, Some(200));
788        assert_eq!(fm.total_cached_tokens, Some(100));
789    }
790
791    #[test]
792    fn step_metrics_from_usage() {
793        let usage = Usage {
794            input_tokens: 1000,
795            cached_input_tokens: 200,
796            cache_creation_tokens: 50,
797            output_tokens: 300,
798        };
799        let metrics = StepMetrics::from_usage(&usage);
800        assert_eq!(metrics.prompt_tokens, Some(1000));
801        assert_eq!(metrics.completion_tokens, Some(300));
802        assert_eq!(metrics.cached_tokens, Some(200));
803        assert!(metrics.extra.is_some());
804    }
805
806    #[test]
807    fn builder_implements_event_emitter() {
808        let mut builder = AtifTrajectoryBuilder::new(AtifAgent::vtcode());
809        let event = ThreadEvent::ThreadStarted(ThreadStartedEvent {
810            thread_id: "t-1".to_string(),
811        });
812        // Use EventEmitter trait
813        crate::EventEmitter::emit(&mut builder, &event);
814        assert_eq!(builder.step_count(), 0); // ThreadStarted doesn't create a step
815    }
816
817    #[test]
818    fn skips_lifecycle_events() {
819        let mut builder = AtifTrajectoryBuilder::new(AtifAgent::vtcode());
820        builder.process_event(&ThreadEvent::TurnStarted(crate::TurnStartedEvent {}));
821        assert_eq!(builder.step_count(), 0);
822    }
823}
vtcode_exec_events/atif.rs

vtcode_exec_events/
atif.rs
