fix: Capture session ID from stream headers rather than final event

crates/atuin-ai/src/driver.rs

@@ -1033,6 +1033,9 @@ async fn run_stream_bridge(
                 StreamControl::Done { session_id } => Some(Event::StreamDone { session_id }),
                 StreamControl::Error(msg) => Some(Event::StreamError(msg)),
             },
+            Ok(StreamFrame::SessionIdentity(session_id)) => {
+                Some(Event::SessionIdReceived(session_id))
+            }
             Err(e) => Some(Event::StreamError(e.to_string())),
         };
 

crates/atuin-ai/src/fsm/events.rs

@@ -22,6 +22,7 @@ pub(crate) enum Event {
     InterruptTools,
 
     // ─── Stream lifecycle ───────────────────────────────────────
+    SessionIdReceived(String),
     /// Stream connection established, first frame received.
     StreamStarted,
     /// Received a chunk of streamed text content.
@@ -43,14 +44,19 @@ pub(crate) enum Event {
     /// Stream status changed (e.g. "thinking", "searching").
     StreamStatusChanged(String),
     /// Stream ended normally.
-    StreamDone { session_id: String },
+    StreamDone {
+        session_id: String,
+    },
     /// Stream encountered an error.
     StreamError(String),
 
     // ─── Suggest command (terminal tool call) ───────────────────
     /// The suggest_command tool call acts as a stream terminal event.
     /// This is the server signaling "turn complete, here's the command."
-    SuggestCommand { id: String, input: Value },
+    SuggestCommand {
+        id: String,
+        input: Value,
+    },
 
     // ─── Tool lifecycle ─────────────────────────────────────────
     /// Permission resolver completed for a tool.
@@ -79,9 +85,14 @@ pub(crate) enum Event {
 
     // ─── Timers ─────────────────────────────────────────────────
     /// Confirmation timeout expired.
-    ConfirmationTimeout { timeout_id: u64 },
+    ConfirmationTimeout {
+        timeout_id: u64,
+    },
     /// Shell tool execution timeout expired.
-    ToolExecutionTimeout { timeout_id: u64, tool_id: String },
+    ToolExecutionTimeout {
+        timeout_id: u64,
+        tool_id: String,
+    },
 
     // ─── Session management ─────────────────────────────────────
     /// User ran /new to start a fresh session.
@@ -91,7 +102,10 @@ pub(crate) enum Event {
     /// User submitted a slash command (other than /new).
     /// The driver resolves known commands (like /help) and passes the
     /// rendered content; the FSM just pushes an OOB event.
-    SlashCommand { command: String, content: String },
+    SlashCommand {
+        command: String,
+        content: String,
+    },
 
     // ─── Skills ────────────────────────────────────────────────
     /// User invoked a skill via /skill-name. FSM emits a LoadSkill

crates/atuin-ai/src/fsm/mod.rs

@@ -199,6 +199,15 @@ impl AgentFsm {
 
     /// Handle an event, returning effects to execute.
     pub fn handle(&mut self, event: Event) -> Vec<Effect> {
+        // From all states: if the session ID arrives and isn't set, set it, then continue as normal.
+        // This event fires from the stream response headers, rather than having to wait until the end
+        // of a turn when StreamDone arrives, which can be lost for cancelled turns.
+        if let Event::SessionIdReceived(session_id) = &event {
+            self.ctx
+                .session_id
+                .get_or_insert_with(|| session_id.clone());
+        }
+
         match (&self.state, event) {
             // ================================================================
             // Idle state

crates/atuin-ai/src/stream.rs

@@ -47,6 +47,7 @@ pub(crate) enum StreamContent {
 /// A frame from the SSE stream, classified as control or content.
 #[derive(Debug, Clone)]
 pub(crate) enum StreamFrame {
+    SessionIdentity(String),
     Content(StreamContent),
     Control(StreamControl),
 }
@@ -196,6 +197,14 @@ pub(crate) fn create_chat_stream(
             return;
         }
 
+        if let Some(sess_id) = response
+            .headers()
+            .get("x-atuin-ai-session-id")
+            .and_then(|v| v.to_str().ok())
+            .filter(|s| !s.is_empty()) {
+                yield Ok(StreamFrame::SessionIdentity(sess_id.to_string()));
+            }
+
         let byte_stream = response.bytes_stream();
         let mut stream = byte_stream.eventsource();