Fix Codex transcript streaming and images

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);
+    }
 }