Fix Codex transcript streaming and images

server/src/session/claude/translate.rs

@@ -1148,23 +1148,13 @@ pub(in crate::session) fn save_image(session_dir: &Path, part: &Value) -> Option
     let bytes = base64::engine::general_purpose::STANDARD
         .decode(data)
         .ok()?;
-    // Screenshots are the overwhelming case and they are PNG; an unrecognized
-    // type is more likely a dialect change than a JPEG.
-    let extension = source
+    // Screenshots are the overwhelming case and they are PNG; an unrecognized type is more likely
+    // a dialect change than a JPEG. `store_image` owns that fallback.
+    let media_type = source
         .get("media_type")
         .and_then(Value::as_str)
-        .and_then(crate::media::extension_for)
-        .unwrap_or("png");
-    let name = format!("{}.{extension}", super::super::random_hex());
-    let dir = session_dir.join("files");
-    if let Err(err) = wg_app_link::private::create_dir(&dir)
-        .map_err(std::io::Error::other)
-        .and_then(|()| std::fs::write(dir.join(&name), bytes))
-    {
-        tracing::error!("couldn't save produced image: {err}");
-        return None;
-    }
-    Some(name)
+        .unwrap_or("image/png");
+    super::super::driver::store_image(session_dir, media_type, &bytes)
 }
 
 #[cfg(test)]

server/src/session/codex.rs

@@ -18,7 +18,9 @@ use serde_json::{Value, json};
 use tokio::io::AsyncWriteExt;
 use tokio::sync::mpsc;
 
-use super::driver::{AttachmentRef, Driver, Event, EventSink, SessionStatus, Unqueued};
+use super::driver::{
+    AttachmentRef, Driver, Event, EventSink, SessionStatus, Unqueued, store_image,
+};
 use super::process;
 use super::transport::{Launch, Streams, Transport};
 use crate::config::{ProviderConfig, SessionConfig};
@@ -459,10 +461,15 @@ fn input_for(inner: &Inner, message: &Waiting) -> Result<Vec<Value>> {
     let mut input = Vec::new();
     for attachment in &message.attachments {
         let path = attachment_path(&inner.session_dir, attachment)?;
-        if crate::media::media_type_for(attachment).is_some()
-            && matches!(inner.transport, Transport::Here)
-        {
-            input.push(json!({"type": "localImage", "path": path}));
+        if let Some(media_type) = crate::media::media_type_for(attachment) {
+            if matches!(inner.transport, Transport::Here) {
+                input.push(json!({"type": "localImage", "path": path}));
+            } else {
+                // The upload is on the server, not on the machine reached over ssh. Inline image
+                // input carries those bytes across the app-server connection just as Claude's
+                // image block does; naming the server path left Codex unable to see it.
+                input.push(inline_image(&path, media_type)?);
+            }
         } else {
             if !text.is_empty() {
                 text.push_str("\n\n");
@@ -476,6 +483,17 @@ fn input_for(inner: &Inner, message: &Waiting) -> Result<Vec<Value>> {
     Ok(input)
 }
 
+fn inline_image(path: &Path, media_type: &str) -> Result<Value> {
+    use base64::Engine;
+    let bytes =
+        std::fs::read(path).with_context(|| format!("read attachment {}", path.display()))?;
+    let data = base64::engine::general_purpose::STANDARD.encode(bytes);
+    Ok(json!({
+        "type": "image",
+        "url": format!("data:{media_type};base64,{data}")
+    }))
+}
+
 fn dispatch_waiting(inner: &Arc<Inner>) {
     let Some(thread_id) = read_thread(&inner.session_dir) else {
         return;
@@ -560,8 +578,14 @@ async fn follow(inner: Arc<Inner>, mut record: process::Record, mut offset: u64)
     let stdout = inner.session_dir.join(STDOUT_LOG);
     let stderr = inner.session_dir.join(STDERR_LOG);
     let mut translator = Translator::default();
+    // `offset` is the durable boundary after the last complete record. `read_at` may move beyond
+    // it while app-server is still writing one record. Image-bearing tool results can be several
+    // megabytes long, and rereading their incomplete prefix every 50 ms made arrival over ssh
+    // quadratic in time and allocation.
+    let mut read_at = offset;
+    let mut pending = Vec::new();
     while inner.reading.load(Ordering::SeqCst) {
-        let (bytes, _) = match process::read_from(&stdout, offset) {
+        let (bytes, next) = match process::read_from(&stdout, read_at) {
             Ok(read) => read,
             Err(err) => {
                 let _ = inner.sink.send(Event::Error {
@@ -570,12 +594,14 @@ async fn follow(inner: Arc<Inner>, mut record: process::Record, mut offset: u64)
                 return;
             }
         };
-        let complete = bytes
+        read_at = next;
+        pending.extend_from_slice(&bytes);
+        let complete = pending
             .iter()
             .rposition(|byte| *byte == b'\n')
             .map(|at| at + 1)
             .unwrap_or(0);
-        for line in String::from_utf8_lossy(&bytes[..complete]).lines() {
+        for line in String::from_utf8_lossy(&pending[..complete]).lines() {
             let Ok(value) = serde_json::from_str::<Value>(line) else {
                 tracing::warn!(
                     "unparseable Codex JSONL line: {}",
@@ -587,6 +613,7 @@ async fn follow(inner: Arc<Inner>, mut record: process::Record, mut offset: u64)
         }
         if complete > 0 {
             offset += complete as u64;
+            pending.drain(..complete);
             record.detail = process::Detail::Stdio {
                 stdout_read: offset,
             };
@@ -667,6 +694,9 @@ fn handle_line(inner: &Arc<Inner>, translator: &mut Translator, line: &Value) {
             if item.get("type").and_then(Value::as_str) == Some("userMessage") {
                 announce_user(inner, item);
             }
+            for event in image_events(inner, item) {
+                let _ = inner.sink.send(event);
+            }
         }
         Some("turn/completed") => {
             let mut state = inner.state.lock().unwrap();
@@ -686,6 +716,123 @@ fn handle_line(inner: &Arc<Inner>, translator: &mut Translator, line: &Value) {
     }
 }
 
+/// Images embedded in a structured tool result, copied into the session before the translator's
+/// `ToolEnd` is emitted so they stay attached to that call in transcript order.
+fn image_events(inner: &Inner, item: &Value) -> Vec<Event> {
+    let Some(id) = item.get("id").and_then(Value::as_str) else {
+        return Vec::new();
+    };
+    let mut images = Vec::new();
+    match item.get("type").and_then(Value::as_str) {
+        Some("dynamicToolCall") => save_data_images(
+            &inner.session_dir,
+            item.get("contentItems").and_then(Value::as_array),
+            "inputImage",
+            "imageUrl",
+            &mut images,
+        ),
+        Some("mcpToolCall") => {
+            if let Some(parts) = item.pointer("/result/content").and_then(Value::as_array) {
+                for part in parts {
+                    if part.get("type").and_then(Value::as_str) != Some("image") {
+                        continue;
+                    }
+                    if let (Some(data), Some(media_type)) = (
+                        part.get("data").and_then(Value::as_str),
+                        part.get("mimeType")
+                            .or_else(|| part.get("mime_type"))
+                            .and_then(Value::as_str),
+                    ) && let Some(image) =
+                        save_base64_image(&inner.session_dir, media_type, data)
+                    {
+                        images.push(image);
+                    }
+                }
+            }
+        }
+        Some("functionCallOutput") => save_data_images(
+            &inner.session_dir,
+            item.get("output").and_then(Value::as_array),
+            "input_image",
+            "image_url",
+            &mut images,
+        ),
+        Some("imageView") => {
+            if let Some(path) = item.get("path").and_then(Value::as_str)
+                && let Some(image) = save_viewed_image(inner, path)
+            {
+                images.push(image);
+            }
+        }
+        _ => {}
+    }
+    images
+        .into_iter()
+        .map(|image| Event::Image {
+            image,
+            about: Some(id.to_string()),
+        })
+        .collect()
+}
+
+fn save_data_images(
+    session_dir: &Path,
+    parts: Option<&Vec<Value>>,
+    image_kind: &str,
+    url_field: &str,
+    images: &mut Vec<String>,
+) {
+    for part in parts.into_iter().flatten() {
+        if part.get("type").and_then(Value::as_str) != Some(image_kind) {
+            continue;
+        }
+        let Some(url) = part.get(url_field).and_then(Value::as_str) else {
+            continue;
+        };
+        let Some((header, data)) = url.split_once(',') else {
+            continue;
+        };
+        let Some(media_type) = header
+            .strip_prefix("data:")
+            .and_then(|header| header.strip_suffix(";base64"))
+        else {
+            continue;
+        };
+        if let Some(image) = save_base64_image(session_dir, media_type, data) {
+            images.push(image);
+        }
+    }
+}
+
+fn save_base64_image(session_dir: &Path, media_type: &str, data: &str) -> Option<String> {
+    use base64::Engine;
+    let bytes = base64::engine::general_purpose::STANDARD
+        .decode(data)
+        .ok()?;
+    store_image(session_dir, media_type, &bytes)
+}
+
+fn save_viewed_image(inner: &Inner, path: &str) -> Option<String> {
+    let media_type = crate::media::media_type_for(path).unwrap_or("image/png");
+    let bytes = match &inner.transport {
+        Transport::Here => std::fs::read(path).map_err(anyhow::Error::from),
+        Transport::Ssh { .. } => tokio::task::block_in_place(|| {
+            inner.transport.capture_bytes_blocking(&Launch::new(
+                "cat",
+                vec![path.to_string()],
+                None,
+            ))
+        }),
+    };
+    match bytes {
+        Ok(bytes) => store_image(&inner.session_dir, media_type, &bytes),
+        Err(err) => {
+            tracing::error!("couldn't save image Codex read from {path}: {err:#}");
+            None
+        }
+    }
+}
+
 fn handle_response(inner: &Arc<Inner>, line: &Value) {
     let Some(id) = line.get("id").and_then(Value::as_str) else {
         return;
@@ -707,6 +854,9 @@ fn handle_response(inner: &Arc<Inner>, line: &Value) {
         let method = match read_thread(&inner.session_dir) {
             Some(thread) => {
                 params["threadId"] = Value::String(thread);
+                // This app already owns and pages its common transcript. Hydrating the complete
+                // Codex history here sends it a second time, including every base64 screenshot.
+                params["excludeTurns"] = Value::Bool(true);
                 "thread/resume"
             }
             None => "thread/start",
@@ -918,6 +1068,42 @@ done
 printf '%s\n' "$state"
 "#;
 
+const READ_CONTEXT_SCRIPT: &str = r#"
+for f in "$HOME"/.codex/sessions/*/*/*/rollout-*-${1}.jsonl; do
+  [ -f "$f" ] || continue
+  grep '"type":"token_count"' "$f" | tail -1
+  exit 0
+done
+"#;
+
+/// The last measured prompt size from Codex's own rollout, for a session whose common transcript
+/// predates context events. This is the same `last.inputTokens` app-server reports live, under the
+/// rollout writer's snake-case names.
+pub async fn context_of(transport: &Transport, id: &str) -> Option<u64> {
+    if !valid_thread_id(id) {
+        return None;
+    }
+    let launch = Launch::new(
+        "sh",
+        vec![
+            "-c".to_string(),
+            READ_CONTEXT_SCRIPT.to_string(),
+            "sh".to_string(),
+            id.to_string(),
+        ],
+        None,
+    );
+    let line = transport.capture(&launch).await.ok()?;
+    context_from_rollout(&line)
+}
+
+fn context_from_rollout(line: &str) -> Option<u64> {
+    serde_json::from_str::<Value>(line)
+        .ok()?
+        .pointer("/payload/info/last_token_usage/input_tokens")?
+        .as_u64()
+}
+
 /// Removes the rollout whose suffix is this thread id.
 pub async fn delete_transcript(transport: &Transport, id: &str) -> Result<()> {
     if !valid_thread_id(id) {
@@ -964,6 +1150,39 @@ mod tests {
         assert_eq!(state.waiting[0].id, "q1");
     }
 
+    #[test]
+    fn a_structured_image_is_saved_outside_the_transcript() {
+        let dir = tempfile::tempdir().expect("tempdir");
+        let parts = vec![json!({
+            "type": "inputImage",
+            "imageUrl": "data:image/png;base64,aGVsbG8="
+        })];
+        let mut images = Vec::new();
+        save_data_images(
+            dir.path(),
+            Some(&parts),
+            "inputImage",
+            "imageUrl",
+            &mut images,
+        );
+        assert_eq!(images.len(), 1);
+        assert_eq!(
+            std::fs::read(dir.path().join("files").join(&images[0])).expect("saved image"),
+            b"hello"
+        );
+    }
+
+    #[test]
+    fn an_inline_image_carries_its_bytes_to_a_remote_codex() {
+        let dir = tempfile::tempdir().expect("tempdir");
+        let path = dir.path().join("shot.png");
+        std::fs::write(&path, b"hello").expect("image");
+        assert_eq!(
+            inline_image(&path, "image/png").expect("inline image"),
+            json!({"type": "image", "url": "data:image/png;base64,aGVsbG8="})
+        );
+    }
+
     #[test]
     fn transcript_delete_resolves_only_the_named_codex_rollout() {
         use std::process::Command;
@@ -985,4 +1204,15 @@ mod tests {
         assert!(other.exists());
         assert!(!valid_thread_id("../../something"));
     }
+
+    #[test]
+    fn context_is_read_from_the_last_codex_model_call() {
+        assert_eq!(
+            context_from_rollout(
+                r#"{"type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":118866,"cached_input_tokens":118144,"output_tokens":37,"total_tokens":118903}}}}"#
+            ),
+            Some(118_866)
+        );
+        assert_eq!(context_from_rollout(""), None);
+    }
 }

server/src/session/codex/translate.rs

@@ -4,8 +4,6 @@
 //! therefore match only the records that have a useful common equivalent and
 //! ignore the rest; an added Codex item must not make a live session go deaf.
 
-use std::collections::HashSet;
-
 use serde_json::{Value, json};
 
 use super::super::driver::{Event, SessionStatus};
@@ -15,8 +13,12 @@ pub(super) struct Translator {
     pub(super) thread_id: Option<String>,
     completed: bool,
     limited: bool,
-    streamed_messages: HashSet<String>,
-    pending_usage: Option<u64>,
+    pending_usage: Option<Usage>,
+}
+
+struct Usage {
+    tokens: u64,
+    context: Option<u64>,
 }
 
 impl Translator {
@@ -43,23 +45,23 @@ impl Translator {
                 self.completed = false;
                 self.limited = false;
                 self.pending_usage = None;
-                self.streamed_messages.clear();
                 vec![Event::Status {
                     state: SessionStatus::Running,
                 }]
             }
             Some("item.started") | Some("item/started") => start_item(&body["item"]),
             Some("item.updated") => update_item(&line["item"]),
-            Some("item.completed") | Some("item/completed") => {
-                complete_item(&body["item"], &self.streamed_messages)
-            }
+            // The old `codex exec --json` dialect reports only the completed message. App-server
+            // reports every message through durable delta notifications and its completed copy
+            // must always be skipped. That rule cannot live in an in-memory set: after a backend
+            // restart the deltas are behind the persisted log cursor while the completion is not,
+            // which used to append the whole message again after its already-recorded prefix.
+            Some("item.completed") => complete_item(&body["item"], true),
+            Some("item/completed") => complete_item(&body["item"], false),
             Some("item/agentMessage/delta") => {
                 let Some(delta) = body.get("delta").and_then(Value::as_str) else {
                     return Vec::new();
                 };
-                if let Some(id) = body.get("itemId").and_then(Value::as_str) {
-                    self.streamed_messages.insert(id.to_string());
-                }
                 vec![Event::AssistantText {
                     delta: delta.to_string(),
                 }]
@@ -77,18 +79,24 @@ impl Translator {
                 }]
             }
             Some("thread/tokenUsage/updated") => {
-                self.pending_usage = body
-                    .pointer("/tokenUsage/last/totalTokens")
-                    .and_then(Value::as_u64);
+                let last = &body["tokenUsage"]["last"];
+                self.pending_usage =
+                    last.get("totalTokens")
+                        .and_then(Value::as_u64)
+                        .map(|tokens| Usage {
+                            tokens,
+                            // Cached input is a subset of this figure, not an additional count.
+                            context: last.get("inputTokens").and_then(Value::as_u64),
+                        });
                 Vec::new()
             }
             Some("turn.completed") | Some("turn/completed") => {
                 self.completed = true;
                 let mut events = Vec::new();
-                if let Some(tokens) = self.pending_usage.take() {
+                if let Some(usage) = self.pending_usage.take() {
                     events.push(Event::UsageDelta {
-                        tokens,
-                        context: None,
+                        tokens: usage.tokens,
+                        context: usage.context,
                     });
                 } else if let Some(usage) = line.get("usage") {
                     let input = number(usage, "input_tokens");
@@ -96,9 +104,7 @@ impl Translator {
                     if input.is_some() || output.is_some() {
                         events.push(Event::UsageDelta {
                             tokens: input.unwrap_or(0) + output.unwrap_or(0),
-                            // `exec` reports the sum across every model call in
-                            // a turn, not the final call's context.
-                            context: None,
+                            context: input,
                         });
                     }
                 }
@@ -169,17 +175,13 @@ fn update_item(item: &Value) -> Vec<Event> {
         .unwrap_or_default()
 }
 
-fn complete_item(item: &Value, streamed_messages: &HashSet<String>) -> Vec<Event> {
+fn complete_item(item: &Value, include_agent_message: bool) -> Vec<Event> {
     match item.get("type").and_then(Value::as_str) {
         Some("agent_message" | "agentMessage") => item
             .get("text")
             .and_then(Value::as_str)
             .filter(|text| !text.is_empty())
-            .filter(|_| {
-                item.get("id")
-                    .and_then(Value::as_str)
-                    .is_none_or(|id| !streamed_messages.contains(id))
-            })
+            .filter(|_| include_agent_message)
             .map(|delta| {
                 vec![Event::AssistantText {
                     delta: delta.to_string(),
@@ -217,6 +219,24 @@ fn tool(item: &Value) -> Option<(String, String, Value)> {
                 .unwrap_or_else(|| "mcp".to_string()),
             item.get("arguments").cloned().unwrap_or(Value::Null),
         ),
+        "dynamicToolCall" => (
+            item.get("tool")
+                .and_then(Value::as_str)
+                .unwrap_or("tool")
+                .to_string(),
+            item.get("arguments").cloned().unwrap_or(Value::Null),
+        ),
+        "functionCallOutput" => (
+            item.get("name")
+                .and_then(Value::as_str)
+                .unwrap_or("tool")
+                .to_string(),
+            Value::Null,
+        ),
+        "imageView" => (
+            "view_image".to_string(),
+            json!({"path": item.get("path").cloned().unwrap_or(Value::Null)}),
+        ),
         "web_search" | "webSearch" => (
             "web_search".to_string(),
             json!({"query": item.get("query").cloned().unwrap_or(Value::Null)}),
@@ -228,6 +248,20 @@ fn tool(item: &Value) -> Option<(String, String, Value)> {
 }
 
 fn tool_output(item: &Value) -> String {
+    match item.get("type").and_then(Value::as_str) {
+        Some("dynamicToolCall") => {
+            return content_text(item.get("contentItems"), "inputText");
+        }
+        Some("mcpToolCall") => {
+            if let Some(message) = item.pointer("/error/message").and_then(Value::as_str) {
+                return message.to_string();
+            }
+            return content_text(item.pointer("/result/content"), "text");
+        }
+        Some("functionCallOutput") => return content_text(item.get("output"), "input_text"),
+        Some("imageView") => return String::new(),
+        _ => {}
+    }
     for key in [
         "aggregated_output",
         "aggregatedOutput",
@@ -247,6 +281,22 @@ fn tool_output(item: &Value) -> String {
     }
 }
 
+/// Text from a structured result, deliberately excluding its image data. The driver saves images
+/// beside the transcript; serializing a data URL here makes a screenshot a megabytes-long line and
+/// still cannot draw it.
+fn content_text(value: Option<&Value>, text_kind: &str) -> String {
+    match value {
+        Some(Value::String(text)) => text.clone(),
+        Some(Value::Array(parts)) => parts
+            .iter()
+            .filter(|part| part.get("type").and_then(Value::as_str) == Some(text_kind))
+            .filter_map(|part| part.get("text").and_then(Value::as_str))
+            .collect::<Vec<_>>()
+            .join("\n"),
+        _ => String::new(),
+    }
+}
+
 fn value_text(value: &Value) -> Option<String> {
     value
         .as_str()
@@ -316,7 +366,7 @@ mod tests {
             events[0],
             Event::UsageDelta {
                 tokens: 18,
-                context: None
+                context: Some(13)
             }
         );
         assert!(translator.completed());
@@ -380,7 +430,7 @@ mod tests {
         assert!(
             translator
                 .translate(&line(
-                    r#"{"method":"thread/tokenUsage/updated","params":{"tokenUsage":{"last":{"totalTokens":42}}}}"#
+                    r#"{"method":"thread/tokenUsage/updated","params":{"tokenUsage":{"last":{"inputTokens":39,"cachedInputTokens":30,"outputTokens":3,"reasoningOutputTokens":1,"totalTokens":42},"total":{"inputTokens":100,"cachedInputTokens":80,"outputTokens":9,"reasoningOutputTokens":2,"totalTokens":109},"modelContextWindow":258400}}}"#
                 ))
                 .is_empty()
         );
@@ -391,7 +441,7 @@ mod tests {
             vec![
                 Event::UsageDelta {
                     tokens: 42,
-                    context: None
+                    context: Some(39)
                 },
                 Event::Status {
                     state: SessionStatus::Idle
@@ -399,4 +449,38 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn structured_tool_results_keep_text_but_not_image_data() {
+        let mut translator = Translator::default();
+        let started = translator.translate(&line(
+            r#"{"method":"item/started","params":{"item":{"id":"tool-1","type":"dynamicToolCall","tool":"view_image","arguments":{"path":"shot.png"},"status":"inProgress"}}}"#,
+        ));
+        assert!(matches!(&started[0], Event::ToolStart { tool, .. } if tool == "view_image"));
+
+        let ended = translator.translate(&line(
+            r#"{"method":"item/completed","params":{"item":{"id":"tool-1","type":"dynamicToolCall","tool":"view_image","arguments":{},"status":"completed","contentItems":[{"type":"inputText","text":"looked"},{"type":"inputImage","imageUrl":"data:image/png;base64,aGVsbG8="}]}}}"#,
+        ));
+        assert_eq!(
+            ended,
+            vec![Event::ToolEnd {
+                id: "tool-1".to_string(),
+                output: "looked".to_string()
+            }]
+        );
+    }
+
+    #[test]
+    fn an_app_server_completion_never_repeats_streamed_text_after_adoption() {
+        // A newly adopted translator has not seen the deltas already recorded by the previous
+        // backend. The dialect, rather than process-local memory, decides that this is a copy.
+        let mut adopted = Translator::default();
+        assert!(
+            adopted
+                .translate(&line(
+                    r#"{"method":"item/completed","params":{"item":{"id":"message-1","type":"agentMessage","text":"the complete message"}}}"#
+                ))
+                .is_empty()
+        );
+    }
 }

server/src/session/driver.rs

@@ -15,6 +15,27 @@ use tokio::sync::mpsc;
 /// renders them identically.
 pub type ImageRef = String;
 
+/// Stores image bytes where the files route serves them and returns their transcript reference.
+/// Both CLI dialects produce images in different envelopes; the durable file and naming rule are
+/// part of the common event model and must not vary with that envelope.
+pub(in crate::session) fn store_image(
+    session_dir: &std::path::Path,
+    media_type: &str,
+    bytes: &[u8],
+) -> Option<ImageRef> {
+    let extension = crate::media::extension_for(media_type).unwrap_or("png");
+    let name = format!("{}.{extension}", super::random_hex());
+    let dir = session_dir.join("files");
+    if let Err(err) = wg_app_link::private::create_dir(&dir)
+        .map_err(std::io::Error::other)
+        .and_then(|()| std::fs::write(dir.join(&name), bytes))
+    {
+        tracing::error!("couldn't save produced image: {err}");
+        return None;
+    }
+    Some(name)
+}
+
 /// The name an upload is stored and served under: an image is
 /// `<hex>.<extension>` and is an [`ImageRef`] like any other; any other file
 /// keeps its own name after the hex, `<hex>-<name>`, because the name is what

server/src/session/mod.rs

@@ -2416,6 +2416,22 @@ fn launch(
             }
         });
     }
+    if provider.kind == DriverKind::CodexCli
+        && shared.context_tokens.lock().unwrap().is_none()
+        && let Some(thread_id) = codex::read_thread(&dir)
+    {
+        let transport = Transport::for_setup(setup);
+        let shared = Arc::clone(&shared);
+        tokio::spawn(async move {
+            if let Some(context) = codex::context_of(&transport, &thread_id).await {
+                // A live turn may have answered while the rollout read was in flight.
+                let mut held = shared.context_tokens.lock().unwrap();
+                if held.is_none() {
+                    *held = Some(context);
+                }
+            }
+        });
+    }
 
     // An imported session shares its transcript file with the CLI, so work
     // done at a terminal belongs in this session too and arrives without

server/src/session/transport.rs

@@ -252,6 +252,12 @@ impl Transport {
     /// would otherwise need a runtime to ask a machine a question. Both build the
     /// invocation the same way.
     pub fn capture_blocking(&self, launch: &Launch) -> Result<String> {
+        Ok(String::from_utf8_lossy(&self.capture_bytes_blocking(launch)?).into_owned())
+    }
+
+    /// The byte-preserving form of [`capture_blocking`](Self::capture_blocking), used when a
+    /// driver copies a file back from the machine it runs on.
+    pub fn capture_bytes_blocking(&self, launch: &Launch) -> Result<Vec<u8>> {
         let host = match self {
             Self::Here => None,
             Self::Ssh { ssh, .. } => Some(ssh),
@@ -273,7 +279,7 @@ impl Transport {
                 stderr
             });
         }
-        Ok(String::from_utf8_lossy(&output.stdout).into_owned())
+        Ok(output.stdout)
     }
 
     /// Runs `launch` with `input` on its stdin and reports everything it