From 30fcee98c220fc37ce76638a6e66872abb0d7b4e Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Thu, 9 Apr 2026 12:39:40 -0400 Subject: [PATCH] refactor(events): type stage/parallel ids with newtypes Promote RunEvent.stage_id / parallel_group_id / parallel_branch_id and the internal Event enum's matching fields from stringly-typed Option to Option / Option. The wire contract is now self-enforcing: malformed strings are rejected at the serde seam, not quietly round-tripped, and the three StageId::new(...).to_string() calls in stored_event_fields() just drop the .to_string() since the newtypes flow straight through. - fabro-types/src/stage_id.rs: new ParallelBranchId { group: StageId, index: u32 } mirroring StageId's Display / FromStr / serde string form. "{group}:{index}" (e.g. "fanout@2:0"). Tests for round-trip and parse rejections. - fabro-types/src/lib.rs: re-export ParallelBranchId. - fabro-types/src/run_event/mod.rs: RunEvent, RunEventRaw, and RunEventParts take Option / Option. from_ref gains a small generic opt_field helper that also replaces the bespoke actor null-handling branch. to_value uses serde_json::to_value(value) for the three typed fields. - fabro-workflow/src/event.rs: Event::Stage{Started,Completed, Failed,Retrying} and Event::Agent take Option / Option. Event::ParallelBranch{Started,Completed} take the required (non-Option) typed forms. StoredEventFields and stored_event_fields() plumb the newtypes end-to-end. - fabro-workflow/src/context.rs: WorkflowContext::parallel_group_id() returns Option, parallel_branch_id() returns Option. Read via serde_json::from_value which validates the shape on the way out. - fabro-workflow/src/handler/parallel.rs: builds typed values directly, stores in context via serde_json::to_value (still produces a JSON string through the custom Serialize). BranchSetup holds a ParallelBranchId. - fabro-workflow/src/handler/llm/api.rs: StageEventScope holds typed ids. - fabro-workflow/src/lifecycle/event.rs: stage_parallel_ids returns typed tuple. Wire JSON is byte-identical before and after (StageId serializes as "{node_id}@{visit}", ParallelBranchId as "{node_id}@{visit}:{index}", matching the existing spec). Progenitor-generated types and OpenAPI schema untouched. Existing None-only fixtures in runtime_store, git, pipeline, error, run_state, rewind, pr_view, and store/dump didn't need any edit because None fits any Option. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/commands/run/run_progress/mod.rs | 22 ++-- lib/crates/fabro-store/src/types.rs | 4 +- lib/crates/fabro-types/src/lib.rs | 2 +- lib/crates/fabro-types/src/run_event/mod.rs | 57 +++++---- lib/crates/fabro-types/src/stage_id.rs | 120 +++++++++++++++++- lib/crates/fabro-workflow/src/context.rs | 20 +-- lib/crates/fabro-workflow/src/event.rs | 87 +++++++------ .../fabro-workflow/src/handler/llm/api.rs | 5 +- .../fabro-workflow/src/handler/parallel.rs | 16 ++- .../fabro-workflow/src/lifecycle/event.rs | 4 +- 10 files changed, 241 insertions(+), 96 deletions(-) diff --git a/lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs b/lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs index 5c5ced36d..5e15deb73 100644 --- a/lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs +++ b/lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs @@ -420,7 +420,7 @@ mod tests { use fabro_agent::{AgentEvent, SandboxEvent}; use fabro_llm::types::TokenCounts; use fabro_model::Provider; - use fabro_types::fixtures; + use fabro_types::{ParallelBranchId, StageId, fixtures}; use fabro_workflow::event::{Event, RunNoticeLevel, to_run_event, to_run_event_at}; use fabro_workflow::outcome::billed_model_usage_from_llm; @@ -569,8 +569,8 @@ mod tests { emit( &mut ui, Event::ParallelBranchStarted { - parallel_group_id: "fork1@1".into(), - parallel_branch_id: "fork1@1:0".into(), + parallel_group_id: StageId::new("fork1", 1), + parallel_branch_id: ParallelBranchId::new(StageId::new("fork1", 1), 0), branch: "security".into(), index: 0, }, @@ -586,8 +586,8 @@ mod tests { emit( &mut ui, Event::ParallelBranchCompleted { - parallel_group_id: "fork1@1".into(), - parallel_branch_id: "fork1@1:0".into(), + parallel_group_id: StageId::new("fork1", 1), + parallel_branch_id: ParallelBranchId::new(StageId::new("fork1", 1), 0), branch: "security".into(), index: 0, duration_ms: 2000, @@ -619,8 +619,8 @@ mod tests { emit( &mut ui, Event::ParallelBranchStarted { - parallel_group_id: "fork1@1".into(), - parallel_branch_id: "fork1@1:0".into(), + parallel_group_id: StageId::new("fork1", 1), + parallel_branch_id: ParallelBranchId::new(StageId::new("fork1", 1), 0), branch: "security".into(), index: 0, }, @@ -1128,8 +1128,8 @@ mod tests { emit( &mut ui, Event::ParallelBranchStarted { - parallel_group_id: "fork1@1".into(), - parallel_branch_id: "fork1@1:0".into(), + parallel_group_id: StageId::new("fork1", 1), + parallel_branch_id: ParallelBranchId::new(StageId::new("fork1", 1), 0), branch: "security".into(), index: 0, }, @@ -1137,8 +1137,8 @@ mod tests { emit( &mut ui, Event::ParallelBranchCompleted { - parallel_group_id: "fork1@1".into(), - parallel_branch_id: "fork1@1:0".into(), + parallel_group_id: StageId::new("fork1", 1), + parallel_branch_id: ParallelBranchId::new(StageId::new("fork1", 1), 0), branch: "security".into(), index: 0, duration_ms: 500, diff --git a/lib/crates/fabro-store/src/types.rs b/lib/crates/fabro-store/src/types.rs index 4c3d23f75..35b485f01 100644 --- a/lib/crates/fabro-store/src/types.rs +++ b/lib/crates/fabro-store/src/types.rs @@ -91,7 +91,7 @@ pub struct EventEnvelope { mod tests { use chrono::{TimeZone, Utc}; - use fabro_types::{EventBody, RunEvent, fixtures, run_event::RunCompletedProps}; + use fabro_types::{EventBody, RunEvent, StageId, fixtures, run_event::RunCompletedProps}; use super::{EventEnvelope, EventPayload}; @@ -103,7 +103,7 @@ mod tests { run_id: fixtures::RUN_1, node_id: Some("code".to_string()), node_label: Some("Code".to_string()), - stage_id: Some("code@1".to_string()), + stage_id: Some(StageId::new("code", 1)), parallel_group_id: None, parallel_branch_id: None, session_id: None, diff --git a/lib/crates/fabro-types/src/lib.rs b/lib/crates/fabro-types/src/lib.rs index c6dda288f..7dd665897 100644 --- a/lib/crates/fabro-types/src/lib.rs +++ b/lib/crates/fabro-types/src/lib.rs @@ -54,7 +54,7 @@ pub use run_id::RunId; pub use run_id::fixtures; pub use sandbox_record::SandboxRecord; pub use settings::{ArtifactStorageBackend, ArtifactStorageSettings, Settings}; -pub use stage_id::StageId; +pub use stage_id::{ParallelBranchId, StageId}; pub use start::StartRecord; pub use status::{ InvalidTransition, ParseRunStatusError, RunControlAction, RunStatus, RunStatusRecord, diff --git a/lib/crates/fabro-types/src/run_event/mod.rs b/lib/crates/fabro-types/src/run_event/mod.rs index d8a1bc1db..80fe77849 100644 --- a/lib/crates/fabro-types/src/run_event/mod.rs +++ b/lib/crates/fabro-types/src/run_event/mod.rs @@ -11,7 +11,7 @@ use serde::ser::Error as SerError; use serde::{Deserialize, Deserializer, Serialize, Serializer}; use serde_json::{Map, Value, json}; -use crate::RunId; +use crate::{ParallelBranchId, RunId, StageId}; pub use agent::*; pub use infra::*; @@ -51,9 +51,9 @@ pub struct RunEvent { pub run_id: RunId, pub node_id: Option, pub node_label: Option, - pub stage_id: Option, - pub parallel_group_id: Option, - pub parallel_branch_id: Option, + pub stage_id: Option, + pub parallel_group_id: Option, + pub parallel_branch_id: Option, pub session_id: Option, pub parent_session_id: Option, pub tool_call_id: Option, @@ -293,11 +293,11 @@ struct RunEventRaw { #[serde(default)] node_label: Option, #[serde(default)] - stage_id: Option, + stage_id: Option, #[serde(default)] - parallel_group_id: Option, + parallel_group_id: Option, #[serde(default)] - parallel_branch_id: Option, + parallel_branch_id: Option, #[serde(default)] session_id: Option, #[serde(default)] @@ -321,9 +321,9 @@ struct RunEventParts<'a> { run_id: RunId, node_id: Option, node_label: Option, - stage_id: Option, - parallel_group_id: Option, - parallel_branch_id: Option, + stage_id: Option, + parallel_group_id: Option, + parallel_branch_id: Option, session_id: Option, parent_session_id: Option, tool_call_id: Option, @@ -592,6 +592,16 @@ impl RunEvent { } pub fn from_ref(value: &Value) -> serde_json::Result { + fn opt_field Deserialize<'a>>( + obj: &Map, + key: &str, + ) -> serde_json::Result> { + match obj.get(key) { + Some(value) if !value.is_null() => Ok(Some(T::deserialize(value)?)), + _ => Ok(None), + } + } + let obj = value.as_object().ok_or_else(|| { ::custom("run event must be a JSON object") })?; @@ -614,23 +624,19 @@ impl RunEvent { .get("properties") .cloned() .unwrap_or_else(default_properties); - let actor = match obj.get("actor") { - Some(value) if !value.is_null() => Some(ActorRef::deserialize(value)?), - _ => None, - }; Self::from_parts(RunEventParts { id: id.to_string(), ts, run_id, node_id: opt_str("node_id"), node_label: opt_str("node_label"), - stage_id: opt_str("stage_id"), - parallel_group_id: opt_str("parallel_group_id"), - parallel_branch_id: opt_str("parallel_branch_id"), + stage_id: opt_field(obj, "stage_id")?, + parallel_group_id: opt_field(obj, "parallel_group_id")?, + parallel_branch_id: opt_field(obj, "parallel_branch_id")?, session_id: opt_str("session_id"), parent_session_id: opt_str("parent_session_id"), tool_call_id: opt_str("tool_call_id"), - actor, + actor: opt_field(obj, "actor")?, event, properties: &properties, }) @@ -695,18 +701,18 @@ impl RunEvent { map.insert("node_label".to_string(), Value::String(value.clone())); } if let Some(value) = &self.stage_id { - map.insert("stage_id".to_string(), Value::String(value.clone())); + map.insert("stage_id".to_string(), serde_json::to_value(value)?); } if let Some(value) = &self.parallel_group_id { map.insert( "parallel_group_id".to_string(), - Value::String(value.clone()), + serde_json::to_value(value)?, ); } if let Some(value) = &self.parallel_branch_id { map.insert( "parallel_branch_id".to_string(), - Value::String(value.clone()), + serde_json::to_value(value)?, ); } if let Some(value) = &self.tool_call_id { @@ -981,9 +987,12 @@ mod tests { }); let parsed = RunEvent::from_value(value.clone()).unwrap(); - assert_eq!(parsed.stage_id.as_deref(), Some("code@1")); - assert_eq!(parsed.parallel_group_id.as_deref(), Some("code@1")); - assert_eq!(parsed.parallel_branch_id.as_deref(), Some("code@1:0")); + assert_eq!(parsed.stage_id, Some(StageId::new("code", 1))); + assert_eq!(parsed.parallel_group_id, Some(StageId::new("code", 1))); + assert_eq!( + parsed.parallel_branch_id, + Some(ParallelBranchId::new(StageId::new("code", 1), 0)) + ); assert_eq!(parsed.tool_call_id.as_deref(), Some("call_1")); let actor = parsed.actor.as_ref().expect("actor present"); assert_eq!(actor.kind, ActorKind::Agent); diff --git a/lib/crates/fabro-types/src/stage_id.rs b/lib/crates/fabro-types/src/stage_id.rs index 747133c78..846d61cb8 100644 --- a/lib/crates/fabro-types/src/stage_id.rs +++ b/lib/crates/fabro-types/src/stage_id.rs @@ -90,9 +90,90 @@ impl<'de> Deserialize<'de> for StageId { } } +#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct ParallelBranchId { + group: StageId, + index: u32, +} + +impl ParallelBranchId { + #[must_use] + pub fn new(group: StageId, index: u32) -> Self { + Self { group, index } + } + + #[must_use] + pub fn group(&self) -> &StageId { + &self.group + } + + #[must_use] + pub fn index(&self) -> u32 { + self.index + } +} + +impl fmt::Display for ParallelBranchId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}:{}", self.group, self.index) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ParseParallelBranchIdError(String); + +impl fmt::Display for ParseParallelBranchIdError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl std::error::Error for ParseParallelBranchIdError {} + +impl FromStr for ParallelBranchId { + type Err = ParseParallelBranchIdError; + + fn from_str(s: &str) -> Result { + let (group, index) = s.rsplit_once(':').ok_or_else(|| { + ParseParallelBranchIdError("parallel branch id must contain ':'".to_string()) + })?; + let group = group.parse::().map_err(|err| { + ParseParallelBranchIdError(format!("invalid parallel group id: {err}")) + })?; + if index.is_empty() { + return Err(ParseParallelBranchIdError( + "parallel branch id index must not be empty".to_string(), + )); + } + let index = index.parse().map_err(|err| { + ParseParallelBranchIdError(format!("invalid parallel branch index: {err}")) + })?; + Ok(Self::new(group, index)) + } +} + +impl Serialize for ParallelBranchId { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&self.to_string()) + } +} + +impl<'de> Deserialize<'de> for ParallelBranchId { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + value.parse().map_err(D::Error::custom) + } +} + #[cfg(test)] mod tests { - use super::StageId; + use super::{ParallelBranchId, StageId}; #[test] fn display_and_parse_round_trip() { @@ -151,4 +232,41 @@ mod tests { let err = "@3".parse::().unwrap_err(); assert_eq!(err.to_string(), "stage id node_id must not be empty"); } + + #[test] + fn parallel_branch_id_display_and_parse_round_trip() { + let branch = ParallelBranchId::new(StageId::new("fanout", 2), 3); + assert_eq!(branch.to_string(), "fanout@2:3"); + assert_eq!("fanout@2:3".parse::().unwrap(), branch); + } + + #[test] + fn parallel_branch_id_serde_round_trip_uses_string_form() { + let branch = ParallelBranchId::new(StageId::new("fanout", 2), 0); + let value = serde_json::to_value(&branch).unwrap(); + assert_eq!(value, serde_json::json!("fanout@2:0")); + let decoded: ParallelBranchId = serde_json::from_value(value).unwrap(); + assert_eq!(decoded, branch); + } + + #[test] + fn parallel_branch_id_rejects_missing_colon() { + let err = "fanout@2".parse::().unwrap_err(); + assert_eq!(err.to_string(), "parallel branch id must contain ':'"); + } + + #[test] + fn parallel_branch_id_rejects_bad_group() { + let err = "fanout:0".parse::().unwrap_err(); + assert!(err.to_string().starts_with("invalid parallel group id:")); + } + + #[test] + fn parallel_branch_id_rejects_non_numeric_index() { + let err = "fanout@2:zero".parse::().unwrap_err(); + assert!( + err.to_string() + .starts_with("invalid parallel branch index:") + ); + } } diff --git a/lib/crates/fabro-workflow/src/context.rs b/lib/crates/fabro-workflow/src/context.rs index 2b33ceedb..153d8911c 100644 --- a/lib/crates/fabro-workflow/src/context.rs +++ b/lib/crates/fabro-workflow/src/context.rs @@ -136,6 +136,7 @@ pub mod keys { pub use fabro_core::Context; use fabro_graphviz::Fidelity; +use fabro_types::{ParallelBranchId, StageId}; /// Domain-specific typed accessors for workflow context values. pub trait WorkflowContext { @@ -143,8 +144,8 @@ pub trait WorkflowContext { fn thread_id(&self) -> Option; fn preamble(&self) -> String; fn run_id(&self) -> String; - fn parallel_group_id(&self) -> Option; - fn parallel_branch_id(&self) -> Option; + fn parallel_group_id(&self) -> Option; + fn parallel_branch_id(&self) -> Option; } impl WorkflowContext for Context { @@ -167,14 +168,14 @@ impl WorkflowContext for Context { self.get_string(keys::INTERNAL_RUN_ID, "unknown") } - fn parallel_group_id(&self) -> Option { + fn parallel_group_id(&self) -> Option { self.get(keys::INTERNAL_PARALLEL_GROUP_ID) - .and_then(|value| value.as_str().map(String::from)) + .and_then(|value| serde_json::from_value(value).ok()) } - fn parallel_branch_id(&self) -> Option { + fn parallel_branch_id(&self) -> Option { self.get(keys::INTERNAL_PARALLEL_BRANCH_ID) - .and_then(|value| value.as_str().map(String::from)) + .and_then(|value| serde_json::from_value(value).ok()) } } @@ -341,8 +342,11 @@ mod tests { keys::INTERNAL_PARALLEL_BRANCH_ID, serde_json::json!("fanout@2:1"), ); - assert_eq!(ctx.parallel_group_id(), Some("fanout@2".to_string())); - assert_eq!(ctx.parallel_branch_id(), Some("fanout@2:1".to_string())); + assert_eq!(ctx.parallel_group_id(), Some(StageId::new("fanout", 2))); + assert_eq!( + ctx.parallel_branch_id(), + Some(ParallelBranchId::new(StageId::new("fanout", 2), 1)) + ); } #[test] diff --git a/lib/crates/fabro-workflow/src/event.rs b/lib/crates/fabro-workflow/src/event.rs index b018e97f1..b2e0f96ef 100644 --- a/lib/crates/fabro-workflow/src/event.rs +++ b/lib/crates/fabro-workflow/src/event.rs @@ -5,8 +5,8 @@ use std::sync::atomic::{AtomicI64, Ordering}; use ::fabro_types::run_event as fabro_types; use ::fabro_types::{ - ActorKind, ActorRef, BilledTokenCounts, RunBlobId, RunControlAction, RunEvent, RunId, - RunProvenance, StageId, StageStatus, StatusReason, + ActorKind, ActorRef, BilledTokenCounts, ParallelBranchId, RunBlobId, RunControlAction, + RunEvent, RunId, RunProvenance, StageId, StageStatus, StatusReason, }; use anyhow::{Context, Result}; use chrono::Utc; @@ -140,9 +140,9 @@ pub enum Event { index: usize, visit: u32, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_group_id: Option, + parallel_group_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_branch_id: Option, + parallel_branch_id: Option, handler_type: String, attempt: usize, max_attempts: usize, @@ -153,9 +153,9 @@ pub enum Event { index: usize, visit: u32, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_group_id: Option, + parallel_group_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_branch_id: Option, + parallel_branch_id: Option, duration_ms: u64, status: String, preferred_label: Option, @@ -188,9 +188,9 @@ pub enum Event { index: usize, visit: u32, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_group_id: Option, + parallel_group_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_branch_id: Option, + parallel_branch_id: Option, failure: FailureDetail, will_retry: bool, }, @@ -200,9 +200,9 @@ pub enum Event { index: usize, visit: u32, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_group_id: Option, + parallel_group_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_branch_id: Option, + parallel_branch_id: Option, attempt: usize, max_attempts: usize, delay_ms: u64, @@ -214,14 +214,14 @@ pub enum Event { join_policy: String, }, ParallelBranchStarted { - parallel_group_id: String, - parallel_branch_id: String, + parallel_group_id: StageId, + parallel_branch_id: ParallelBranchId, branch: String, index: usize, }, ParallelBranchCompleted { - parallel_group_id: String, - parallel_branch_id: String, + parallel_group_id: StageId, + parallel_branch_id: ParallelBranchId, branch: String, index: usize, duration_ms: u64, @@ -378,9 +378,9 @@ pub enum Event { #[serde(default, skip_serializing_if = "Option::is_none")] parent_session_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_group_id: Option, + parallel_group_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - parallel_branch_id: Option, + parallel_branch_id: Option, }, SubgraphStarted { node_id: String, @@ -1285,9 +1285,9 @@ struct StoredEventFields { parent_session_id: Option, node_id: Option, node_label: Option, - stage_id: Option, - parallel_group_id: Option, - parallel_branch_id: Option, + stage_id: Option, + parallel_group_id: Option, + parallel_branch_id: Option, tool_call_id: Option, actor: Option, } @@ -1361,7 +1361,7 @@ fn stored_event_fields(event: &Event) -> StoredEventFields { } => { let node_id_str = node_id.clone(); let node_label = default_node_label(Some(&node_id_str), Some(name.clone())); - let stage_id = Some(StageId::new(node_id_str.clone(), *visit).to_string()); + let stage_id = Some(StageId::new(node_id_str.clone(), *visit)); StoredEventFields { node_id: Some(node_id_str), node_label, @@ -1375,7 +1375,7 @@ fn stored_event_fields(event: &Event) -> StoredEventFields { | Event::ParallelCompleted { node_id, visit, .. } => { let node_id_str = node_id.clone(); let node_label = default_node_label(Some(&node_id_str), None); - let parallel_group_id = Some(StageId::new(node_id_str.clone(), *visit).to_string()); + let parallel_group_id = Some(StageId::new(node_id_str.clone(), *visit)); StoredEventFields { node_id: Some(node_id_str), node_label, @@ -1404,7 +1404,7 @@ fn stored_event_fields(event: &Event) -> StoredEventFields { } => { let node_id = Some(stage.clone()); let node_label = default_node_label(node_id.as_ref(), None); - let stage_id = Some(StageId::new(stage.clone(), *visit).to_string()); + let stage_id = Some(StageId::new(stage.clone(), *visit)); let tool_call_id = agent_tool_call_id(agent_event).map(str::to_string); let actor = agent_actor_for_event(agent_event, session_id.as_deref()); StoredEventFields { @@ -2891,7 +2891,7 @@ mod tests { assert_eq!(stored.run_id, fixtures::RUN_2); assert_eq!(stored.node_id.as_deref(), Some("plan")); assert_eq!(stored.node_label.as_deref(), Some("Plan")); - assert_eq!(stored.stage_id.as_deref(), Some("plan@1")); + assert_eq!(stored.stage_id, Some(StageId::new("plan", 1))); let properties = stored.properties().unwrap(); assert_eq!(properties["duration_ms"], 5000); assert_eq!(properties["status"], "success"); @@ -3133,8 +3133,8 @@ mod tests { ); assert_eq!( event_name(&Event::ParallelBranchStarted { - parallel_group_id: "plan@1".to_string(), - parallel_branch_id: "plan@1:0".to_string(), + parallel_group_id: StageId::new("plan", 1), + parallel_branch_id: ParallelBranchId::new(StageId::new("plan", 1), 0), branch: "fork".to_string(), index: 0, }), @@ -3167,15 +3167,18 @@ mod tests { name: "review".to_string(), index: 1, visit: 1, - parallel_group_id: Some("fanout@2".to_string()), - parallel_branch_id: Some("fanout@2:1".to_string()), + parallel_group_id: Some(StageId::new("fanout", 2)), + parallel_branch_id: Some(ParallelBranchId::new(StageId::new("fanout", 2), 1)), handler_type: "agent".to_string(), attempt: 1, max_attempts: 1, }, ); - assert_eq!(stored.parallel_group_id.as_deref(), Some("fanout@2")); - assert_eq!(stored.parallel_branch_id.as_deref(), Some("fanout@2:1")); + assert_eq!(stored.parallel_group_id, Some(StageId::new("fanout", 2))); + assert_eq!( + stored.parallel_branch_id, + Some(ParallelBranchId::new(StageId::new("fanout", 2), 1)) + ); } #[test] @@ -3189,7 +3192,7 @@ mod tests { join_policy: "wait_all".to_string(), }, ); - assert_eq!(stored.parallel_group_id.as_deref(), Some("fanout@2")); + assert_eq!(stored.parallel_group_id, Some(StageId::new("fanout", 2))); assert!(stored.parallel_branch_id.is_none()); } @@ -3198,14 +3201,17 @@ mod tests { let stored = to_run_event( &fixtures::RUN_1, &Event::ParallelBranchStarted { - parallel_group_id: "fanout@2".to_string(), - parallel_branch_id: "fanout@2:1".to_string(), + parallel_group_id: StageId::new("fanout", 2), + parallel_branch_id: ParallelBranchId::new(StageId::new("fanout", 2), 1), branch: "review".to_string(), index: 1, }, ); - assert_eq!(stored.parallel_group_id.as_deref(), Some("fanout@2")); - assert_eq!(stored.parallel_branch_id.as_deref(), Some("fanout@2:1")); + assert_eq!(stored.parallel_group_id, Some(StageId::new("fanout", 2))); + assert_eq!( + stored.parallel_branch_id, + Some(ParallelBranchId::new(StageId::new("fanout", 2), 1)) + ); } #[test] @@ -3222,14 +3228,17 @@ mod tests { }, session_id: Some("ses_1".to_string()), parent_session_id: None, - parallel_group_id: Some("fanout@2".to_string()), - parallel_branch_id: Some("fanout@2:0".to_string()), + parallel_group_id: Some(StageId::new("fanout", 2)), + parallel_branch_id: Some(ParallelBranchId::new(StageId::new("fanout", 2), 0)), }, ); - assert_eq!(stored.stage_id.as_deref(), Some("code@3")); + assert_eq!(stored.stage_id, Some(StageId::new("code", 3))); assert_eq!(stored.tool_call_id.as_deref(), Some("call_abc")); - assert_eq!(stored.parallel_group_id.as_deref(), Some("fanout@2")); - assert_eq!(stored.parallel_branch_id.as_deref(), Some("fanout@2:0")); + assert_eq!(stored.parallel_group_id, Some(StageId::new("fanout", 2))); + assert_eq!( + stored.parallel_branch_id, + Some(ParallelBranchId::new(StageId::new("fanout", 2), 0)) + ); } #[test] diff --git a/lib/crates/fabro-workflow/src/handler/llm/api.rs b/lib/crates/fabro-workflow/src/handler/llm/api.rs index a95b0a64d..73390f201 100644 --- a/lib/crates/fabro-workflow/src/handler/llm/api.rs +++ b/lib/crates/fabro-workflow/src/handler/llm/api.rs @@ -13,6 +13,7 @@ use fabro_llm::types::{Message, Request, TokenCounts}; use fabro_mcp::config::McpServerSettings; use fabro_model::FallbackTarget; use fabro_model::Provider; +use fabro_types::{ParallelBranchId, StageId}; use tokio::sync::Mutex as TokioMutex; use super::super::agent::{CodergenBackend, CodergenResult}; @@ -40,8 +41,8 @@ fn build_profile(model: &str, provider: Provider) -> Box { #[derive(Clone)] struct StageEventScope { visit: u32, - parallel_group_id: Option, - parallel_branch_id: Option, + parallel_group_id: Option, + parallel_branch_id: Option, } fn current_stage_event_scope(context: &Context) -> StageEventScope { diff --git a/lib/crates/fabro-workflow/src/handler/parallel.rs b/lib/crates/fabro-workflow/src/handler/parallel.rs index ccd07125d..968e9044b 100644 --- a/lib/crates/fabro-workflow/src/handler/parallel.rs +++ b/lib/crates/fabro-workflow/src/handler/parallel.rs @@ -4,7 +4,7 @@ use std::time::Instant; use async_trait::async_trait; use fabro_agent::{Sandbox, WorktreeOptions, WorktreeSandbox}; -use fabro_types::{RunId, StageId}; +use fabro_types::{ParallelBranchId, RunId, StageId}; use tokio::sync::Semaphore; use crate::context::keys; @@ -132,7 +132,7 @@ impl Handler for ParallelHandler { struct BranchSetup { target_id: String, branch_index: usize, - parallel_branch_id: String, + parallel_branch_id: ParallelBranchId, branch_context: Context, sandbox: Arc, worktree_path: Option, @@ -152,7 +152,7 @@ impl Handler for ParallelHandler { ); let parallel_visit = u32::try_from(visit_from_context(context)).unwrap_or(u32::MAX); - let parallel_group_id = StageId::new(node.id.clone(), parallel_visit).to_string(); + let parallel_group_id = StageId::new(node.id.clone(), parallel_visit); services.emitter.emit(&Event::ParallelStarted { node_id: node.id.clone(), @@ -208,14 +208,18 @@ impl Handler for ParallelHandler { for (branch_index, edge) in branches.iter().enumerate() { let target_id = edge.to.clone(); let branch_context = context.fork(); - let parallel_branch_id = format!("{parallel_group_id}:{branch_index}"); + let parallel_branch_id = ParallelBranchId::new( + parallel_group_id.clone(), + u32::try_from(branch_index).unwrap_or(u32::MAX), + ); branch_context.set( keys::INTERNAL_PARALLEL_GROUP_ID, - serde_json::json!(¶llel_group_id), + serde_json::to_value(¶llel_group_id).expect("StageId serializes as string"), ); branch_context.set( keys::INTERNAL_PARALLEL_BRANCH_ID, - serde_json::json!(¶llel_branch_id), + serde_json::to_value(¶llel_branch_id) + .expect("ParallelBranchId serializes as string"), ); let (branch_sandbox, worktree_path): (Arc, Option) = if let ( diff --git a/lib/crates/fabro-workflow/src/lifecycle/event.rs b/lib/crates/fabro-workflow/src/lifecycle/event.rs index 6bca3b85e..00ce55b28 100644 --- a/lib/crates/fabro-workflow/src/lifecycle/event.rs +++ b/lib/crates/fabro-workflow/src/lifecycle/event.rs @@ -23,7 +23,7 @@ use crate::event::{Emitter, Event}; use crate::graph::WorkflowGraph; use crate::graph::WorkflowNode; use crate::outcome::{BilledModelUsage, FailureCategory, FailureDetail, Outcome, StageStatus}; -use fabro_types::{BilledTokenCounts, RunId, StatusReason}; +use fabro_types::{BilledTokenCounts, ParallelBranchId, RunId, StageId, StatusReason}; type WfRunState = ExecutionState>; type WfNodeResult = NodeResult>; @@ -85,7 +85,7 @@ fn stage_visit(state: &WfRunState, node_id: &str) -> u32 { u32::try_from(visits.max(1)).unwrap_or(u32::MAX) } -fn stage_parallel_ids(state: &WfRunState) -> (Option, Option) { +fn stage_parallel_ids(state: &WfRunState) -> (Option, Option) { ( state.context.parallel_group_id(), state.context.parallel_branch_id(),