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

vtcode_exec_events/
atif.rs
