From c5865d38c541ab1f8a83ed2e4bd0ba0b15e917da Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Tue, 5 May 2026 13:10:33 -0400 Subject: [PATCH] feat(workflow): replace steering lifecycle events --- apps/fabro-web/app/lib/run-events.test.tsx | 2 +- apps/fabro-web/app/lib/run-events.ts | 4 +- docs/internal/events.md | 47 ++- lib/crates/fabro-agent/src/session.rs | 33 +- lib/crates/fabro-server/src/server.rs | 45 ++- lib/crates/fabro-server/src/server/tests.rs | 61 ++++ lib/crates/fabro-store/src/run_state.rs | 68 +++- lib/crates/fabro-types/src/lib.rs | 2 +- lib/crates/fabro-types/src/run_event/agent.rs | 34 +- lib/crates/fabro-types/src/run_event/mod.rs | 45 ++- .../fabro-workflow/src/event/convert.rs | 48 ++- lib/crates/fabro-workflow/src/event/events.rs | 70 +++- lib/crates/fabro-workflow/src/event/names.rs | 6 +- .../fabro-workflow/src/event/stored_fields.rs | 27 +- .../fabro-workflow/src/handler/agent.rs | 25 +- .../fabro-workflow/src/handler/fan_in.rs | 6 + .../src/handler/llm/activation_lease.rs | 248 ++++++++++++++ .../fabro-workflow/src/handler/llm/api.rs | 311 +++++++++++++++--- .../fabro-workflow/src/handler/llm/cli.rs | 4 + .../fabro-workflow/src/handler/llm/mod.rs | 1 + lib/crates/fabro-workflow/src/handler/mod.rs | 10 + .../fabro-workflow/src/handler/prompt.rs | 9 +- .../fabro-workflow/src/operations/fork.rs | 24 +- .../fabro-workflow/src/pipeline/execute.rs | 2 + lib/crates/fabro-workflow/src/steering_hub.rs | 205 +++++++----- 25 files changed, 1125 insertions(+), 212 deletions(-) create mode 100644 lib/crates/fabro-workflow/src/handler/llm/activation_lease.rs diff --git a/apps/fabro-web/app/lib/run-events.test.tsx b/apps/fabro-web/app/lib/run-events.test.tsx index 4f52ea4c3..bf35aad8c 100644 --- a/apps/fabro-web/app/lib/run-events.test.tsx +++ b/apps/fabro-web/app/lib/run-events.test.tsx @@ -63,7 +63,7 @@ describe("queryKeysForRunEvent", () => { }); test("stage-scoped steering events invalidate run events and stage events", () => { - expect(queryKeysForRunEvent("run-1", "agent.steering.injected", "agent@1")).toEqual([ + expect(queryKeysForRunEvent("run-1", "agent.session.activated", "agent@1")).toEqual([ queryKeys.runs.events("run-1", 1000), queryKeys.runs.stageEvents("run-1", "agent@1"), ]); diff --git a/apps/fabro-web/app/lib/run-events.ts b/apps/fabro-web/app/lib/run-events.ts index 6be35b020..348f97c97 100644 --- a/apps/fabro-web/app/lib/run-events.ts +++ b/apps/fabro-web/app/lib/run-events.ts @@ -79,8 +79,8 @@ const INTERVIEW_EVENTS = new Set([ ]); const STEERING_EVENTS = new Set([ "agent.steering.injected", - "agent.steering.attached", - "agent.steering.detached", + "agent.session.activated", + "agent.session.deactivated", "agent.steer.buffered", "agent.steer.dropped", ]); diff --git a/docs/internal/events.md b/docs/internal/events.md index 92e026b43..337274272 100644 --- a/docs/internal/events.md +++ b/docs/internal/events.md @@ -863,7 +863,7 @@ Emitted when execution loops back to an earlier node. ## Agent events -All agent events have `node_id` (the workflow stage), `node_label`, `session_id`, and `parent_session_id` in the envelope. The `properties` contain the inner agent event fields. +Most agent activity events are stage-scoped and carry `node_id` (the workflow stage), `node_label`, `stage_id`, `session_id`, and `parent_session_id` in the envelope. Session object lifecycle events are the exception: `agent.session.started` and `agent.session.ended` are not stage-scoped and intentionally omit `node_id`, `node_label`, `stage_id`, and `visit`. ### `agent.session.started` @@ -871,13 +871,49 @@ All agent events have `node_id` (the workflow stage), `node_label`, `session_id` { "id": "...", "ts": "...", "run_id": "...", "event": "agent.session.started", - "node_id": "code", "node_label": "code", "session_id": "ses_abc", "parent_session_id": null, - "properties": {} + "properties": { + "provider": "openai", + "model": "gpt-5.4" + } } ``` -No properties. +Object-lifecycle event. `session_id` and `parent_session_id` are envelope fields. `properties.provider` and `properties.model` are optional. + +### `agent.session.activated` + +```json +{ + "id": "...", "ts": "...", "run_id": "...", + "event": "agent.session.activated", + "node_id": "code", "node_label": "code", "stage_id": "code@1", + "session_id": "ses_abc", + "properties": { + "thread_id": "main", + "provider": "openai", + "model": "gpt-5.4", + "capabilities": ["steer"], + "visit": 1 + } +} +``` + +Stage-scoped lease event. A stage is steerable while the latest matching `agent.session.activated` lease is active. + +### `agent.session.deactivated` + +```json +{ + "id": "...", "ts": "...", "run_id": "...", + "event": "agent.session.deactivated", + "node_id": "code", "node_label": "code", "stage_id": "code@1", + "session_id": "ses_abc", + "properties": { "visit": 1 } +} +``` + +Stage-scoped lease event. Consumers should pair it by `stage_id` and `session_id` so stale deactivations cannot clear a newer active lease. ### `agent.session.ended` @@ -885,13 +921,12 @@ No properties. { "id": "...", "ts": "...", "run_id": "...", "event": "agent.session.ended", - "node_id": "code", "node_label": "code", "session_id": "ses_abc", "properties": {} } ``` -No properties. +Object-lifecycle event. `session_id` and `parent_session_id` are envelope fields. No properties. ### `agent.processing.end` diff --git a/lib/crates/fabro-agent/src/session.rs b/lib/crates/fabro-agent/src/session.rs index 775d6a610..71187cab8 100644 --- a/lib/crates/fabro-agent/src/session.rs +++ b/lib/crates/fabro-agent/src/session.rs @@ -13,6 +13,7 @@ use fabro_llm::types::{ use fabro_llm::{Error as LlmError, retry}; use fabro_mcp::config::{McpServerSettings, McpTransport}; use fabro_mcp::connection_manager::McpConnectionManager; +use fabro_model::Provider; use fabro_types::{Principal, SteerKind}; use futures::StreamExt; use tokio::sync::{Mutex as AsyncMutex, broadcast}; @@ -248,6 +249,16 @@ impl Session { &self.id } + #[must_use] + pub fn provider(&self) -> Provider { + self.provider_profile.provider() + } + + #[must_use] + pub fn model(&self) -> &str { + self.provider_profile.model() + } + /// Initialize session by discovering project docs and capturing environment /// context. Call before `process_input`. /// @@ -821,8 +832,10 @@ impl Session { self.state = to; } - pub fn close(&mut self) { + pub fn close(&mut self) -> bool { + let was_open = self.state != SessionState::Closed; self.transition(SessionState::Closed); + was_open } pub fn set_reasoning_effort(&mut self, effort: Option) { @@ -2064,6 +2077,24 @@ mod tests { assert!(matches!(result.unwrap_err(), Error::SessionClosed)); } + #[tokio::test] + async fn close_reports_whether_it_transitioned_to_closed() { + let mut session = make_session(vec![]).await; + let mut rx = session.subscribe(); + + assert!(session.close()); + assert!(!session.close()); + + let events: Vec<_> = std::iter::from_fn(|| rx.try_recv().ok()).collect(); + assert_eq!( + events + .iter() + .filter(|event| matches!(event.event, AgentEvent::SessionEnded)) + .count(), + 1 + ); + } + #[tokio::test] async fn closed_session_does_not_emit_session_start() { let mut session = make_session(vec![]).await; diff --git a/lib/crates/fabro-server/src/server.rs b/lib/crates/fabro-server/src/server.rs index 51f1ea935..e07105556 100644 --- a/lib/crates/fabro-server/src/server.rs +++ b/lib/crates/fabro-server/src/server.rs @@ -80,7 +80,7 @@ use fabro_types::settings::server::{ use fabro_types::settings::{InterpString, RunNamespace}; use fabro_types::{ EventBody, InterviewQuestionRecord, Principal, PullRequestRecord, QuestionType, RunBlobId, - RunControlAction, RunEvent, RunId, ServerSettings, + RunControlAction, RunEvent, RunId, ServerSettings, SessionCapability, }; use fabro_util::error::{SharedError, collect_causes, render_with_causes}; use fabro_util::version::FABRO_VERSION; @@ -196,10 +196,10 @@ struct ManagedRun { // Populated when running: answer_transport: Option, accepted_questions: HashSet, - /// Stage IDs of currently running API-mode (SDK) agent sessions, as - /// observed from the worker's `agent.steering.attached/detached` - /// events. Used by the steerability predicate. - active_api_stages: HashSet, + /// Stage IDs of currently steerable API-mode (SDK) agent sessions, + /// keyed to the session id that owns the active lease. Used by the + /// steerability predicate. + active_api_stages: HashMap, /// Stage IDs of currently running CLI-mode agent sessions, observed /// from `agent.cli.started/completed` plus `stage.completed`/ /// `stage.failed` backstops. @@ -2148,7 +2148,7 @@ fn managed_run( enqueued_at: Instant::now(), answer_transport: None, accepted_questions: HashSet::new(), - active_api_stages: HashSet::new(), + active_api_stages: HashMap::new(), active_cli_stages: HashSet::new(), event_tx: None, checkpoint: None, @@ -2266,18 +2266,31 @@ fn update_live_run_from_event(state: &AppState, run_id: RunId, event: &RunEvent) managed_run.status = prior.into(); } } - // Track API-mode steerable sessions. attached/detached fire - // deterministically inside `AgentApiBackend::run` (register/ - // unregister), so they're the authoritative window in which a - // steer can be delivered to a live session. - EventBody::AgentSteeringAttached(_) => { - if let Some(stage_id) = event.stage_id.as_ref() { - managed_run.active_api_stages.insert(stage_id.clone()); + // Track API-mode steerable sessions. Activated/deactivated are + // leased by session id so stale deactivations cannot clear a newer + // binding for the same stage. + EventBody::AgentSessionActivated(props) + if props.capabilities.contains(&SessionCapability::Steer) => + { + if let (Some(stage_id), Some(session_id)) = + (event.stage_id.as_ref(), event.session_id.as_ref()) + { + managed_run + .active_api_stages + .insert(stage_id.clone(), session_id.clone()); } } - EventBody::AgentSteeringDetached(_) => { - if let Some(stage_id) = &event.stage_id { - managed_run.active_api_stages.remove(stage_id); + EventBody::AgentSessionDeactivated(_) => { + if let (Some(stage_id), Some(session_id)) = + (event.stage_id.as_ref(), event.session_id.as_ref()) + { + if managed_run + .active_api_stages + .get(stage_id) + .is_some_and(|current| current == session_id) + { + managed_run.active_api_stages.remove(stage_id); + } } } // Track CLI-mode agent stages. CLI started/completed are coarser diff --git a/lib/crates/fabro-server/src/server/tests.rs b/lib/crates/fabro-server/src/server/tests.rs index eb3c77497..bb3403506 100644 --- a/lib/crates/fabro-server/src/server/tests.rs +++ b/lib/crates/fabro-server/src/server/tests.rs @@ -6120,6 +6120,67 @@ async fn steer_empty_text_returns_bad_request() { ); } +#[test] +fn active_api_stage_projection_ignores_stale_deactivation() { + let state = test_app_state(); + let run_id = fixtures::RUN_1; + let temp_dir = tempfile::tempdir().unwrap(); + { + let mut runs = state.runs.lock().expect("runs lock poisoned"); + runs.insert( + run_id, + managed_run( + String::new(), + RunStatus::Running, + chrono::Utc::now(), + temp_dir.path().join(run_id.to_string()), + RunExecutionMode::Start, + ), + ); + } + + let stage_id = StageId::new("agent", 1); + let activated_a = + workflow_event::to_run_event(&run_id, &workflow_event::Event::AgentSessionActivated { + node_id: "agent".to_string(), + visit: 1, + session_id: "session-a".to_string(), + thread_id: None, + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + capabilities: vec![SessionCapability::Steer], + }); + update_live_run_from_event(&state, run_id, &activated_a); + + let deactivated_a = + workflow_event::to_run_event(&run_id, &workflow_event::Event::AgentSessionDeactivated { + node_id: "agent".to_string(), + visit: 1, + session_id: "session-a".to_string(), + }); + update_live_run_from_event(&state, run_id, &deactivated_a); + + let activated_b = + workflow_event::to_run_event(&run_id, &workflow_event::Event::AgentSessionActivated { + node_id: "agent".to_string(), + visit: 1, + session_id: "session-b".to_string(), + thread_id: None, + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + capabilities: vec![SessionCapability::Steer], + }); + update_live_run_from_event(&state, run_id, &activated_b); + update_live_run_from_event(&state, run_id, &deactivated_a); + + let runs = state.runs.lock().expect("runs lock poisoned"); + let run = runs.get(&run_id).unwrap(); + assert_eq!( + run.active_api_stages.get(&stage_id).map(String::as_str), + Some("session-b") + ); +} + #[tokio::test] async fn get_graph_returns_svg() { let state = test_app_state(); diff --git a/lib/crates/fabro-store/src/run_state.rs b/lib/crates/fabro-store/src/run_state.rs index 16b70e0cc..8ff510b7b 100644 --- a/lib/crates/fabro-store/src/run_state.rs +++ b/lib/crates/fabro-store/src/run_state.rs @@ -3,7 +3,7 @@ use std::str::FromStr; use chrono::{DateTime, Utc}; use fabro_types::run_event::{ - AgentCliStartedProps, AgentSessionStartedProps, CheckpointCompletedProps, RunCompletedProps, + AgentCliStartedProps, AgentSessionActivatedProps, CheckpointCompletedProps, RunCompletedProps, RunFailedProps, StageCompletedProps, StagePromptProps, }; use fabro_types::{ @@ -351,12 +351,12 @@ impl RunProjectionReducer for RunProjection { stage.usage.clone_from(&props.billing); stage.state = Some(StageState::from(outcome)); } - EventBody::AgentSessionStarted(props) => { + EventBody::AgentSessionActivated(props) => { let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) else { return Ok(()); }; - stage.provider_used = Some(provider_used_from_agent_session_started(props)); + stage.provider_used = Some(provider_used_from_agent_session_activated(props)); } EventBody::AgentCliStarted(props) => { let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) @@ -684,7 +684,7 @@ fn provider_used_from_prompt(props: &StagePromptProps) -> Option { (!provider_used.is_empty()).then_some(Value::Object(provider_used)) } -fn provider_used_from_agent_session_started(props: &AgentSessionStartedProps) -> Value { +fn provider_used_from_agent_session_activated(props: &AgentSessionActivatedProps) -> Value { let mut provider_used = serde_json::Map::new(); provider_used.insert("mode".to_string(), Value::String("agent".to_string())); if let Some(provider) = props.provider.clone() { @@ -732,6 +732,7 @@ mod tests { use fabro_types::run_event::run::RunFailedProps; use fabro_types::run_event::{ AgentCliCancelledProps, AgentCliCompletedProps, AgentCliTimedOutProps, + AgentSessionActivatedProps, AgentSessionEndedProps, AgentSessionStartedProps, CheckpointCompletedProps, InterviewCompletedProps, InterviewOption, InterviewStartedProps, RunControlEffectProps, StageCompletedProps, StageFailedProps, StagePromptProps, StageRetryingProps, StageStartedProps, @@ -1024,6 +1025,65 @@ mod tests { .unwrap(); } + #[test] + fn agent_session_activated_updates_stage_provider_used() { + let mut state = RunProjection::default(); + let stage_id = StageId::new("code", 1); + start_stage(&mut state, &stage_id); + + state + .apply_event(&test_stage_event( + 4, + EventBody::AgentSessionActivated(AgentSessionActivatedProps { + thread_id: Some("thread-1".to_string()), + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + capabilities: vec![fabro_types::SessionCapability::Steer], + visit: 1, + }), + stage_id.clone(), + )) + .unwrap(); + + let stage = state.stage(&stage_id).unwrap(); + assert_eq!( + stage.provider_used.as_ref().unwrap(), + &json!({ + "mode": "agent", + "provider": "openai", + "model": "gpt-5.4" + }) + ); + } + + #[test] + fn object_lifecycle_session_events_do_not_update_stage_provider_used() { + let mut state = RunProjection::default(); + let stage_id = StageId::new("code", 1); + start_stage(&mut state, &stage_id); + + state + .apply_event(&test_event( + 4, + EventBody::AgentSessionStarted(AgentSessionStartedProps { + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + }), + None, + )) + .unwrap(); + state + .apply_event(&test_event( + 5, + EventBody::AgentSessionEnded(AgentSessionEndedProps {}), + None, + )) + .unwrap(); + + let stage = state.stage(&stage_id).unwrap(); + assert!(stage.provider_used.is_none()); + } + #[test] fn agent_cli_completed_updates_stage_output_projection() { let mut state = RunProjection::default(); diff --git a/lib/crates/fabro-types/src/lib.rs b/lib/crates/fabro-types/src/lib.rs index 873d328a2..fafb01c05 100644 --- a/lib/crates/fabro-types/src/lib.rs +++ b/lib/crates/fabro-types/src/lib.rs @@ -73,7 +73,7 @@ pub use run::{ pub use run_blob_id::RunBlobId; pub use run_event::{ EventBody, ExecOutputTail, InterviewOption, MetadataSnapshotFailureKind, MetadataSnapshotPhase, - RunEvent, RunNoticeCode, RunNoticeLevel, + RunEvent, RunNoticeCode, RunNoticeLevel, SessionCapability, }; pub use run_id::{RunId, fixtures}; pub use run_projection::{PendingInterviewRecord, RunProjection, StageProjection, first_event_seq}; diff --git a/lib/crates/fabro-types/src/run_event/agent.rs b/lib/crates/fabro-types/src/run_event/agent.rs index 53516b620..9ef85d027 100644 --- a/lib/crates/fabro-types/src/run_event/agent.rs +++ b/lib/crates/fabro-types/src/run_event/agent.rs @@ -10,11 +10,35 @@ pub struct AgentSessionStartedProps { pub provider: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub model: Option, - pub visit: u32, +} + +#[allow( + clippy::empty_structs_with_brackets, + reason = "This type must serialize as {} rather than null." +)] +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct AgentSessionEndedProps {} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SessionCapability { + Steer, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct AgentSessionEndedProps { +pub struct AgentSessionActivatedProps { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub thread_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + pub capabilities: Vec, + pub visit: u32, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct AgentSessionDeactivatedProps { pub visit: u32, } @@ -87,12 +111,6 @@ pub struct AgentSteeringInjectedProps { pub visit: u32, } -#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] -pub struct AgentSteeringAttachedProps; - -#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] -pub struct AgentSteeringDetachedProps; - #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AgentSteerBufferedProps { pub kind: SteerKind, diff --git a/lib/crates/fabro-types/src/run_event/mod.rs b/lib/crates/fabro-types/src/run_event/mod.rs index 99220bce4..9df4a6456 100644 --- a/lib/crates/fabro-types/src/run_event/mod.rs +++ b/lib/crates/fabro-types/src/run_event/mod.rs @@ -148,6 +148,10 @@ pub enum EventBody { PromptCompleted(PromptCompletedProps), #[serde(rename = "agent.session.started")] AgentSessionStarted(AgentSessionStartedProps), + #[serde(rename = "agent.session.activated")] + AgentSessionActivated(AgentSessionActivatedProps), + #[serde(rename = "agent.session.deactivated")] + AgentSessionDeactivated(AgentSessionDeactivatedProps), #[serde(rename = "agent.session.ended")] AgentSessionEnded(AgentSessionEndedProps), #[serde(rename = "agent.processing.end")] @@ -170,10 +174,6 @@ pub enum EventBody { AgentTurnLimitReached(AgentTurnLimitReachedProps), #[serde(rename = "agent.steering.injected")] AgentSteeringInjected(AgentSteeringInjectedProps), - #[serde(rename = "agent.steering.attached")] - AgentSteeringAttached(AgentSteeringAttachedProps), - #[serde(rename = "agent.steering.detached")] - AgentSteeringDetached(AgentSteeringDetachedProps), #[serde(rename = "agent.steer.buffered")] AgentSteerBuffered(AgentSteerBufferedProps), #[serde(rename = "agent.steer.dropped")] @@ -393,6 +393,8 @@ impl EventBody { Self::StagePrompt(_) => "stage.prompt", Self::PromptCompleted(_) => "prompt.completed", Self::AgentSessionStarted(_) => "agent.session.started", + Self::AgentSessionActivated(_) => "agent.session.activated", + Self::AgentSessionDeactivated(_) => "agent.session.deactivated", Self::AgentSessionEnded(_) => "agent.session.ended", Self::AgentProcessingEnd(_) => "agent.processing.end", Self::AgentInput(_) => "agent.input", @@ -404,8 +406,6 @@ impl EventBody { Self::AgentLoopDetected(_) => "agent.loop.detected", Self::AgentTurnLimitReached(_) => "agent.turn.limit", Self::AgentSteeringInjected(_) => "agent.steering.injected", - Self::AgentSteeringAttached(_) => "agent.steering.attached", - Self::AgentSteeringDetached(_) => "agent.steering.detached", Self::AgentSteerBuffered(_) => "agent.steer.buffered", Self::AgentSteerDropped(_) => "agent.steer.dropped", Self::AgentCompactionStarted(_) => "agent.compaction.started", @@ -531,6 +531,8 @@ fn is_known_event_name(event: &str) -> bool { | "stage.prompt" | "prompt.completed" | "agent.session.started" + | "agent.session.activated" + | "agent.session.deactivated" | "agent.session.ended" | "agent.processing.end" | "agent.input" @@ -542,8 +544,6 @@ fn is_known_event_name(event: &str) -> bool { | "agent.loop.detected" | "agent.turn.limit" | "agent.steering.injected" - | "agent.steering.attached" - | "agent.steering.detached" | "agent.steer.buffered" | "agent.steer.dropped" | "agent.compaction.started" @@ -1046,6 +1046,35 @@ mod tests { assert_eq!(serialized["actor"], value["actor"]); } + #[test] + fn agent_session_ended_serializes_empty_properties() { + let event = RunEvent { + id: "evt_session_ended".to_string(), + ts: DateTime::parse_from_rfc3339("2026-04-04T12:00:00.000Z") + .unwrap() + .with_timezone(&Utc), + run_id: fixtures::RUN_1, + node_id: None, + node_label: None, + stage_id: None, + parallel_group_id: None, + parallel_branch_id: None, + session_id: Some("ses_abc".to_string()), + parent_session_id: None, + tool_call_id: None, + actor: None, + body: EventBody::AgentSessionEnded(AgentSessionEndedProps {}), + }; + + let serialized = event.to_value().unwrap(); + + assert_eq!(serialized["event"], "agent.session.ended"); + assert_eq!(serialized["session_id"], "ses_abc"); + assert_eq!(serialized["properties"], json!({})); + assert!(serialized.get("node_id").is_none()); + assert!(serialized.get("stage_id").is_none()); + } + #[test] fn run_event_omits_absent_envelope_fields() { let event = RunEvent { diff --git a/lib/crates/fabro-workflow/src/event/convert.rs b/lib/crates/fabro-workflow/src/event/convert.rs index 454a071e8..7f3c94c21 100644 --- a/lib/crates/fabro-workflow/src/event/convert.rs +++ b/lib/crates/fabro-workflow/src/event/convert.rs @@ -531,16 +531,6 @@ fn event_body_from_event(event: &Event) -> EventBody { billing: billing.clone(), }), Event::Agent { visit, event, .. } => match event { - AgentEvent::SessionStarted { provider, model } => { - EventBody::AgentSessionStarted(fabro_types::AgentSessionStartedProps { - provider: provider.clone(), - model: model.clone(), - visit: *visit, - }) - } - AgentEvent::SessionEnded => { - EventBody::AgentSessionEnded(fabro_types::AgentSessionEndedProps { visit: *visit }) - } AgentEvent::ProcessingEnd => { EventBody::AgentProcessingEnd(fabro_types::AgentProcessingEndProps { visit: *visit, @@ -707,9 +697,11 @@ fn event_body_from_event(event: &Event) -> EventBody { | AgentEvent::TextDelta { .. } | AgentEvent::ReasoningDelta { .. } | AgentEvent::ToolCallOutputDelta { .. } - | AgentEvent::SkillExpanded { .. } => { - panic!("streaming-noise agent event should not be converted to RunEvent") - } + | AgentEvent::SkillExpanded { .. } + | AgentEvent::SessionStarted { .. } + | AgentEvent::SessionEnded => panic!( + "agent event should not be converted through the stage-scoped Event::Agent wrapper" + ), }, Event::SubgraphStarted { start_node, .. } => { EventBody::SubgraphStarted(fabro_types::SubgraphStartedProps { @@ -1009,11 +1001,33 @@ fn event_body_from_event(event: &Event) -> EventBody { exit_code: *exit_code, duration_ms: *duration_ms, }), - Event::AgentSteeringAttached { .. } => { - EventBody::AgentSteeringAttached(fabro_types::AgentSteeringAttachedProps {}) + Event::AgentSessionStarted { + provider, model, .. + } => EventBody::AgentSessionStarted(fabro_types::AgentSessionStartedProps { + provider: provider.clone(), + model: model.clone(), + }), + Event::AgentSessionActivated { + thread_id, + provider, + model, + capabilities, + visit, + .. + } => EventBody::AgentSessionActivated(fabro_types::AgentSessionActivatedProps { + thread_id: thread_id.clone(), + provider: provider.clone(), + model: model.clone(), + capabilities: capabilities.clone(), + visit: *visit, + }), + Event::AgentSessionDeactivated { visit, .. } => { + EventBody::AgentSessionDeactivated(fabro_types::AgentSessionDeactivatedProps { + visit: *visit, + }) } - Event::AgentSteeringDetached { .. } => { - EventBody::AgentSteeringDetached(fabro_types::AgentSteeringDetachedProps {}) + Event::AgentSessionEnded { .. } => { + EventBody::AgentSessionEnded(fabro_types::AgentSessionEndedProps {}) } Event::AgentSteerBuffered { kind, .. } => { EventBody::AgentSteerBuffered(fabro_types::AgentSteerBufferedProps { kind: *kind }) diff --git a/lib/crates/fabro-workflow/src/event/events.rs b/lib/crates/fabro-workflow/src/event/events.rs index 4aa28496d..04bc998b2 100644 --- a/lib/crates/fabro-workflow/src/event/events.rs +++ b/lib/crates/fabro-workflow/src/event/events.rs @@ -524,16 +524,40 @@ pub enum Event { model: String, command: String, }, - /// A `SteeringHub` registered an active API-mode session for a stage. - /// Emitted once per `register` insert (not on replace). - AgentSteeringAttached { - node_id: String, - visit: u32, + /// A top-level agent session object started its lifecycle. + AgentSessionStarted { + session_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + parent_session_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + provider: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + model: Option, }, - /// The corresponding session was unregistered from the hub. - AgentSteeringDetached { - node_id: String, - visit: u32, + /// A stage has a currently steerable API-mode session binding. + AgentSessionActivated { + node_id: String, + visit: u32, + session_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + thread_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + provider: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + model: Option, + capabilities: Vec, + }, + /// A stage's steerable API-mode session binding ended. + AgentSessionDeactivated { + node_id: String, + visit: u32, + session_id: String, + }, + /// A top-level agent session object ended its lifecycle. + AgentSessionEnded { + session_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + parent_session_id: Option, }, /// A steer arrived with no active session and was parked in the run-wide /// pending buffer. The actor (steer author) is lifted to top-level. @@ -1290,11 +1314,31 @@ impl Event { } => { debug!(node_id, exit_code, duration_ms, "Agent CLI completed"); } - Self::AgentSteeringAttached { node_id, visit } => { - debug!(node_id, visit, "Steering hub attached to session"); + Self::AgentSessionStarted { + session_id, + provider, + model, + .. + } => { + debug!(session_id, ?provider, ?model, "Agent session started"); } - Self::AgentSteeringDetached { node_id, visit } => { - debug!(node_id, visit, "Steering hub detached from session"); + Self::AgentSessionActivated { + node_id, + visit, + session_id, + .. + } => { + debug!(node_id, visit, session_id, "Agent session activated"); + } + Self::AgentSessionDeactivated { + node_id, + visit, + session_id, + } => { + debug!(node_id, visit, session_id, "Agent session deactivated"); + } + Self::AgentSessionEnded { session_id, .. } => { + debug!(session_id, "Agent session ended"); } Self::AgentSteerBuffered { kind, .. } => { debug!(kind = kind.as_str(), "Steer buffered (no active session)"); diff --git a/lib/crates/fabro-workflow/src/event/names.rs b/lib/crates/fabro-workflow/src/event/names.rs index eb9e58a12..fda341bc2 100644 --- a/lib/crates/fabro-workflow/src/event/names.rs +++ b/lib/crates/fabro-workflow/src/event/names.rs @@ -116,8 +116,10 @@ pub fn event_name(event: &Event) -> &'static str { Event::CommandCompleted { .. } => "command.completed", Event::AgentCliStarted { .. } => "agent.cli.started", Event::AgentCliCompleted { .. } => "agent.cli.completed", - Event::AgentSteeringAttached { .. } => "agent.steering.attached", - Event::AgentSteeringDetached { .. } => "agent.steering.detached", + Event::AgentSessionStarted { .. } => "agent.session.started", + Event::AgentSessionActivated { .. } => "agent.session.activated", + Event::AgentSessionDeactivated { .. } => "agent.session.deactivated", + Event::AgentSessionEnded { .. } => "agent.session.ended", Event::AgentSteerBuffered { .. } => "agent.steer.buffered", Event::AgentSteerDropped { .. } => "agent.steer.dropped", Event::AgentCliCancelled { .. } => "agent.cli.cancelled", diff --git a/lib/crates/fabro-workflow/src/event/stored_fields.rs b/lib/crates/fabro-workflow/src/event/stored_fields.rs index d188d2133..8408af91f 100644 --- a/lib/crates/fabro-workflow/src/event/stored_fields.rs +++ b/lib/crates/fabro-workflow/src/event/stored_fields.rs @@ -120,11 +120,34 @@ fn stored_event_fields_for_variant(event: &Event) -> StoredEventFields { | Event::AgentCliCompleted { node_id, .. } | Event::AgentCliCancelled { node_id, .. } | Event::AgentCliTimedOut { node_id, .. } => node_stored_fields(Some(node_id.clone())), - Event::AgentSteeringAttached { node_id, visit } - | Event::AgentSteeringDetached { node_id, visit } => { + Event::AgentSessionStarted { + session_id, + parent_session_id, + .. + } + | Event::AgentSessionEnded { + session_id, + parent_session_id, + } => StoredEventFields { + session_id: Some(session_id.clone()), + parent_session_id: parent_session_id.clone(), + ..StoredEventFields::default() + }, + Event::AgentSessionActivated { + node_id, + visit, + session_id, + .. + } + | Event::AgentSessionDeactivated { + node_id, + visit, + session_id, + } => { let node_id_str = node_id.clone(); let node_label = default_node_label(Some(&node_id_str), None); StoredEventFields { + session_id: Some(session_id.clone()), node_id: Some(node_id_str.clone()), node_label, stage_id: Some(StageId::new(node_id_str, *visit)), diff --git a/lib/crates/fabro-workflow/src/handler/agent.rs b/lib/crates/fabro-workflow/src/handler/agent.rs index d66ef8ed6..745fc0b8e 100644 --- a/lib/crates/fabro-workflow/src/handler/agent.rs +++ b/lib/crates/fabro-workflow/src/handler/agent.rs @@ -61,6 +61,8 @@ pub trait CodergenBackend: Send + Sync { "one_shot mode not supported by this backend".into(), )) } + + async fn shutdown(&self, _emitter: &Arc) {} } /// The default handler for LLM task nodes. @@ -223,6 +225,12 @@ pub(crate) fn simulate_llm_handler(node: &Node) -> Outcome { #[async_trait] impl Handler for AgentHandler { + async fn shutdown(&self, emitter: &Arc) { + if let Some(backend) = self.backend.as_ref() { + backend.shutdown(emitter).await; + } + } + async fn simulate( &self, node: &Node, @@ -748,15 +756,14 @@ mod tests { ) -> Result { let scope = StageScope::for_handler(context, &node.id); emitter.emit_scoped( - &crate::event::Event::Agent { - stage: node.id.clone(), - visit: scope.visit, - event: fabro_agent::AgentEvent::SessionStarted { - provider: Some("openai".to_string()), - model: Some("gpt-5.4".to_string()), - }, - session_id: Some("session_123".to_string()), - parent_session_id: None, + &crate::event::Event::AgentSessionActivated { + node_id: node.id.clone(), + visit: scope.visit, + session_id: "session_123".to_string(), + thread_id: None, + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + capabilities: vec![fabro_types::SessionCapability::Steer], }, &scope, ); diff --git a/lib/crates/fabro-workflow/src/handler/fan_in.rs b/lib/crates/fabro-workflow/src/handler/fan_in.rs index 65c838731..11a79f7b0 100644 --- a/lib/crates/fabro-workflow/src/handler/fan_in.rs +++ b/lib/crates/fabro-workflow/src/handler/fan_in.rs @@ -29,6 +29,12 @@ impl FanInHandler { #[async_trait] impl Handler for FanInHandler { + async fn shutdown(&self, emitter: &Arc) { + if let Some(backend) = self.backend.as_ref() { + backend.shutdown(emitter).await; + } + } + async fn simulate( &self, node: &Node, diff --git a/lib/crates/fabro-workflow/src/handler/llm/activation_lease.rs b/lib/crates/fabro-workflow/src/handler/llm/activation_lease.rs new file mode 100644 index 000000000..6d9d089a4 --- /dev/null +++ b/lib/crates/fabro-workflow/src/handler/llm/activation_lease.rs @@ -0,0 +1,248 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use fabro_agent::SessionControlHandle; +use fabro_types::{SessionCapability, StageId}; + +use crate::error::Error; +use crate::event::{Emitter, Event}; +use crate::steering_hub::SteeringHub; + +pub struct ActivationLease { + stage_id: StageId, + session_id: String, + hub: Arc, + emitter: Arc, + released: AtomicBool, +} + +pub struct ActivationLeaseOptions { + pub stage_id: StageId, + pub session_id: String, + pub thread_id: Option, + pub provider: Option, + pub model: Option, + pub capabilities: Vec, + pub hub: Arc, + pub emitter: Arc, +} + +impl ActivationLease { + pub fn activate( + options: ActivationLeaseOptions, + handle: &SessionControlHandle, + ) -> Result, Error> { + if !options + .hub + .attach_handle(&options.stage_id, &options.session_id, handle) + { + return Err(Error::Precondition(format!( + "stage {} already has a different active agent session", + options.stage_id + ))); + } + + options.emitter.emit(&Event::AgentSessionActivated { + node_id: options.stage_id.node_id().to_string(), + visit: options.stage_id.visit(), + session_id: options.session_id.clone(), + thread_id: options.thread_id, + provider: options.provider, + model: options.model, + capabilities: options.capabilities, + }); + options.hub.drain_pending_into(&options.stage_id, handle); + + Ok(Arc::new(Self { + stage_id: options.stage_id, + session_id: options.session_id, + hub: options.hub, + emitter: options.emitter, + released: AtomicBool::new(false), + })) + } + + pub fn release(&self) { + if !self.mark_released() { + return; + } + self.hub.detach(&self.stage_id, &self.session_id); + } + + pub fn release_if_queue_empty(&self, handle: &SessionControlHandle) -> bool { + if self.released.load(Ordering::Acquire) { + return true; + } + if !self + .hub + .detach_if_queue_empty(&self.stage_id, &self.session_id, handle) + { + return false; + } + self.mark_released(); + true + } + + fn mark_released(&self) -> bool { + if self + .released + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return false; + } + self.emitter.emit(&Event::AgentSessionDeactivated { + node_id: self.stage_id.node_id().to_string(), + visit: self.stage_id.visit(), + session_id: self.session_id.clone(), + }); + true + } +} + +impl Drop for ActivationLease { + fn drop(&mut self) { + self.release(); + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use fabro_agent::SessionControlHandle; + use fabro_types::{RunId, SteerKind}; + + use super::*; + + fn collect_event_names(emitter: &Arc) -> Arc>> { + let names = Arc::new(Mutex::new(Vec::new())); + let names_for_listener = Arc::clone(&names); + emitter.on_event(move |event| { + names_for_listener + .lock() + .unwrap() + .push(event.event_name().to_string()); + }); + names + } + + fn options( + stage_id: StageId, + session_id: &str, + hub: Arc, + emitter: Arc, + ) -> ActivationLeaseOptions { + ActivationLeaseOptions { + stage_id, + session_id: session_id.to_string(), + thread_id: None, + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + capabilities: vec![SessionCapability::Steer], + hub, + emitter, + } + } + + #[test] + fn activate_emits_activated_before_draining_pending() { + let emitter = Arc::new(Emitter::new(RunId::new())); + let names = collect_event_names(&emitter); + let hub = Arc::new(SteeringHub::new(Arc::clone(&emitter))); + let stage_id = StageId::new("agent", 1); + let handle = SessionControlHandle::new(); + + hub.deliver("queued".to_string(), SteerKind::Interrupt, None); + let _lease = ActivationLease::activate( + options( + stage_id.clone(), + "session-a", + Arc::clone(&hub), + Arc::clone(&emitter), + ), + &handle, + ) + .unwrap(); + + assert_eq!(handle.queue_len(), 1); + assert_eq!(names.lock().unwrap().as_slice(), [ + "agent.steer.buffered", + "agent.session.activated" + ]); + } + + #[test] + fn activate_rejects_mismatched_existing_session() { + let emitter = Arc::new(Emitter::new(RunId::new())); + let names = collect_event_names(&emitter); + let hub = Arc::new(SteeringHub::new(Arc::clone(&emitter))); + let stage_id = StageId::new("agent", 1); + let handle_a = SessionControlHandle::new(); + let handle_b = SessionControlHandle::new(); + + let _lease = ActivationLease::activate( + options( + stage_id.clone(), + "session-a", + Arc::clone(&hub), + Arc::clone(&emitter), + ), + &handle_a, + ) + .unwrap(); + let result = ActivationLease::activate( + options( + stage_id, + "session-b", + Arc::clone(&hub), + Arc::clone(&emitter), + ), + &handle_b, + ); + + assert!(result.is_err()); + assert_eq!(handle_b.queue_len(), 0); + assert_eq!( + names + .lock() + .unwrap() + .iter() + .filter(|name| name.as_str() == "agent.session.activated") + .count(), + 1 + ); + } + + #[test] + fn release_is_idempotent() { + let emitter = Arc::new(Emitter::new(RunId::new())); + let names = collect_event_names(&emitter); + let hub = Arc::new(SteeringHub::new(Arc::clone(&emitter))); + let stage_id = StageId::new("agent", 1); + let handle = SessionControlHandle::new(); + + let lease = ActivationLease::activate( + options( + stage_id, + "session-a", + Arc::clone(&hub), + Arc::clone(&emitter), + ), + &handle, + ) + .unwrap(); + lease.release(); + lease.release(); + + assert_eq!( + names + .lock() + .unwrap() + .iter() + .filter(|name| name.as_str() == "agent.session.deactivated") + .count(), + 1 + ); + } +} diff --git a/lib/crates/fabro-workflow/src/handler/llm/api.rs b/lib/crates/fabro-workflow/src/handler/llm/api.rs index a714b980e..4ffe27cd7 100644 --- a/lib/crates/fabro-workflow/src/handler/llm/api.rs +++ b/lib/crates/fabro-workflow/src/handler/llm/api.rs @@ -13,12 +13,13 @@ use fabro_llm::client::Client; use fabro_llm::types::{Message, Request, TokenCounts}; use fabro_mcp::config::McpServerSettings; use fabro_model::{FallbackTarget, Provider}; -use fabro_types::StageId; +use fabro_types::{SessionCapability, StageId}; use tokio::sync::Mutex as TokioMutex; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use super::super::agent::{CodergenBackend, CodergenResult}; +use super::activation_lease::{ActivationLease, ActivationLeaseOptions}; use crate::context::keys::Fidelity; use crate::context::{Context, WorkflowContext}; use crate::error::Error; @@ -118,6 +119,37 @@ fn classify_agent_error(err: fabro_agent::Error, allow_failover: bool) -> AgentA } } +fn begin_session_lifecycle( + session: &Session, + emitter: &Arc, + parent_session_id: Option, +) { + emitter.emit(&Event::AgentSessionStarted { + session_id: session.id().to_string(), + parent_session_id, + provider: Some(session.provider().to_string()), + model: Some(session.model().to_string()), + }); +} + +fn discard_session( + session: &mut Session, + lease: &mut Option>, + emitter: &Arc, + parent_session_id: Option, +) { + if let Some(lease) = lease.take() { + lease.release(); + } + let session_id = session.id().to_string(); + if session.close() { + emitter.emit(&Event::AgentSessionEnded { + session_id, + parent_session_id, + }); + } +} + fn build_profile(model: &str, provider: Provider) -> Box { match provider { Provider::OpenAi => Box::new(OpenAiProfile::new(model)), @@ -190,6 +222,10 @@ fn spawn_event_forwarder( // Forward non-streaming agent events to pipeline if !event.event.is_streaming_noise() && !matches!(&event.event, AgentEvent::ProcessingEnd) + && !matches!( + &event.event, + AgentEvent::SessionStarted { .. } | AgentEvent::SessionEnded + ) { emitter.emit_scoped( &Event::Agent { @@ -378,22 +414,62 @@ impl AgentApiBackend { Ok(session) } - /// Register `session` with the steering hub under `stage_id` and wire - /// up the completion coordinator. Used both at initial setup and on - /// failover (re-register replaces silently — no re-drain, no event). - fn attach_session_to_hub(&self, session: &mut Session, stage_id: &StageId) { + /// Activate `session` with the steering hub under `stage_id` and wire up + /// the completion coordinator. + fn attach_session_to_hub( + &self, + session: &mut Session, + stage_id: &StageId, + thread_id: Option<&str>, + emitter: &Arc, + ) -> Result, Error> { let handle = session.control_handle(); - self.steering_hub.register(stage_id, &handle); + let lease = ActivationLease::activate( + ActivationLeaseOptions { + stage_id: stage_id.clone(), + session_id: session.id().to_string(), + thread_id: thread_id.map(str::to_string), + provider: Some(session.provider().to_string()), + model: Some(session.model().to_string()), + capabilities: vec![SessionCapability::Steer], + hub: Arc::clone(&self.steering_hub), + emitter: Arc::clone(emitter), + }, + &handle, + )?; session.set_completion_coordinator(Arc::new(SteeringCompletionCoordinator { - hub: Arc::clone(&self.steering_hub), - stage_id: stage_id.clone(), handle, + lease: Mutex::new(Some(Arc::clone(&lease))), })); + Ok(lease) + } + + fn shutdown_cached_sessions(&self, emitter: &Arc) { + let sessions: Vec = self + .sessions + .lock() + .unwrap() + .drain() + .map(|(_, s)| s) + .collect(); + for mut session in sessions { + let session_id = session.id().to_string(); + if session.close() { + emitter.emit(&Event::AgentSessionEnded { + session_id, + parent_session_id: None, + }); + } + } } } #[async_trait] impl CodergenBackend for AgentApiBackend { + async fn shutdown(&self, emitter: &Arc) { + self.shutdown_cached_sessions(emitter); + } + async fn one_shot( &self, node: &Node, @@ -617,29 +693,29 @@ impl CodergenBackend for AgentApiBackend { // Record turn count before processing so we only aggregate new usage. let mut turns_before = session.history().turns().len(); - // Register with the steering hub so HTTP `POST /runs/{id}/steer` - // calls reach this session. The RAII guard below unregisters on - // every exit path (success, error, failover replace). + // Activate with the steering hub after initialization so HTTP + // `POST /runs/{id}/steer` calls reach this session. The activation + // lease is shared with the natural-completion coordinator and is + // released on every exit path. let stage_id = stage_scope.stage_id(); - let _hub_guard = { - let hub = Arc::clone(&self.steering_hub); - let sid = stage_id.clone(); - scopeguard::guard((), move |()| hub.unregister(&sid)) - }; + let mut lease: Option> = None; let allow_failover_primary = !self.fallback_chain.is_empty(); let init_result = if is_reused { Ok(()) } else { + begin_session_lifecycle(&session, emitter, None); match session.initialize().await { Ok(()) => Ok(()), Err(err) => match classify_agent_error(err, allow_failover_primary) { AgentApiErrorDisposition::Cancelled => { bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); return Err(Error::Cancelled); } AgentApiErrorDisposition::Terminal(err) => { bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); return Err(err); } AgentApiErrorDisposition::FailoverEligible(sdk_err) => { @@ -653,7 +729,14 @@ impl CodergenBackend for AgentApiBackend { // process_input failover trigger; otherwise run process_input. let result = match init_result { Ok(()) => { - self.attach_session_to_hub(&mut session, &stage_id); + match self.attach_session_to_hub(&mut session, &stage_id, thread_id, emitter) { + Ok(active_lease) => lease = Some(active_lease), + Err(err) => { + bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); + return Err(err); + } + } session.process_input(prompt).await } Err(err) => Err(err), @@ -665,10 +748,12 @@ impl CodergenBackend for AgentApiBackend { Err(err) => match classify_agent_error(err, allow_failover_primary) { AgentApiErrorDisposition::Cancelled => { bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); return Err(Error::Cancelled); } AgentApiErrorDisposition::Terminal(err) => { bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); return Err(err); } AgentApiErrorDisposition::FailoverEligible(sdk_err) => { @@ -679,6 +764,9 @@ impl CodergenBackend for AgentApiBackend { let mut last_err = Error::Llm(sdk_err); let mut succeeded = false; + bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); + for (index, target) in self.fallback_chain.iter().enumerate() { emitter.emit_scoped( &Event::Failover { @@ -697,9 +785,6 @@ impl CodergenBackend for AgentApiBackend { Err(_) => continue, }; - // Detach the bridge from the failing session before - // refreshing credentials and building a new one. - bridge.abort(); if cancel_token.is_cancelled() { return Err(Error::Cancelled); } @@ -738,23 +823,40 @@ impl CodergenBackend for AgentApiBackend { ); let allow_failover_next = index + 1 < self.fallback_chain.len(); + begin_session_lifecycle(&session, emitter, None); if let Err(err) = session.initialize().await { match classify_agent_error(err, allow_failover_next) { AgentApiErrorDisposition::Cancelled => { bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); return Err(Error::Cancelled); } AgentApiErrorDisposition::Terminal(err) => { bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); return Err(err); } AgentApiErrorDisposition::FailoverEligible(sdk_err) => { last_err = Error::Llm(sdk_err); + bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); continue; } } } - self.attach_session_to_hub(&mut session, &stage_id); + match self.attach_session_to_hub( + &mut session, + &stage_id, + thread_id, + emitter, + ) { + Ok(active_lease) => lease = Some(active_lease), + Err(err) => { + bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); + return Err(err); + } + } match session.process_input(prompt).await { Ok(()) => { succeeded = true; @@ -763,14 +865,18 @@ impl CodergenBackend for AgentApiBackend { Err(err) => match classify_agent_error(err, allow_failover_next) { AgentApiErrorDisposition::Cancelled => { bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); return Err(Error::Cancelled); } AgentApiErrorDisposition::Terminal(err) => { bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); return Err(err); } AgentApiErrorDisposition::FailoverEligible(sdk_err) => { last_err = Error::Llm(sdk_err); + bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); } }, } @@ -781,9 +887,13 @@ impl CodergenBackend for AgentApiBackend { }, }; - // On error, drop the session (don't cache failed state). The bridge's - // `Drop` will abort the spawned task on early return. - result?; + // On error, discard the session (don't cache failed state). The + // bridge's `Drop` will abort the spawned task on early return. + if let Err(err) = result { + bridge.abort(); + discard_session(&mut session, &mut lease, emitter, None); + return Err(err); + } // Aggregate token usage only from new turns (prevents double-counting on // reuse). @@ -825,11 +935,23 @@ impl CodergenBackend for AgentApiBackend { (v, s.last.clone()) }; + if let Some(lease) = lease.take() { + lease.release(); + } + // Cache session back for reuse on success. Detach the bridge first so // the cached session is not left wired to this run's cancel token. if let Some(key) = reuse_key { bridge.abort(); self.sessions.lock().unwrap().insert(key, session); + } else { + let session_id = session.id().to_string(); + if session.close() { + emitter.emit(&Event::AgentSessionEnded { + session_id, + parent_session_id: None, + }); + } } Ok(CodergenResult::Text { @@ -843,37 +965,102 @@ impl CodergenBackend for AgentApiBackend { /// Coordinator that lets the agent loop ask the workflow layer whether to /// keep iterating after a no-tool natural completion. Implements the -/// "close-the-door" pattern: unregister, check the queue, then either -/// break or re-register and report `true` so the loop drains. +/// "close-the-door" pattern: detach only if the queue is empty, otherwise +/// report `true` so the loop drains. struct SteeringCompletionCoordinator { - hub: Arc, - stage_id: StageId, - handle: SessionControlHandle, + handle: SessionControlHandle, + lease: Mutex>>, } impl CompletionCoordinator for SteeringCompletionCoordinator { fn on_natural_completion(&self) -> bool { - // Atomic close-the-door: under the hub's active write lock, check - // the queue. If empty → unregister + emit `detached`, return - // `false` (loop breaks). If non-empty → leave registration in - // place (no event flap), return `true` (loop iterates once more - // and drains). - !self - .hub - .unregister_if_queue_empty(&self.stage_id, &self.handle) + let mut lease = self.lease.lock().expect("activation lease lock poisoned"); + let Some(active_lease) = lease.as_ref() else { + return false; + }; + if active_lease.release_if_queue_empty(&self.handle) { + lease.take(); + false + } else { + true + } } } #[cfg(test)] mod tests { use fabro_agent::subagent::SessionFactory; + use fabro_agent::{AgentProfile, ToolRegistry}; use fabro_auth::{AuthCredential, AuthDetails, VaultCredentialSource}; + use fabro_llm::provider::{ProviderAdapter, StreamEventStream}; use fabro_llm::{Error as LlmError, ProviderErrorDetail, ProviderErrorKind}; use fabro_vault::{SecretType, Vault}; + use futures::stream; use tokio::sync::RwLock as AsyncRwLock; use super::*; + struct ShutdownTestProfile { + registry: ToolRegistry, + } + + impl ShutdownTestProfile { + fn new() -> Self { + Self { + registry: ToolRegistry::new(), + } + } + } + + impl AgentProfile for ShutdownTestProfile { + fn provider(&self) -> Provider { + Provider::OpenAi + } + + fn model(&self) -> &str { + "gpt-5.4" + } + + fn tool_registry(&self) -> &ToolRegistry { + &self.registry + } + + fn tool_registry_mut(&mut self) -> &mut ToolRegistry { + &mut self.registry + } + + fn build_system_prompt( + &self, + _env: &dyn fabro_agent::Sandbox, + _env_context: &fabro_agent::EnvContext, + _memory: &[String], + _user_instructions: Option<&str>, + _skills: &[fabro_agent::Skill], + ) -> String { + "test".to_string() + } + } + + struct ShutdownTestProvider; + + #[async_trait] + impl ProviderAdapter for ShutdownTestProvider { + fn name(&self) -> &str { + "openai" + } + + async fn complete( + &self, + _request: &Request, + ) -> Result { + unreachable!("shutdown test never calls LLM completion") + } + + async fn stream(&self, _request: &Request) -> Result { + Ok(Box::pin(stream::empty())) + } + } + #[test] fn agent_backend_stores_config() { let backend = AgentApiBackend::new_from_env( @@ -1053,6 +1240,56 @@ mod tests { assert_eq!(client.provider_names(), vec!["anthropic"]); } + #[tokio::test] + async fn api_backend_shutdown_closes_cached_sessions_once() { + let backend = AgentApiBackend::new_from_env( + "gpt-5.4".to_string(), + Provider::OpenAi, + Vec::new(), + SteeringHub::for_tests(), + ); + let emitter = Arc::new(Emitter::new(fabro_types::RunId::new())); + let event_names = Arc::new(Mutex::new(Vec::new())); + let event_names_for_listener = Arc::clone(&event_names); + emitter.on_event(move |event| { + event_names_for_listener + .lock() + .unwrap() + .push(event.event_name().to_string()); + }); + + let mut providers = HashMap::new(); + providers.insert( + "openai".to_string(), + Arc::new(ShutdownTestProvider) as Arc, + ); + let client = Client::new(providers, Some("openai".to_string()), Vec::new()); + let session = Session::new( + client, + Arc::new(ShutdownTestProfile::new()), + Arc::new(fabro_agent::LocalSandbox::new( + tempfile::tempdir().unwrap().path().to_path_buf(), + )), + SessionOptions::default(), + None, + ); + begin_session_lifecycle(&session, &emitter, None); + backend + .sessions + .lock() + .unwrap() + .insert("thread-1".to_string(), session); + + backend.shutdown(&emitter).await; + backend.shutdown(&emitter).await; + + assert_eq!(event_names.lock().unwrap().as_slice(), [ + "agent.session.started", + "agent.session.ended" + ]); + assert!(backend.sessions.lock().unwrap().is_empty()); + } + // --- Bridge guard tests --- fn failover_eligible_llm_error() -> LlmError { diff --git a/lib/crates/fabro-workflow/src/handler/llm/cli.rs b/lib/crates/fabro-workflow/src/handler/llm/cli.rs index 90832063a..58fdb3633 100644 --- a/lib/crates/fabro-workflow/src/handler/llm/cli.rs +++ b/lib/crates/fabro-workflow/src/handler/llm/cli.rs @@ -909,6 +909,10 @@ impl CodergenBackend for BackendRouter { .one_shot(node, prompt, system_prompt, emitter, stage_scope) .await } + + async fn shutdown(&self, emitter: &Arc) { + self.api_backend.shutdown(emitter).await; + } } #[cfg(test)] diff --git a/lib/crates/fabro-workflow/src/handler/llm/mod.rs b/lib/crates/fabro-workflow/src/handler/llm/mod.rs index 9e7570560..f210015ff 100644 --- a/lib/crates/fabro-workflow/src/handler/llm/mod.rs +++ b/lib/crates/fabro-workflow/src/handler/llm/mod.rs @@ -1,3 +1,4 @@ +pub mod activation_lease; pub mod api; pub mod cli; pub mod preamble; diff --git a/lib/crates/fabro-workflow/src/handler/mod.rs b/lib/crates/fabro-workflow/src/handler/mod.rs index 8dc13ae6e..671d5c2b7 100644 --- a/lib/crates/fabro-workflow/src/handler/mod.rs +++ b/lib/crates/fabro-workflow/src/handler/mod.rs @@ -22,6 +22,7 @@ use fabro_interview::Interviewer; use crate::context::Context; use crate::error::Error; +use crate::event::Emitter; use crate::outcome::{Outcome, OutcomeExt}; pub use crate::services::{EngineServices, RunServices}; @@ -55,6 +56,8 @@ pub trait Handler: Send + Sync { fn should_retry(&self, err: &Error) -> bool { err.is_retryable() } + + async fn shutdown(&self, _emitter: &Arc) {} } /// Extract a human-readable message from a panic payload. @@ -130,6 +133,13 @@ impl HandlerRegistry { // 3. Default self.default_handler.as_ref() } + + pub async fn shutdown_all(&self, emitter: &Arc) { + self.default_handler.shutdown(emitter).await; + for handler in self.handlers.values() { + handler.shutdown(emitter).await; + } + } } /// Build a [`HandlerRegistry`] with all built-in handler types registered. diff --git a/lib/crates/fabro-workflow/src/handler/prompt.rs b/lib/crates/fabro-workflow/src/handler/prompt.rs index 50dc04e06..7f6e9b691 100644 --- a/lib/crates/fabro-workflow/src/handler/prompt.rs +++ b/lib/crates/fabro-workflow/src/handler/prompt.rs @@ -1,4 +1,5 @@ use std::path::Path; +use std::sync::Arc; use async_trait::async_trait; use fabro_graphviz::graph::{Graph, Node}; @@ -10,7 +11,7 @@ use super::agent::{ use super::{EngineServices, Handler}; use crate::context::{Context, WorkflowContext, keys}; use crate::error::Error; -use crate::event::{Event, StageScope}; +use crate::event::{Emitter, Event, StageScope}; use crate::outcome::Outcome; /// Handler for single-shot LLM calls (no tools, no agent loop). @@ -27,6 +28,12 @@ impl PromptHandler { #[async_trait] impl Handler for PromptHandler { + async fn shutdown(&self, emitter: &Arc) { + if let Some(backend) = self.backend.as_ref() { + backend.shutdown(emitter).await; + } + } + async fn simulate( &self, node: &Node, diff --git a/lib/crates/fabro-workflow/src/operations/fork.rs b/lib/crates/fabro-workflow/src/operations/fork.rs index f2a88c899..66f6baa67 100644 --- a/lib/crates/fabro-workflow/src/operations/fork.rs +++ b/lib/crates/fabro-workflow/src/operations/fork.rs @@ -248,7 +248,7 @@ fn replay_event_for_fork_projection(body: &EventBody) -> bool { | EventBody::InterviewCompleted(_) | EventBody::InterviewTimeout(_) | EventBody::InterviewInterrupted(_) - | EventBody::AgentSessionStarted(_) + | EventBody::AgentSessionActivated(_) | EventBody::AgentCliStarted(_) | EventBody::AgentCliCancelled(_) | EventBody::AgentCliTimedOut(_) @@ -314,6 +314,28 @@ mod tests { ) } + #[test] + fn fork_replay_keeps_stage_scoped_session_activation_only() { + assert!(replay_event_for_fork_projection( + &EventBody::AgentSessionActivated(fabro_types::run_event::AgentSessionActivatedProps { + thread_id: None, + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + capabilities: vec![fabro_types::SessionCapability::Steer], + visit: 1, + }) + )); + assert!(!replay_event_for_fork_projection( + &EventBody::AgentSessionStarted(fabro_types::run_event::AgentSessionStartedProps { + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + }) + )); + assert!(!replay_event_for_fork_projection( + &EventBody::AgentSessionEnded(fabro_types::run_event::AgentSessionEndedProps {}) + )); + } + #[tokio::test] async fn fork_persists_historical_node_projection_through_target_checkpoint() { let store = test_store(); diff --git a/lib/crates/fabro-workflow/src/pipeline/execute.rs b/lib/crates/fabro-workflow/src/pipeline/execute.rs index 57ed86cfc..d1788456d 100644 --- a/lib/crates/fabro-workflow/src/pipeline/execute.rs +++ b/lib/crates/fabro-workflow/src/pipeline/execute.rs @@ -291,6 +291,8 @@ pub async fn execute(init: Initialized) -> Executed { Err(e) => (Err(Error::engine(e.to_string())), initial_context), }; + engine.registry.shutdown_all(&engine.run.emitter).await; + let duration_ms = crate::millis_u64(start.elapsed()); Executed { diff --git a/lib/crates/fabro-workflow/src/steering_hub.rs b/lib/crates/fabro-workflow/src/steering_hub.rs index cb3764f27..db9b33ebd 100644 --- a/lib/crates/fabro-workflow/src/steering_hub.rs +++ b/lib/crates/fabro-workflow/src/steering_hub.rs @@ -2,7 +2,7 @@ //! `Session`s. The hub owns: //! //! - A map of currently steerable API-mode sessions, keyed by `StageId` → -//! `SessionControlHandle`. +//! active `(session_id, SessionControlHandle)` entries. //! - A bounded run-wide pending buffer for steers that arrive when no session //! is registered (between stages, before the first agent stage, or after a //! session ends but before the next registers). @@ -13,7 +13,7 @@ //! - `pending` is `std::sync::Mutex` taken under the active read lock. //! - All methods are sync — no `.await` while holding any lock — so the //! `CompletionCoordinator::on_natural_completion` close-the-door dance can -//! call `unregister(...)` synchronously from the agent loop. +//! call `detach_if_queue_empty(...)` synchronously from the agent loop. use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex, RwLock}; @@ -38,12 +38,18 @@ struct PendingSteer { actor: Option, } +#[derive(Clone)] +struct ActiveEntry { + handle: SessionControlHandle, + session_id: String, +} + #[allow( clippy::module_name_repetitions, reason = "external callers refer to it as SteeringHub" )] pub struct SteeringHub { - active: RwLock>, + active: RwLock>, pending: Mutex>, emitter: Arc, } @@ -80,87 +86,83 @@ impl SteeringHub { self.active.read().expect("active lock poisoned").len() } - /// Register an API-mode session as steerable for this stage. If no - /// entry existed for `stage_id`, emits `agent.steering.attached` and - /// drains pending into the new handle as `Append`-kind messages. If - /// an entry already existed (e.g. failover replaced the underlying - /// session), the handle is overwritten silently — no drain, no event. - pub fn register(&self, stage_id: &StageId, handle: &SessionControlHandle) { - let was_new = { - let mut active = self.active.write().expect("active lock poisoned"); - let was_new = !active.contains_key(stage_id); - active.insert(stage_id.clone(), handle.clone()); - was_new - }; - if was_new { - // Emit attached *before* draining so any `agent.steer.dropped` - // events from cap-evictions during the drain follow the - // attached event in the stream (UI consumers can attribute - // drops to a known session). - self.emitter.emit(&Event::AgentSteeringAttached { - node_id: stage_id.node_id().to_string(), - visit: stage_id.visit(), - }); - let pending: Vec = { - let mut pending = self.pending.lock().expect("pending lock poisoned"); - pending.drain(..).collect() - }; - for item in pending { - // Buffered steers always flush as Append — the original - // Interrupt semantics no longer make sense once the round - // has rolled over. - Self::enqueue_into_session_queue( - handle, - (item.text, SteerKind::Append, item.actor), - &self.emitter, - Some(stage_id), - ); + /// Attach an API-mode session as steerable for this stage. Returns + /// `false` when a different session is already active for the stage. + pub fn attach_handle( + &self, + stage_id: &StageId, + session_id: &str, + handle: &SessionControlHandle, + ) -> bool { + let mut active = self.active.write().expect("active lock poisoned"); + match active.get_mut(stage_id) { + Some(entry) if entry.session_id != session_id => false, + Some(entry) => { + entry.handle = handle.clone(); + true + } + None => { + active.insert(stage_id.clone(), ActiveEntry { + handle: handle.clone(), + session_id: session_id.to_string(), + }); + true } } } - /// Unregister the session previously registered for this stage. Emits - /// `agent.steering.detached` only when an entry was actually removed - /// (idempotent — safe to call multiple times from RAII guards). - pub fn unregister(&self, stage_id: &StageId) { - let removed = { - let mut active = self.active.write().expect("active lock poisoned"); - active.remove(stage_id).is_some() + /// Drain pending run-wide steers into `handle` as `Append` messages. + pub fn drain_pending_into(&self, stage_id: &StageId, handle: &SessionControlHandle) { + let pending: Vec = { + let mut pending = self.pending.lock().expect("pending lock poisoned"); + pending.drain(..).collect() }; - if removed { - self.emitter.emit(&Event::AgentSteeringDetached { - node_id: stage_id.node_id().to_string(), - visit: stage_id.visit(), - }); + for item in pending { + // Buffered steers always flush as Append — the original + // Interrupt semantics no longer make sense once the round has + // rolled over. + Self::enqueue_into_session_queue( + handle, + (item.text, SteerKind::Append, item.actor), + &self.emitter, + Some(stage_id), + ); } } + /// Detach the session for this stage. Stale session ids are ignored. + pub fn detach(&self, stage_id: &StageId, session_id: &str) -> bool { + let mut active = self.active.write().expect("active lock poisoned"); + let Some(entry) = active.get(stage_id) else { + return false; + }; + if entry.session_id != session_id { + return false; + } + active.remove(stage_id); + true + } + /// Atomic close-the-door check used by the agent loop's natural- - /// completion path. Under the `active` write lock: if `handle`'s - /// queue is empty, remove the stage and return `true` (loop should - /// break — emits `detached`). If the queue is non-empty, leave the - /// registration intact and return `false` (loop should iterate once - /// more — no event emitted, so no detach/attach flap). - pub fn unregister_if_queue_empty( + /// completion path. Under the `active` write lock: if `handle`'s queue + /// is empty and the active session id matches, remove the stage and + /// return `true`. If the queue is non-empty, leave the registration + /// intact and return `false`. + pub fn detach_if_queue_empty( &self, stage_id: &StageId, + session_id: &str, handle: &SessionControlHandle, ) -> bool { - let removed = { - let mut active = self.active.write().expect("active lock poisoned"); - if handle.queue_is_empty() { - active.remove(stage_id).is_some() - } else { - false - } + let mut active = self.active.write().expect("active lock poisoned"); + let Some(entry) = active.get(stage_id) else { + return false; }; - if removed { - self.emitter.emit(&Event::AgentSteeringDetached { - node_id: stage_id.node_id().to_string(), - visit: stage_id.visit(), - }); + if entry.session_id != session_id || !handle.queue_is_empty() { + return false; } - removed + active.remove(stage_id); + true } /// Deliver a steer from the HTTP control plane. Broadcasts to every @@ -201,9 +203,9 @@ impl SteeringHub { } // Broadcast to every active session. - for (stage_id, handle) in active.iter() { + for (stage_id, entry) in active.iter() { Self::enqueue_into_session_queue( - handle, + &entry.handle, (text.clone(), kind, actor.clone()), &self.emitter, Some(stage_id), @@ -297,12 +299,12 @@ mod tests { fn unregister_is_idempotent() { let hub = SteeringHub::for_tests(); let stage = StageId::new("agent-node", 1); - hub.unregister(&stage); - hub.unregister(&stage); + hub.detach(&stage, "session-a"); + hub.detach(&stage, "session-a"); } #[test] - fn register_drains_pending_into_first_session() { + fn attach_and_drain_pending_into_first_session() { let hub = SteeringHub::for_tests(); hub.deliver("queued1".into(), SteerKind::Append, None); hub.deliver("queued2".into(), SteerKind::Interrupt, None); @@ -310,7 +312,8 @@ mod tests { let stage = StageId::new("agent-node", 1); let handle = SessionControlHandle::new(); - hub.register(&stage, &handle); + assert!(hub.attach_handle(&stage, "session-a", &handle)); + hub.drain_pending_into(&stage, &handle); assert_eq!(handle.queue_len(), 2); assert_eq!(hub.pending_len(), 0); @@ -324,8 +327,8 @@ mod tests { let stage_b = StageId::new("b", 1); let handle_a = SessionControlHandle::new(); let handle_b = SessionControlHandle::new(); - hub.register(&stage_a, &handle_a); - hub.register(&stage_b, &handle_b); + assert!(hub.attach_handle(&stage_a, "session-a", &handle_a)); + assert!(hub.attach_handle(&stage_b, "session-b", &handle_b)); hub.deliver("hello".into(), SteerKind::Append, None); @@ -335,27 +338,63 @@ mod tests { } #[test] - fn re_register_same_stage_does_not_redrain() { + fn attach_rejects_different_session_for_same_stage() { let hub = SteeringHub::for_tests(); let stage = StageId::new("a", 1); let handle1 = SessionControlHandle::new(); - hub.register(&stage.clone(), &handle1); + assert!(hub.attach_handle(&stage, "session-a", &handle1)); hub.deliver("x".into(), SteerKind::Append, None); assert_eq!(handle1.queue_len(), 1); - // Replace handle (failover) — must not redrain pending or emit - // attached again. let handle2 = SessionControlHandle::new(); - hub.register(&stage, &handle2); + assert!(!hub.attach_handle(&stage, "session-b", &handle2)); assert_eq!(handle2.queue_len(), 0); } + #[test] + fn stale_detach_does_not_remove_active_session() { + let hub = SteeringHub::for_tests(); + let stage = StageId::new("a", 1); + let handle = SessionControlHandle::new(); + assert!(hub.attach_handle(&stage, "session-a", &handle)); + + assert!(!hub.detach(&stage, "session-b")); + hub.deliver("still-active".into(), SteerKind::Append, None); + + assert_eq!(handle.queue_len(), 1); + assert_eq!(hub.active_count(), 1); + } + + #[test] + fn detach_if_queue_empty_respects_session_id_and_queue_state() { + let hub = SteeringHub::for_tests(); + let stage = StageId::new("a", 1); + let handle = SessionControlHandle::new(); + assert!(hub.attach_handle(&stage, "session-a", &handle)); + + assert!(!hub.detach_if_queue_empty(&stage, "session-b", &handle)); + hub.deliver("queued".into(), SteerKind::Append, None); + assert!(!hub.detach_if_queue_empty(&stage, "session-a", &handle)); + assert_eq!(hub.active_count(), 1); + } + + #[test] + fn detach_if_queue_empty_removes_matching_empty_session() { + let hub = SteeringHub::for_tests(); + let stage = StageId::new("a", 1); + let handle = SessionControlHandle::new(); + assert!(hub.attach_handle(&stage, "session-a", &handle)); + + assert!(hub.detach_if_queue_empty(&stage, "session-a", &handle)); + assert_eq!(hub.active_count(), 0); + } + #[test] fn per_session_queue_evicts_oldest_at_cap() { let hub = SteeringHub::for_tests(); let stage = StageId::new("a", 1); let handle = SessionControlHandle::new(); - hub.register(&stage, &handle); + assert!(hub.attach_handle(&stage, "session-a", &handle)); for i in 0..(super::PER_SESSION_QUEUE_CAP + 5) { hub.deliver(format!("m{i}"), SteerKind::Append, None);