iris/ai-app
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix Codex transcript streaming and images

Browse Source
This commit is contained in:
iris committed 2026-09-09 15:14:24 -04:00
1 parent 14dd520719
commit 4dc3e3d784
7 files changed
+422 -53
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/session/codex/translate.rs
+113 -29
View File
@@ -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()
);
}
}