fleetforge_telemetry/
context.rs

1//! Shared tracing contract helpers for FleetForge runtime instrumentation.
2
3use std::collections::BTreeMap;
4use std::fmt;
5use std::str::FromStr;
6
7use crate::telemetry_core::propagation;
8use anyhow::{anyhow, Result};
9use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState};
10use opentelemetry::{Context as OtelContext, KeyValue};
11use serde_json::{Map, Number, Value};
12use uuid::Uuid;
13
14/// Attribute prefix used for FleetForge-specific telemetry dimensions.
15pub const ATTR_PREFIX: &str = "fleetforge";
16
17/// Trace context tying together run, step, tool, and policy scopes.
18#[derive(Debug, Clone)]
19pub struct TraceContext {
20    trace_id: TraceId,
21    span_id: SpanId,
22    parent_span_id: Option<SpanId>,
23    trace_flags: TraceFlags,
24    tracestate: Option<String>,
25    pub run: RunScope,
26    pub step: Option<StepScope>,
27    pub tool_call: Option<ToolCallScope>,
28    pub policy_decision: Option<PolicyDecisionScope>,
29    pub budget: Option<BudgetScope>,
30    pub seed: Option<i64>,
31    pub baggage: Vec<(String, String)>,
32}
33
34impl TraceContext {
35    /// Create a new root trace context for a run.
36    pub fn new(run: RunScope) -> Self {
37        Self {
38            trace_id: new_trace_id(),
39            span_id: new_span_id(),
40            parent_span_id: None,
41            trace_flags: TraceFlags::SAMPLED,
42            tracestate: None,
43            run,
44            step: None,
45            tool_call: None,
46            policy_decision: None,
47            budget: None,
48            seed: None,
49            baggage: Vec::new(),
50        }
51    }
52
53    /// Override the generated trace identifier (used for deterministic trace IDs).
54    pub fn with_trace_id(mut self, trace_id: TraceId) -> Self {
55        self.trace_id = trace_id;
56        self
57    }
58
59    /// Override the generated span identifier.
60    pub fn with_span_id(mut self, span_id: SpanId) -> Self {
61        self.span_id = span_id;
62        self
63    }
64
65    /// Set an explicit parent span identifier.
66    pub fn with_parent_span_id(mut self, parent_span_id: impl Into<Option<SpanId>>) -> Self {
67        self.parent_span_id = parent_span_id.into();
68        self
69    }
70
71    /// Create a trace context from an existing `trace_id`/`span_id` pair, typically
72    /// when continuing a trace injected by an upstream caller.
73    pub fn from_parent(run: RunScope, trace_id: TraceId, parent_span_id: SpanId) -> Self {
74        Self {
75            trace_id,
76            span_id: new_span_id(),
77            parent_span_id: Some(parent_span_id),
78            trace_flags: TraceFlags::SAMPLED,
79            tracestate: None,
80            run,
81            step: None,
82            tool_call: None,
83            policy_decision: None,
84            budget: None,
85            seed: None,
86            baggage: Vec::new(),
87        }
88    }
89
90    /// Attach a deterministic seed to the trace.
91    pub fn with_seed(mut self, seed: impl Into<Option<i64>>) -> Self {
92        self.seed = seed.into();
93        self
94    }
95
96    /// Override the trace flags (e.g., when sampling is disabled).
97    pub fn with_trace_flags(mut self, flags: TraceFlags) -> Self {
98        self.trace_flags = flags;
99        self
100    }
101
102    /// Set an explicit W3C tracestate header value.
103    pub fn with_tracestate(mut self, tracestate: impl Into<Option<String>>) -> Self {
104        self.tracestate = tracestate.into();
105        self
106    }
107
108    /// Construct a trace context from W3C headers captured on ingress.
109    pub fn from_w3c(
110        run: RunScope,
111        traceparent: &str,
112        tracestate: Option<&str>,
113        baggage: Option<&str>,
114    ) -> Result<Self> {
115        let (trace_id, parent_span_id, trace_flags) = parse_traceparent(traceparent)?;
116        let mut ctx = Self::new(run)
117            .with_trace_id(trace_id)
118            .with_parent_span_id(parent_span_id)
119            .with_trace_flags(trace_flags);
120
121        if let Some(state) = tracestate {
122            let trimmed = state.trim();
123            if !trimmed.is_empty() {
124                ctx.tracestate = Some(trimmed.to_string());
125            }
126        }
127
128        if let Some(raw) = baggage {
129            ctx.baggage = parse_baggage(raw)?;
130        }
131
132        Ok(ctx)
133    }
134
135    /// Replace baggage entries (for example to propagate budgets or experiment IDs).
136    pub fn with_baggage<I, K, V>(mut self, baggage: I) -> Self
137    where
138        I: IntoIterator<Item = (K, V)>,
139        K: Into<String>,
140        V: Into<String>,
141    {
142        self.baggage = baggage
143            .into_iter()
144            .map(|(k, v)| (k.into(), v.into()))
145            .collect();
146        self
147    }
148
149    /// Attach budget scope metadata.
150    pub fn with_budget(mut self, budget: BudgetScope) -> Self {
151        self.budget = Some(budget);
152        self
153    }
154
155    /// Generate a child span context representing a step.
156    pub fn child_step(&self, scope: StepScope) -> Self {
157        let mut derived = self.derived();
158        derived.step = Some(scope);
159        derived.tool_call = None;
160        derived.policy_decision = None;
161        derived
162    }
163
164    /// Generate a child span context representing a tool call.
165    pub fn child_tool_call(&self, scope: ToolCallScope) -> Self {
166        let mut derived = self.derived();
167        derived.step = self.step.clone();
168        derived.tool_call = Some(scope);
169        derived.policy_decision = None;
170        derived
171    }
172
173    /// Generate a child span context representing a policy decision.
174    pub fn child_policy_decision(&self, scope: PolicyDecisionScope) -> Self {
175        let mut derived = self.derived();
176        derived.step = self.step.clone();
177        derived.tool_call = self.tool_call.clone();
178        derived.policy_decision = Some(scope);
179        derived
180    }
181
182    /// Trace identifier as a `TraceId`.
183    pub fn trace_id(&self) -> TraceId {
184        self.trace_id
185    }
186
187    /// Current span identifier.
188    pub fn span_id(&self) -> SpanId {
189        self.span_id
190    }
191
192    /// Trace flags used when constructing spans.
193    pub fn trace_flags(&self) -> TraceFlags {
194        self.trace_flags
195    }
196
197    /// Parent span identifier, if any.
198    pub fn parent_span_id(&self) -> Option<SpanId> {
199        self.parent_span_id
200    }
201
202    /// Compute the W3C `traceparent` header value.
203    pub fn traceparent(&self) -> String {
204        propagation::traceparent(self.trace_id, self.span_id, self.trace_flags)
205    }
206
207    /// Return the stored W3C `tracestate` header value, if any.
208    pub fn tracestate(&self) -> Option<String> {
209        self.tracestate.clone()
210    }
211
212    /// Return the formatted W3C `baggage` header, if any entries are present.
213    pub fn baggage_header(&self) -> Option<String> {
214        propagation::baggage_header(&self.baggage)
215    }
216
217    /// Produce the standard set of W3C headers (traceparent, tracestate, baggage).
218    pub fn w3c_headers(&self) -> Vec<(String, String)> {
219        propagation::w3c_headers(
220            self.trace_id,
221            self.span_id,
222            self.trace_flags,
223            self.tracestate.as_deref(),
224            &self.baggage,
225        )
226    }
227
228    /// Build an OpenTelemetry context representing this trace (as a remote parent).
229    pub fn as_parent_context(&self) -> OtelContext {
230        let trace_state = self
231            .tracestate
232            .as_ref()
233            .and_then(|state| TraceState::from_str(state).ok())
234            .unwrap_or_else(TraceState::default);
235
236        let span_context = SpanContext::new(
237            self.trace_id,
238            self.span_id,
239            self.trace_flags,
240            true,
241            trace_state,
242        );
243
244        OtelContext::current().with_remote_span_context(span_context)
245    }
246
247    /// Construct a trace context from a recorded span (e.g., after creating a tracing span).
248    pub fn from_span_context(
249        run: RunScope,
250        step: Option<StepScope>,
251        tool_call: Option<ToolCallScope>,
252        policy_decision: Option<PolicyDecisionScope>,
253        budget: Option<BudgetScope>,
254        parent_span_id: Option<SpanId>,
255        span_context: SpanContext,
256        seed: Option<i64>,
257        baggage: Vec<(String, String)>,
258    ) -> Self {
259        let tracestate = span_context.trace_state().header();
260        Self {
261            trace_id: span_context.trace_id(),
262            span_id: span_context.span_id(),
263            parent_span_id,
264            trace_flags: span_context.trace_flags(),
265            tracestate: if tracestate.is_empty() {
266                None
267            } else {
268                Some(tracestate)
269            },
270            run,
271            step,
272            tool_call,
273            policy_decision,
274            budget,
275            seed,
276            baggage,
277        }
278    }
279
280    /// Convert the trace context into OpenTelemetry attributes using the
281    /// `fleetforge.*` namespace.
282    pub fn attributes(&self) -> Vec<KeyValue> {
283        let mut attrs = Vec::with_capacity(24);
284
285        attrs.push(KeyValue::new(
286            format!("{ATTR_PREFIX}.trace.trace_id"),
287            self.trace_id.to_string(),
288        ));
289        attrs.push(KeyValue::new(
290            format!("{ATTR_PREFIX}.trace.span_id"),
291            self.span_id.to_string(),
292        ));
293        if let Some(parent) = self.parent_span_id {
294            attrs.push(KeyValue::new(
295                format!("{ATTR_PREFIX}.trace.parent_span_id"),
296                parent.to_string(),
297            ));
298        }
299
300        attrs.push(KeyValue::new(
301            format!("{ATTR_PREFIX}.run.id"),
302            self.run.run_id.clone(),
303        ));
304        if let Some(workspace) = &self.run.workspace_id {
305            attrs.push(KeyValue::new(
306                format!("{ATTR_PREFIX}.workspace.id"),
307                workspace.clone(),
308            ));
309        }
310        if let Some(app_id) = &self.run.app_id {
311            attrs.push(KeyValue::new(
312                format!("{ATTR_PREFIX}.app.id"),
313                app_id.clone(),
314            ));
315        }
316        if let Some(attempt_id) = &self.run.attempt_id {
317            attrs.push(KeyValue::new(
318                format!("{ATTR_PREFIX}.run.attempt_id"),
319                attempt_id.clone(),
320            ));
321        }
322        if let Some(seed) = self.seed {
323            attrs.push(KeyValue::new(
324                format!("{ATTR_PREFIX}.run.seed"),
325                seed as i64,
326            ));
327        }
328        for (key, value) in &self.run.labels {
329            attrs.push(KeyValue::new(
330                format!("{ATTR_PREFIX}.run.label.{key}"),
331                value.clone(),
332            ));
333        }
334
335        if let Some(step) = &self.step {
336            attrs.push(KeyValue::new(
337                format!("{ATTR_PREFIX}.step.id"),
338                step.step_id.clone(),
339            ));
340            if let Some(index) = step.index {
341                attrs.push(KeyValue::new(
342                    format!("{ATTR_PREFIX}.step.index"),
343                    index as i64,
344                ));
345            }
346            if let Some(attempt) = step.attempt {
347                attrs.push(KeyValue::new(
348                    format!("{ATTR_PREFIX}.step.attempt"),
349                    attempt as i64,
350                ));
351            }
352            if let Some(kind) = &step.kind {
353                attrs.push(KeyValue::new(
354                    format!("{ATTR_PREFIX}.step.kind"),
355                    kind.clone(),
356                ));
357            }
358            if let Some(scheduler) = &step.scheduler {
359                attrs.push(KeyValue::new(
360                    format!("{ATTR_PREFIX}.step.scheduler"),
361                    scheduler.clone(),
362                ));
363            }
364        }
365
366        if let Some(tool) = &self.tool_call {
367            attrs.push(KeyValue::new(
368                format!("{ATTR_PREFIX}.tool.id"),
369                tool.tool_call_id.clone(),
370            ));
371            if let Some(name) = &tool.tool_name {
372                attrs.push(KeyValue::new(
373                    format!("{ATTR_PREFIX}.tool.name"),
374                    name.clone(),
375                ));
376            }
377            if let Some(variant) = &tool.tool_variant {
378                attrs.push(KeyValue::new(
379                    format!("{ATTR_PREFIX}.tool.variant"),
380                    variant.clone(),
381                ));
382            }
383            if let Some(provider) = &tool.provider {
384                attrs.push(KeyValue::new(
385                    format!("{ATTR_PREFIX}.tool.provider"),
386                    provider.clone(),
387                ));
388            }
389        }
390
391        if let Some(policy) = &self.policy_decision {
392            attrs.push(KeyValue::new(
393                format!("{ATTR_PREFIX}.policy.decision_id"),
394                policy.policy_decision_id.clone(),
395            ));
396            if let Some(pack_id) = &policy.policy_pack_id {
397                attrs.push(KeyValue::new(
398                    format!("{ATTR_PREFIX}.policy.pack_id"),
399                    pack_id.clone(),
400                ));
401            }
402            if let Some(name) = &policy.policy_name {
403                attrs.push(KeyValue::new(
404                    format!("{ATTR_PREFIX}.policy.name"),
405                    name.clone(),
406                ));
407            }
408            if let Some(version) = &policy.policy_version {
409                attrs.push(KeyValue::new(
410                    format!("{ATTR_PREFIX}.policy.version"),
411                    version.clone(),
412                ));
413            }
414        }
415
416        if let Some(budget) = &self.budget {
417            attrs.push(KeyValue::new(
418                format!("{ATTR_PREFIX}.budget.id"),
419                budget.budget_id.clone(),
420            ));
421            if let Some(kind) = &budget.budget_kind {
422                attrs.push(KeyValue::new(
423                    format!("{ATTR_PREFIX}.budget.kind"),
424                    kind.clone(),
425                ));
426            }
427            if let Some(remaining) = budget.amount_remaining {
428                attrs.push(KeyValue::new(
429                    format!("{ATTR_PREFIX}.budget.remaining"),
430                    remaining,
431                ));
432            }
433            if let Some(allocated) = budget.amount_allocated {
434                attrs.push(KeyValue::new(
435                    format!("{ATTR_PREFIX}.budget.allocated"),
436                    allocated,
437                ));
438            }
439        }
440
441        attrs
442    }
443
444    /// Serialises the trace context into a JSON object suitable for inclusion in event payloads.
445    pub fn to_json(&self) -> Value {
446        let mut root = Map::new();
447        root.insert(
448            "trace_id".to_string(),
449            Value::String(self.trace_id.to_string()),
450        );
451        root.insert(
452            "span_id".to_string(),
453            Value::String(self.span_id.to_string()),
454        );
455        if let Some(parent) = self.parent_span_id {
456            root.insert(
457                "parent_span_id".to_string(),
458                Value::String(parent.to_string()),
459            );
460        }
461        if !self.tracestate.as_deref().unwrap_or_default().is_empty() {
462            if let Some(state) = &self.tracestate {
463                root.insert("tracestate".to_string(), Value::String(state.clone()));
464            }
465        }
466        root.insert(
467            "trace_flags".to_string(),
468            Value::Number(Number::from(u64::from(self.trace_flags.to_u8()))),
469        );
470        if let Some(seed) = self.seed {
471            root.insert("seed".to_string(), Value::Number(Number::from(seed)));
472        }
473
474        let mut run_map = Map::new();
475        run_map.insert("run_id".to_string(), Value::String(self.run.run_id.clone()));
476        if let Some(workspace) = &self.run.workspace_id {
477            run_map.insert("workspace_id".to_string(), Value::String(workspace.clone()));
478        }
479        if let Some(app_id) = &self.run.app_id {
480            run_map.insert("app_id".to_string(), Value::String(app_id.clone()));
481        }
482        if let Some(attempt) = &self.run.attempt_id {
483            run_map.insert("attempt_id".to_string(), Value::String(attempt.clone()));
484        }
485        if !self.run.labels.is_empty() {
486            let labels = self
487                .run
488                .labels
489                .iter()
490                .map(|(k, v)| (k.clone(), Value::String(v.clone())))
491                .collect::<Map<_, _>>();
492            run_map.insert("labels".to_string(), Value::Object(labels));
493        }
494        root.insert("run".to_string(), Value::Object(run_map));
495
496        if let Some(step) = &self.step {
497            let mut step_map = Map::new();
498            step_map.insert("step_id".to_string(), Value::String(step.step_id.clone()));
499            if let Some(index) = step.index {
500                step_map.insert("index".to_string(), Value::Number(Number::from(index)));
501            }
502            if let Some(attempt) = step.attempt {
503                step_map.insert("attempt".to_string(), Value::Number(Number::from(attempt)));
504            }
505            if let Some(kind) = &step.kind {
506                step_map.insert("kind".to_string(), Value::String(kind.clone()));
507            }
508            if let Some(scheduler) = &step.scheduler {
509                step_map.insert("scheduler".to_string(), Value::String(scheduler.clone()));
510            }
511            root.insert("step".to_string(), Value::Object(step_map));
512        }
513
514        if let Some(tool) = &self.tool_call {
515            let mut tool_map = Map::new();
516            tool_map.insert(
517                "tool_call_id".to_string(),
518                Value::String(tool.tool_call_id.clone()),
519            );
520            if let Some(name) = &tool.tool_name {
521                tool_map.insert("tool_name".to_string(), Value::String(name.clone()));
522            }
523            if let Some(variant) = &tool.tool_variant {
524                tool_map.insert("tool_variant".to_string(), Value::String(variant.clone()));
525            }
526            if let Some(provider) = &tool.provider {
527                tool_map.insert("provider".to_string(), Value::String(provider.clone()));
528            }
529            root.insert("tool_call".to_string(), Value::Object(tool_map));
530        }
531
532        if let Some(policy) = &self.policy_decision {
533            let mut policy_map = Map::new();
534            policy_map.insert(
535                "policy_decision_id".to_string(),
536                Value::String(policy.policy_decision_id.clone()),
537            );
538            if let Some(pack) = &policy.policy_pack_id {
539                policy_map.insert("policy_pack_id".to_string(), Value::String(pack.clone()));
540            }
541            if let Some(name) = &policy.policy_name {
542                policy_map.insert("policy_name".to_string(), Value::String(name.clone()));
543            }
544            if let Some(version) = &policy.policy_version {
545                policy_map.insert("policy_version".to_string(), Value::String(version.clone()));
546            }
547            root.insert("policy_decision".to_string(), Value::Object(policy_map));
548        }
549
550        if let Some(budget) = &self.budget {
551            let mut budget_map = Map::new();
552            budget_map.insert(
553                "budget_id".to_string(),
554                Value::String(budget.budget_id.clone()),
555            );
556            if let Some(kind) = &budget.budget_kind {
557                budget_map.insert("budget_kind".to_string(), Value::String(kind.clone()));
558            }
559            if let Some(remaining) = budget
560                .amount_remaining
561                .and_then(|value| Number::from_f64(value))
562            {
563                budget_map.insert("amount_remaining".to_string(), Value::Number(remaining));
564            }
565            if let Some(allocated) = budget
566                .amount_allocated
567                .and_then(|value| Number::from_f64(value))
568            {
569                budget_map.insert("amount_allocated".to_string(), Value::Number(allocated));
570            }
571            root.insert("budget".to_string(), Value::Object(budget_map));
572        }
573
574        if !self.baggage.is_empty() {
575            let baggage = self
576                .baggage
577                .iter()
578                .map(|(key, value)| {
579                    let mut map = Map::new();
580                    map.insert("key".to_string(), Value::String(key.clone()));
581                    map.insert("value".to_string(), Value::String(value.clone()));
582                    Value::Object(map)
583                })
584                .collect();
585            root.insert("baggage".to_string(), Value::Array(baggage));
586        }
587
588        Value::Object(root)
589    }
590
591    fn derived(&self) -> Self {
592        Self {
593            trace_id: self.trace_id,
594            span_id: new_span_id(),
595            parent_span_id: Some(self.span_id),
596            trace_flags: self.trace_flags,
597            tracestate: self.tracestate.clone(),
598            run: self.run.clone(),
599            step: None,
600            tool_call: None,
601            policy_decision: None,
602            budget: self.budget.clone(),
603            seed: self.seed,
604            baggage: self.baggage.clone(),
605        }
606    }
607}
608
609/// Run-level identifiers.
610#[derive(Debug, Clone)]
611pub struct RunScope {
612    pub run_id: String,
613    pub workspace_id: Option<String>,
614    pub app_id: Option<String>,
615    pub attempt_id: Option<String>,
616    pub labels: BTreeMap<String, String>,
617}
618
619impl RunScope {
620    pub fn new(run_id: impl Into<String>) -> Self {
621        Self {
622            run_id: run_id.into(),
623            workspace_id: None,
624            app_id: None,
625            attempt_id: None,
626            labels: BTreeMap::new(),
627        }
628    }
629
630    pub fn with_workspace(mut self, workspace_id: impl Into<String>) -> Self {
631        self.workspace_id = Some(workspace_id.into());
632        self
633    }
634
635    pub fn with_app(mut self, app_id: impl Into<String>) -> Self {
636        self.app_id = Some(app_id.into());
637        self
638    }
639
640    pub fn with_attempt(mut self, attempt_id: impl Into<String>) -> Self {
641        self.attempt_id = Some(attempt_id.into());
642        self
643    }
644
645    pub fn with_labels<I, K, V>(mut self, labels: I) -> Self
646    where
647        I: IntoIterator<Item = (K, V)>,
648        K: Into<String>,
649        V: Into<String>,
650    {
651        for (key, value) in labels.into_iter() {
652            self.labels.insert(key.into(), value.into());
653        }
654        self
655    }
656}
657
658/// Step-level identifiers and metadata.
659#[derive(Debug, Clone)]
660pub struct StepScope {
661    pub step_id: String,
662    pub index: Option<u32>,
663    pub attempt: Option<u32>,
664    pub kind: Option<String>,
665    pub scheduler: Option<String>,
666}
667
668impl StepScope {
669    pub fn new(step_id: impl Into<String>) -> Self {
670        Self {
671            step_id: step_id.into(),
672            index: None,
673            attempt: None,
674            kind: None,
675            scheduler: None,
676        }
677    }
678
679    pub fn with_index(mut self, index: u32) -> Self {
680        self.index = Some(index);
681        self
682    }
683
684    pub fn with_attempt(mut self, attempt: u32) -> Self {
685        self.attempt = Some(attempt);
686        self
687    }
688
689    pub fn with_kind(mut self, kind: impl Into<String>) -> Self {
690        self.kind = Some(kind.into());
691        self
692    }
693
694    pub fn with_scheduler(mut self, scheduler: impl Into<String>) -> Self {
695        self.scheduler = Some(scheduler.into());
696        self
697    }
698}
699
700/// Tool call identifiers.
701#[derive(Debug, Clone)]
702pub struct ToolCallScope {
703    pub tool_call_id: String,
704    pub tool_name: Option<String>,
705    pub tool_variant: Option<String>,
706    pub provider: Option<String>,
707}
708
709impl ToolCallScope {
710    pub fn new(tool_call_id: impl Into<String>) -> Self {
711        Self {
712            tool_call_id: tool_call_id.into(),
713            tool_name: None,
714            tool_variant: None,
715            provider: None,
716        }
717    }
718
719    pub fn with_name(mut self, name: impl Into<String>) -> Self {
720        self.tool_name = Some(name.into());
721        self
722    }
723
724    pub fn with_variant(mut self, variant: impl Into<String>) -> Self {
725        self.tool_variant = Some(variant.into());
726        self
727    }
728
729    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
730        self.provider = Some(provider.into());
731        self
732    }
733}
734
735/// Policy decision identifiers.
736#[derive(Debug, Clone)]
737pub struct PolicyDecisionScope {
738    pub policy_decision_id: String,
739    pub policy_pack_id: Option<String>,
740    pub policy_name: Option<String>,
741    pub policy_version: Option<String>,
742}
743
744impl PolicyDecisionScope {
745    pub fn new(policy_decision_id: impl Into<String>) -> Self {
746        Self {
747            policy_decision_id: policy_decision_id.into(),
748            policy_pack_id: None,
749            policy_name: None,
750            policy_version: None,
751        }
752    }
753
754    pub fn with_pack(mut self, pack_id: impl Into<String>) -> Self {
755        self.policy_pack_id = Some(pack_id.into());
756        self
757    }
758
759    pub fn with_name(mut self, name: impl Into<String>) -> Self {
760        self.policy_name = Some(name.into());
761        self
762    }
763
764    pub fn with_version(mut self, version: impl Into<String>) -> Self {
765        self.policy_version = Some(version.into());
766        self
767    }
768}
769
770/// Budget identifiers for policy/billing scopes.
771#[derive(Debug, Clone)]
772pub struct BudgetScope {
773    pub budget_id: String,
774    pub budget_kind: Option<String>,
775    pub amount_remaining: Option<f64>,
776    pub amount_allocated: Option<f64>,
777}
778
779impl BudgetScope {
780    pub fn new(budget_id: impl Into<String>) -> Self {
781        Self {
782            budget_id: budget_id.into(),
783            budget_kind: None,
784            amount_remaining: None,
785            amount_allocated: None,
786        }
787    }
788
789    pub fn with_kind(mut self, kind: impl Into<String>) -> Self {
790        self.budget_kind = Some(kind.into());
791        self
792    }
793
794    pub fn with_remaining(mut self, remaining: f64) -> Self {
795        self.amount_remaining = Some(remaining);
796        self
797    }
798
799    pub fn with_allocated(mut self, allocated: f64) -> Self {
800        self.amount_allocated = Some(allocated);
801        self
802    }
803}
804
805fn new_trace_id() -> TraceId {
806    TraceId::from_bytes(*Uuid::new_v4().as_bytes())
807}
808
809fn new_span_id() -> SpanId {
810    let raw = Uuid::new_v4().as_u128();
811    let high = (raw >> 64) as u64;
812    SpanId::from_bytes(high.to_be_bytes())
813}
814
815fn parse_traceparent(header: &str) -> Result<(TraceId, Option<SpanId>, TraceFlags)> {
816    let mut parts = header.split('-');
817    let version = parts
818        .next()
819        .ok_or_else(|| anyhow!("traceparent missing version"))?;
820    if version.len() != 2 {
821        return Err(anyhow!("traceparent version must be 2 hex chars"));
822    }
823
824    let trace_id_str = parts
825        .next()
826        .ok_or_else(|| anyhow!("traceparent missing trace id"))?;
827    if trace_id_str.len() != 32 {
828        return Err(anyhow!("traceparent trace id must be 32 hex chars"));
829    }
830    let trace_id = TraceId::from_hex(trace_id_str)
831        .map_err(|err| anyhow!("invalid trace id '{}': {err}", trace_id_str))?;
832
833    let span_id_str = parts
834        .next()
835        .ok_or_else(|| anyhow!("traceparent missing parent span id"))?;
836    if span_id_str.len() != 16 {
837        return Err(anyhow!("traceparent span id must be 16 hex chars"));
838    }
839    let parent_span_id = SpanId::from_hex(span_id_str)
840        .map_err(|err| anyhow!("invalid span id '{}': {err}", span_id_str))?;
841
842    let flags_str = parts
843        .next()
844        .ok_or_else(|| anyhow!("traceparent missing trace flags"))?;
845    if flags_str.len() != 2 {
846        return Err(anyhow!("traceparent trace flags must be 2 hex chars"));
847    }
848    let flags_byte = u8::from_str_radix(flags_str, 16)
849        .map_err(|err| anyhow!("invalid trace flags '{}': {err}", flags_str))?;
850    let trace_flags = TraceFlags::new(flags_byte);
851
852    Ok((trace_id, Some(parent_span_id), trace_flags))
853}
854
855fn parse_baggage(header: &str) -> Result<Vec<(String, String)>> {
856    let mut baggage = Vec::new();
857
858    for item in header.split(',') {
859        let trimmed = item.trim();
860        if trimmed.is_empty() {
861            continue;
862        }
863        let mut parts = trimmed.splitn(2, '=');
864        let raw_key = parts
865            .next()
866            .ok_or_else(|| anyhow!("baggage member missing key"))?
867            .trim();
868        let raw_value = parts
869            .next()
870            .ok_or_else(|| anyhow!("baggage member missing value"))?;
871
872        let value_without_props = raw_value
873            .split(';')
874            .next()
875            .ok_or_else(|| anyhow!("baggage member missing value component"))?
876            .trim();
877
878        let key = decode_baggage_token(raw_key)?;
879        let value = decode_baggage_token(value_without_props)?;
880
881        baggage.push((key, value));
882    }
883
884    Ok(baggage)
885}
886
887fn decode_baggage_token(input: &str) -> Result<String> {
888    let bytes = input.as_bytes();
889    let mut output = String::with_capacity(bytes.len());
890    let mut idx = 0;
891
892    while idx < bytes.len() {
893        match bytes[idx] {
894            b'%' => {
895                if idx + 2 >= bytes.len() {
896                    return Err(anyhow!("invalid percent-encoding in baggage token"));
897                }
898                let hex = &bytes[idx + 1..idx + 3];
899                let hex_str = std::str::from_utf8(hex)
900                    .map_err(|err| anyhow!("invalid utf8 in baggage escape: {err}"))?;
901                let value = u8::from_str_radix(hex_str, 16)
902                    .map_err(|err| anyhow!("invalid baggage escape '%{}': {err}", hex_str))?;
903                output.push(value as char);
904                idx += 3;
905            }
906            b'+' => {
907                output.push(' ');
908                idx += 1;
909            }
910            byte => {
911                output.push(byte as char);
912                idx += 1;
913            }
914        }
915    }
916
917    Ok(output)
918}
919
920impl fmt::Display for TraceContext {
921    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
922        write!(
923            f,
924            "TraceContext(trace_id={}, span_id={}, run_id={})",
925            self.trace_id.to_string(),
926            self.span_id.to_string(),
927            self.run.run_id
928        )
929    }
930}
931
932#[cfg(test)]
933mod tests {
934    use super::*;
935    use opentelemetry::trace::TraceFlags;
936    use opentelemetry::Value;
937
938    #[test]
939    fn root_context_produces_traceparent() {
940        let ctx = TraceContext::new(RunScope::new("run-123"));
941        let traceparent = ctx.traceparent();
942        assert!(traceparent.starts_with("00-"));
943        assert_eq!(traceparent.len(), 55);
944        assert!(traceparent.ends_with("-01") || traceparent.ends_with("-00"));
945    }
946
947    #[test]
948    fn child_step_inherits_trace_id() {
949        let root = TraceContext::new(RunScope::new("run-123"));
950        let child = root.child_step(StepScope::new("step-1").with_index(1));
951
952        assert_eq!(root.trace_id(), child.trace_id());
953        assert_ne!(root.span_id(), child.span_id());
954        assert_eq!(child.parent_span_id(), Some(root.span_id()));
955        assert!(child.step.is_some());
956        assert!(child.tool_call.is_none());
957    }
958
959    #[test]
960    fn attributes_include_scope_metadata() {
961        let ctx = TraceContext::new(
962            RunScope::new("run-123")
963                .with_workspace("ws-1")
964                .with_app("app-9")
965                .with_labels([("env", "prod"), ("team", "mlops")]),
966        )
967        .with_seed(Some(42))
968        .child_step(
969            StepScope::new("step-1")
970                .with_index(3)
971                .with_attempt(2)
972                .with_kind("llm"),
973        )
974        .child_tool_call(
975            ToolCallScope::new("tool-1")
976                .with_name("search")
977                .with_provider("internal"),
978        );
979
980        let attrs = ctx.attributes();
981        assert!(attrs.iter().any(|kv| kv.key.as_str() == "fleetforge.run.id"
982            && matches!(kv.value, Value::String(ref v) if v == "run-123")));
983        assert!(attrs
984            .iter()
985            .any(|kv| kv.key.as_str() == "fleetforge.step.kind"
986                && matches!(kv.value, Value::String(ref v) if v == "llm")));
987        assert!(attrs
988            .iter()
989            .any(|kv| kv.key.as_str() == "fleetforge.tool.provider"
990                && matches!(kv.value, Value::String(ref v) if v == "internal")));
991        assert!(attrs.iter().any(
992            |kv| kv.key.as_str() == "fleetforge.run.seed" && matches!(kv.value, Value::I64(42))
993        ));
994    }
995
996    #[test]
997    fn baggage_header_is_percent_encoded() {
998        let ctx = TraceContext::new(RunScope::new("run-abc"))
999            .with_baggage([("Budget Tier", "gold;primary")]);
1000        let baggage = ctx.baggage_header().expect("baggage header");
1001        assert_eq!(baggage, "Budget%20Tier=gold%3Bprimary");
1002    }
1003
1004    #[test]
1005    fn from_w3c_parses_headers() {
1006        let traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
1007        let baggage = "tenant-id=acme,env=prod";
1008        let ctx = TraceContext::from_w3c(
1009            RunScope::new("run-xyz"),
1010            traceparent,
1011            Some("ro=1"),
1012            Some(baggage),
1013        )
1014        .expect("parse traceparent");
1015
1016        assert_eq!(
1017            ctx.trace_id().to_string(),
1018            "4bf92f3577b34da6a3ce929d0e0e4736"
1019        );
1020        assert_eq!(
1021            ctx.parent_span_id().unwrap().to_string(),
1022            "00f067aa0ba902b7"
1023        );
1024        assert_eq!(ctx.trace_flags(), TraceFlags::SAMPLED);
1025        assert_eq!(ctx.tracestate(), Some("ro=1".to_string()));
1026        assert_eq!(
1027            ctx.baggage,
1028            vec![
1029                ("tenant-id".to_string(), "acme".to_string()),
1030                ("env".to_string(), "prod".to_string()),
1031            ]
1032        );
1033    }
1034
1035    #[test]
1036    fn to_json_includes_basic_fields() {
1037        let ctx = TraceContext::new(RunScope::new("run-xyz"))
1038            .with_seed(Some(42))
1039            .child_step(
1040                StepScope::new("step-123")
1041                    .with_index(3)
1042                    .with_attempt(2)
1043                    .with_kind("tool"),
1044            );
1045        let json = ctx.to_json();
1046        let run_id = json
1047            .get("run")
1048            .and_then(|value| value.as_object())
1049            .and_then(|map| map.get("run_id"))
1050            .and_then(|value| value.as_str())
1051            .unwrap();
1052        assert_eq!(run_id, "run-xyz");
1053        let step_kind = json
1054            .get("step")
1055            .and_then(|value| value.as_object())
1056            .and_then(|map| map.get("kind"))
1057            .and_then(|value| value.as_str())
1058            .unwrap();
1059        assert_eq!(step_kind, "tool");
1060    }
1061}
fleetforge_telemetry/context.rs

fleetforge_telemetry/
context.rs
