mirror of
https://github.com/openai/codex.git
synced 2026-04-24 14:45:27 +00:00
try something 3
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
|
||||
use bytes::Bytes;
|
||||
use codex_otel::otel_event_manager::OtelEventManager;
|
||||
@@ -30,9 +31,16 @@ pub async fn process_sse_wire<S, D>(
|
||||
let mut buffer = String::new();
|
||||
|
||||
loop {
|
||||
let start = Instant::now();
|
||||
let result = timeout(max_idle_duration, stream.next()).await;
|
||||
let duration = start.elapsed();
|
||||
match result {
|
||||
Err(_) => {
|
||||
otel_event_manager.sse_event_failed(
|
||||
None,
|
||||
duration,
|
||||
&"idle timeout waiting for SSE",
|
||||
);
|
||||
let _ = tx_event
|
||||
.send(Err(Error::Stream(
|
||||
"stream idle timeout fired before Completed event".to_string(),
|
||||
@@ -42,6 +50,7 @@ pub async fn process_sse_wire<S, D>(
|
||||
return;
|
||||
}
|
||||
Ok(Some(Err(err))) => {
|
||||
otel_event_manager.sse_event_failed(None, duration, &err);
|
||||
let _ = tx_event.send(Err(err)).await;
|
||||
return;
|
||||
}
|
||||
@@ -49,6 +58,11 @@ pub async fn process_sse_wire<S, D>(
|
||||
let chunk_str = match std::str::from_utf8(&chunk) {
|
||||
Ok(s) => s,
|
||||
Err(err) => {
|
||||
otel_event_manager.sse_event_failed(
|
||||
None,
|
||||
duration,
|
||||
&format!("UTF8 error: {err}"),
|
||||
);
|
||||
let _ = tx_event
|
||||
.send(Err(Error::Other(format!(
|
||||
"Invalid UTF-8 in SSE chunk: {err}"
|
||||
@@ -62,13 +76,29 @@ pub async fn process_sse_wire<S, D>(
|
||||
|
||||
buffer.push_str(&chunk_str);
|
||||
while let Some(frame) = next_frame(&mut buffer) {
|
||||
if !handle_frame(frame, &mut decoder, &tx_event, &otel_event_manager).await {
|
||||
if !handle_frame(
|
||||
frame,
|
||||
&mut decoder,
|
||||
&tx_event,
|
||||
&otel_event_manager,
|
||||
duration,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
if !drain_buffer(&mut buffer, &mut decoder, &tx_event, &otel_event_manager).await {
|
||||
if !drain_buffer(
|
||||
&mut buffer,
|
||||
&mut decoder,
|
||||
&tx_event,
|
||||
&otel_event_manager,
|
||||
duration,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return;
|
||||
}
|
||||
return;
|
||||
@@ -99,12 +129,13 @@ async fn drain_buffer<D>(
|
||||
decoder: &mut D,
|
||||
tx_event: &mpsc::Sender<Result<WireEvent>>,
|
||||
otel_event_manager: &OtelEventManager,
|
||||
duration: Duration,
|
||||
) -> bool
|
||||
where
|
||||
D: crate::client::WireResponseDecoder + Send,
|
||||
{
|
||||
while let Some(frame) = next_frame(buffer) {
|
||||
if !handle_frame(frame, decoder, tx_event, otel_event_manager).await {
|
||||
if !handle_frame(frame, decoder, tx_event, otel_event_manager, duration).await {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
@@ -114,7 +145,7 @@ where
|
||||
}
|
||||
|
||||
let remainder = std::mem::take(buffer);
|
||||
handle_frame(remainder, decoder, tx_event, otel_event_manager).await
|
||||
handle_frame(remainder, decoder, tx_event, otel_event_manager, duration).await
|
||||
}
|
||||
|
||||
async fn handle_frame<D>(
|
||||
@@ -122,23 +153,38 @@ async fn handle_frame<D>(
|
||||
decoder: &mut D,
|
||||
tx_event: &mpsc::Sender<Result<WireEvent>>,
|
||||
otel_event_manager: &OtelEventManager,
|
||||
duration: Duration,
|
||||
) -> bool
|
||||
where
|
||||
D: crate::client::WireResponseDecoder + Send,
|
||||
{
|
||||
if let Some(data) = parse_sse_frame(&frame) {
|
||||
if let Err(e) = decoder.on_frame(&data, tx_event, otel_event_manager).await {
|
||||
let _ = tx_event.send(Err(e)).await;
|
||||
return false;
|
||||
if let Some(frame) = parse_sse_frame(&frame) {
|
||||
if frame.data.trim() == "[DONE]" {
|
||||
otel_event_manager.sse_event_kind(&frame.event);
|
||||
return true;
|
||||
}
|
||||
|
||||
match decoder
|
||||
.on_frame(&frame.data, tx_event, otel_event_manager)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
otel_event_manager.sse_event_kind(&frame.event);
|
||||
}
|
||||
Err(e) => {
|
||||
otel_event_manager.sse_event_failed(Some(&frame.event), duration, &e);
|
||||
let _ = tx_event.send(Err(e)).await;
|
||||
return false;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
fn parse_sse_frame(frame: &str) -> Option<String> {
|
||||
fn parse_sse_frame(frame: &str) -> Option<SseFrame> {
|
||||
let mut data = String::new();
|
||||
let mut saw_event = false;
|
||||
let mut event: Option<String> = None;
|
||||
let mut saw_data_line = false;
|
||||
|
||||
for raw_line in frame.split('\n') {
|
||||
@@ -148,10 +194,10 @@ fn parse_sse_frame(frame: &str) -> Option<String> {
|
||||
}
|
||||
|
||||
if let Some(rest) = line.strip_prefix("event:") {
|
||||
if rest.trim_start().is_empty() && !saw_data_line {
|
||||
continue;
|
||||
let trimmed = rest.trim_start();
|
||||
if !trimmed.is_empty() {
|
||||
event = Some(trimmed.to_string());
|
||||
}
|
||||
saw_event = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -171,11 +217,19 @@ fn parse_sse_frame(frame: &str) -> Option<String> {
|
||||
}
|
||||
}
|
||||
|
||||
if data.is_empty() && !saw_event && !saw_data_line {
|
||||
if data.is_empty() && event.is_none() && !saw_data_line {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(data)
|
||||
Some(SseFrame {
|
||||
event: event.unwrap_or_else(|| "message".to_string()),
|
||||
data,
|
||||
})
|
||||
}
|
||||
|
||||
struct SseFrame {
|
||||
event: String,
|
||||
data: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,10 +1,14 @@
|
||||
use async_trait::async_trait;
|
||||
use codex_otel::otel_event_manager::OtelEventManager;
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::ReasoningItemContent;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::client::WireResponseDecoder;
|
||||
use crate::error::Error;
|
||||
use crate::error::Result;
|
||||
use crate::stream::WireEvent;
|
||||
|
||||
@@ -16,15 +20,19 @@ struct FunctionCallState {
|
||||
arguments: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct WireChatSseDecoder {
|
||||
fn_call_state: FunctionCallState,
|
||||
created_emitted: bool,
|
||||
assistant_started: bool,
|
||||
assistant_text: String,
|
||||
reasoning_started: bool,
|
||||
reasoning_text: String,
|
||||
}
|
||||
|
||||
impl WireChatSseDecoder {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
fn_call_state: FunctionCallState::default(),
|
||||
}
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -36,11 +44,11 @@ impl WireResponseDecoder for WireChatSseDecoder {
|
||||
tx: &mpsc::Sender<crate::error::Result<WireEvent>>,
|
||||
_otel: &OtelEventManager,
|
||||
) -> Result<()> {
|
||||
// Chat sends a terminal "[DONE]" frame; ignore it.
|
||||
let Ok(parsed_chunk) = serde_json::from_str::<Value>(json) else {
|
||||
// Chat sends a terminal "[DONE]" frame; ignore it. Treat other parse errors as failures.
|
||||
let parsed_chunk = serde_json::from_str::<Value>(json).map_err(|err| {
|
||||
debug!("failed to parse Chat SSE JSON: {}", json);
|
||||
return Ok(());
|
||||
};
|
||||
Error::Other(format!("failed to parse Chat SSE JSON: {err}"))
|
||||
})?;
|
||||
|
||||
let choices = parsed_chunk
|
||||
.get("choices")
|
||||
@@ -49,10 +57,29 @@ impl WireResponseDecoder for WireChatSseDecoder {
|
||||
.unwrap_or_default();
|
||||
|
||||
for choice in choices {
|
||||
if !self.created_emitted {
|
||||
let _ = tx.send(Ok(WireEvent::Created)).await;
|
||||
self.created_emitted = true;
|
||||
}
|
||||
|
||||
if let Some(delta) = choice.get("delta") {
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_array()) {
|
||||
for piece in content {
|
||||
if let Some(text) = piece.get("text").and_then(|t| t.as_str()) {
|
||||
if !self.assistant_started {
|
||||
self.assistant_started = true;
|
||||
let message = ResponseItem::Message {
|
||||
id: None,
|
||||
role: "assistant".to_string(),
|
||||
content: vec![ContentItem::OutputText {
|
||||
text: String::new(),
|
||||
}],
|
||||
};
|
||||
let value = serde_json::to_value(message)
|
||||
.unwrap_or_else(|_| Value::String(String::new()));
|
||||
let _ = tx.send(Ok(WireEvent::OutputItemAdded(value))).await;
|
||||
}
|
||||
self.assistant_text.push_str(text);
|
||||
let _ = tx
|
||||
.send(Ok(WireEvent::OutputTextDelta(text.to_string())))
|
||||
.await;
|
||||
@@ -80,6 +107,19 @@ impl WireResponseDecoder for WireChatSseDecoder {
|
||||
if let Some(reasoning) = delta.get("reasoning_content").and_then(|c| c.as_array()) {
|
||||
for entry in reasoning {
|
||||
if let Some(text) = entry.get("text").and_then(|t| t.as_str()) {
|
||||
if !self.reasoning_started {
|
||||
self.reasoning_started = true;
|
||||
let reasoning_item = ResponseItem::Reasoning {
|
||||
id: String::new(),
|
||||
summary: vec![],
|
||||
content: None,
|
||||
encrypted_content: None,
|
||||
};
|
||||
let value = serde_json::to_value(reasoning_item)
|
||||
.unwrap_or_else(|_| Value::String(String::new()));
|
||||
let _ = tx.send(Ok(WireEvent::OutputItemAdded(value))).await;
|
||||
}
|
||||
self.reasoning_text.push_str(text);
|
||||
let _ = tx
|
||||
.send(Ok(WireEvent::ReasoningContentDelta(text.to_string())))
|
||||
.await;
|
||||
@@ -88,23 +128,69 @@ impl WireResponseDecoder for WireChatSseDecoder {
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(finish_reason) = choice.get("finish_reason").and_then(|f| f.as_str())
|
||||
&& finish_reason == "tool_calls"
|
||||
&& self.fn_call_state.active
|
||||
{
|
||||
let function_name = self.fn_call_state.name.take().unwrap_or_default();
|
||||
let call_id = self.fn_call_state.call_id.take().unwrap_or_default();
|
||||
let arguments = self.fn_call_state.arguments.clone();
|
||||
self.fn_call_state = FunctionCallState::default();
|
||||
if let Some(finish_reason) = choice.get("finish_reason").and_then(|f| f.as_str()) {
|
||||
match finish_reason {
|
||||
"tool_calls" if self.fn_call_state.active => {
|
||||
let function_name = self.fn_call_state.name.take().unwrap_or_default();
|
||||
let call_id = self.fn_call_state.call_id.take().unwrap_or_default();
|
||||
let arguments = self.fn_call_state.arguments.clone();
|
||||
self.fn_call_state = FunctionCallState::default();
|
||||
|
||||
let item = serde_json::json!({
|
||||
"type": "function_call",
|
||||
"id": call_id,
|
||||
"call_id": call_id,
|
||||
"name": function_name,
|
||||
"arguments": arguments,
|
||||
});
|
||||
let _ = tx.send(Ok(WireEvent::OutputItemDone(item))).await;
|
||||
let item = serde_json::json!({
|
||||
"type": "function_call",
|
||||
"id": call_id,
|
||||
"call_id": call_id,
|
||||
"name": function_name,
|
||||
"arguments": arguments,
|
||||
});
|
||||
let _ = tx.send(Ok(WireEvent::OutputItemDone(item))).await;
|
||||
}
|
||||
"stop" | "length" => {
|
||||
if self.reasoning_started {
|
||||
let mut content = Vec::new();
|
||||
if !self.reasoning_text.is_empty() {
|
||||
content.push(ReasoningItemContent::ReasoningText {
|
||||
text: self.reasoning_text.clone(),
|
||||
});
|
||||
}
|
||||
let reasoning_item = ResponseItem::Reasoning {
|
||||
id: String::new(),
|
||||
summary: vec![],
|
||||
content: Some(content),
|
||||
encrypted_content: None,
|
||||
};
|
||||
let value = serde_json::to_value(reasoning_item)
|
||||
.unwrap_or_else(|_| Value::String(String::new()));
|
||||
let _ = tx.send(Ok(WireEvent::OutputItemDone(value))).await;
|
||||
}
|
||||
|
||||
if self.assistant_started {
|
||||
let message = ResponseItem::Message {
|
||||
id: None,
|
||||
role: "assistant".to_string(),
|
||||
content: vec![ContentItem::OutputText {
|
||||
text: self.assistant_text.clone(),
|
||||
}],
|
||||
};
|
||||
let value = serde_json::to_value(message)
|
||||
.unwrap_or_else(|_| Value::String(String::new()));
|
||||
let _ = tx.send(Ok(WireEvent::OutputItemDone(value))).await;
|
||||
}
|
||||
|
||||
let _ = tx
|
||||
.send(Ok(WireEvent::Completed {
|
||||
response_id: String::new(),
|
||||
token_usage: None,
|
||||
}))
|
||||
.await;
|
||||
|
||||
self.assistant_started = false;
|
||||
self.assistant_text.clear();
|
||||
self.reasoning_started = false;
|
||||
self.reasoning_text.clear();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -35,10 +35,10 @@ impl WireResponseDecoder for WireResponsesSseDecoder {
|
||||
tx: &mpsc::Sender<Result<WireEvent>>,
|
||||
otel: &OtelEventManager,
|
||||
) -> Result<()> {
|
||||
let Ok(event) = serde_json::from_str::<StreamEvent>(json) else {
|
||||
let event = serde_json::from_str::<StreamEvent>(json).map_err(|err| {
|
||||
debug!("failed to parse Responses SSE JSON: {}", json);
|
||||
return Ok(());
|
||||
};
|
||||
Error::Other(format!("failed to parse Responses SSE JSON: {err}"))
|
||||
})?;
|
||||
|
||||
match event.event_type.as_str() {
|
||||
"response.created" => {
|
||||
|
||||
@@ -302,13 +302,20 @@ impl ModelProviderInfo {
|
||||
account_id: None,
|
||||
})
|
||||
} else {
|
||||
match self.api_key()? {
|
||||
Some(key) => Some(AuthContext {
|
||||
match self.api_key() {
|
||||
Ok(Some(key)) => Some(AuthContext {
|
||||
mode: AuthMode::ApiKey,
|
||||
bearer_token: Some(key),
|
||||
account_id: None,
|
||||
}),
|
||||
None => auth.clone(),
|
||||
Ok(None) => auth.clone(),
|
||||
Err(err) => {
|
||||
if auth.is_some() {
|
||||
auth.clone()
|
||||
} else {
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user