From 74772b2ef18626faaa33db4fb974419e08af84a7 Mon Sep 17 00:00:00 2001 From: Milhous Date: Tue, 7 Jul 2026 13:48:26 +0800 Subject: [PATCH 1/4] fix(openai): retry copilot 401 and stream completions (#501) --- crates/puffer-core/runtime/copilot.rs | 35 + crates/puffer-core/runtime/openai.rs | 126 +++- .../runtime/openai/completions_session.rs | 688 +++++++++++++++--- crates/puffer-core/runtime/tests.rs | 254 ++++++- .../runtime/tests/agent_loop_e2e.rs | 52 ++ 5 files changed, 1043 insertions(+), 112 deletions(-) diff --git a/crates/puffer-core/runtime/copilot.rs b/crates/puffer-core/runtime/copilot.rs index a11e7942c..e84e95265 100644 --- a/crates/puffer-core/runtime/copilot.rs +++ b/crates/puffer-core/runtime/copilot.rs @@ -11,6 +11,8 @@ use anyhow::{bail, Context, Result}; use puffer_provider_registry::{apply_copilot_client_identity, COPILOT_TOKEN_URL}; use reqwest::blocking::Client; use std::collections::HashMap; +#[cfg(test)] +use std::sync::Arc; use std::sync::{Mutex, OnceLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -79,6 +81,34 @@ fn cache() -> &'static Mutex> { CACHE.get_or_init(|| Mutex::new(HashMap::new())) } +#[cfg(test)] +type TestBearerExchange = Arc Result + Send + Sync>; + +#[cfg(test)] +fn test_bearer_exchange() -> &'static Mutex> { + static EXCHANGE: OnceLock>> = OnceLock::new(); + EXCHANGE.get_or_init(|| Mutex::new(None)) +} + +#[cfg(test)] +pub(super) struct TestBearerExchangeGuard; + +#[cfg(test)] +impl Drop for TestBearerExchangeGuard { + fn drop(&mut self) { + *test_bearer_exchange().lock().unwrap() = None; + } +} + +#[cfg(test)] +pub(super) fn install_test_bearer_exchange(exchange: F) -> TestBearerExchangeGuard +where + F: Fn(&str) -> Result + Send + Sync + 'static, +{ + *test_bearer_exchange().lock().unwrap() = Some(Arc::new(exchange)); + TestBearerExchangeGuard +} + /// Drops the cached exchanged bearer for a GitHub token. Called when the chat /// endpoint rejects the bearer with 401 (e.g. it was invalidated before its /// cached expiry — revoked Copilot seat, server-side rotation), and on @@ -109,6 +139,11 @@ fn now_secs() -> u64 { /// small burst at first use / expiry; a shared in-flight map would add locking /// complexity for little gain, so we accept it. pub(crate) fn copilot_bearer_token(github_token: &str) -> Result { + #[cfg(test)] + if let Some(exchange) = test_bearer_exchange().lock().unwrap().clone() { + return exchange(github_token); + } + let now = now_secs(); if let Some(cached) = cache().lock().unwrap().get(github_token) { if cached.expires_at_secs > now + EXPIRY_SKEW_SECS { diff --git a/crates/puffer-core/runtime/openai.rs b/crates/puffer-core/runtime/openai.rs index 64813b050..51abbf9af 100644 --- a/crates/puffer-core/runtime/openai.rs +++ b/crates/puffer-core/runtime/openai.rs @@ -616,9 +616,7 @@ pub(super) fn resolve_openai_execution_config( // Auto-only plans (Free/Student): chat only works inside an // auto-mode session — the server rejects direct model selection // with `model_not_supported` without this header. - if let Some(session) = &copilot.session { - custom_headers.push(("Copilot-Session-Token".to_string(), session.token.clone())); - } + set_copilot_session_header(&mut custom_headers, copilot.session.as_ref()); return Ok(OpenAIExecutionConfig { provider_id: provider.id.clone(), request_config: OpenAIRequestConfig { @@ -730,22 +728,45 @@ fn codex_style_for_provider(provider: &ProviderDescriptor, oauth: bool) -> bool != Some("1") } +fn set_copilot_session_header( + custom_headers: &mut Vec<(String, String)>, + session: Option<&super::copilot::CopilotSession>, +) { + custom_headers.retain(|(key, _)| !key.eq_ignore_ascii_case("copilot-session-token")); + if let Some(session) = session { + custom_headers.push(("Copilot-Session-Token".to_string(), session.token.clone())); + } +} + /// On a 401 from the GitHub Copilot chat endpoint, drop the cached exchanged -/// bearer so the next turn re-exchanges the stored GitHub token. Copilot carries -/// no `refresh_token` (the OAuth refresh path is skipped), so without this a -/// bearer invalidated before its ~25min cached expiry would keep 401-ing every -/// turn until expiry with no auto-recovery. -fn evict_copilot_bearer_on_unauthorized( +/// bearer, re-exchange the stored GitHub OAuth token, update the active request +/// config, and let the caller retry the same request once. Copilot carries no +/// `refresh_token`, but the stored GitHub OAuth token is the durable credential +/// for minting a fresh short-lived Copilot bearer. +fn reexchange_copilot_bearer_on_unauthorized( auth_store: &AuthStore, - execution: &OpenAIExecutionConfig, + execution: &mut OpenAIExecutionConfig, unauthorized: bool, -) { +) -> Result { if !unauthorized || execution.provider_id != "github-copilot" { - return; - } - if let Some(StoredCredential::OAuth(credential)) = auth_store.get("github-copilot") { - super::copilot::invalidate_bearer(&credential.access_token); + return Ok(false); } + let github_token = match auth_store.get("github-copilot") { + Some(StoredCredential::OAuth(credential)) if !credential.access_token.is_empty() => { + credential.access_token.clone() + } + _ => return Ok(false), + }; + super::copilot::invalidate_bearer(&github_token); + let copilot = super::copilot::copilot_bearer_token(&github_token) + .context("failed to re-exchange GitHub OAuth token for Copilot bearer after 401")?; + execution.request_config.base_url = copilot.api_url; + execution.request_config.auth = OpenAIAuth::ApiKey(copilot.token); + set_copilot_session_header( + &mut execution.request_config.custom_headers, + copilot.session.as_ref(), + ); + Ok(true) } /// Sends a blocking OpenAI request and refreshes OAuth credentials once after a 401. @@ -781,11 +802,21 @@ where false, proxy, )?; - evict_copilot_bearer_on_unauthorized( + if reexchange_copilot_bearer_on_unauthorized( auth_store, execution, response.status == StatusCode::UNAUTHORIZED, - ); + )? { + let retry = build_request(&execution.request_config)?; + let retry_response = super::send_http_request_raw_with_proxy( + &retry.url, + &retry.headers, + &retry.body, + false, + proxy, + )?; + return parse_http_json_response(&retry.url, false, retry_response); + } if response.status != StatusCode::UNAUTHORIZED || execution.refresh_token.is_none() { return parse_http_json_response(&request.url, false, response); } @@ -832,6 +863,29 @@ pub(super) fn send_openai_request_with_refresh_streaming( where F: Fn(&OpenAIRequestConfig) -> Result, G: FnMut(TurnStreamEvent), +{ + send_openai_request_with_refresh_streaming_using_parser( + auth_store, + execution, + proxy, + build_request, + on_event, + parse_openai_stream_response, + ) +} + +pub(super) fn send_openai_request_with_refresh_streaming_using_parser( + auth_store: &mut AuthStore, + execution: &mut OpenAIExecutionConfig, + proxy: &ProxyConfig, + build_request: F, + on_event: &mut G, + parse_response: P, +) -> Result +where + F: Fn(&OpenAIRequestConfig) -> Result, + G: FnMut(TurnStreamEvent) + ?Sized, + P: Fn(&str, Response, &mut G) -> Result + Copy, { let request = build_request(&execution.request_config)?; // Layered retry: inner = connection-level (`retry_openai_transport`) @@ -870,13 +924,45 @@ where ); }, )?; - evict_copilot_bearer_on_unauthorized( + if reexchange_copilot_bearer_on_unauthorized( auth_store, execution, response.status() == StatusCode::UNAUTHORIZED, - ); + )? { + let retry = build_request(&execution.request_config)?; + let retry_response = super::retry_on_5xx( + || { + retry_openai_transport( + || { + send_openai_request_stream_raw( + &retry.url, + &retry.headers, + &retry.body, + proxy, + ) + }, + |attempt, max, error| { + on_event(TurnStreamEvent::RetryAttempt { + attempt, + max_attempts: max, + error: error.to_string(), + kind: RetryAttemptKind::Transport, + }); + }, + ) + }, + |attempt, max, status| { + tracing::warn!( + target: "puffer::runtime::openai", + "5xx retry (post-copilot-401-reexchange): attempt {attempt}/{max}, HTTP {}", + status.as_u16() + ); + }, + )?; + return parse_response(&retry.url, retry_response, on_event); + } if response.status() != StatusCode::UNAUTHORIZED || execution.refresh_token.is_none() { - return parse_openai_stream_response(&request.url, response, on_event); + return parse_response(&request.url, response, on_event); } let refresh_token = execution @@ -922,7 +1008,7 @@ where ); }, )?; - parse_openai_stream_response(&retry.url, retry_response, on_event) + parse_response(&retry.url, retry_response, on_event) } fn send_openai_request_stream_raw( diff --git a/crates/puffer-core/runtime/openai/completions_session.rs b/crates/puffer-core/runtime/openai/completions_session.rs index dac299061..dd33c1631 100644 --- a/crates/puffer-core/runtime/openai/completions_session.rs +++ b/crates/puffer-core/runtime/openai/completions_session.rs @@ -1,30 +1,27 @@ //! [`TurnSession`] impl for the OpenAI Chat Completions API. //! -//! No live SSE parser yet — the response comes back as one JSON -//! payload via `send_openai_request_with_refresh`. Streaming and -//! non-streaming `one_turn_*` variants both go through the same -//! request path; the streaming path additionally fires -//! `ThinkingDelta` and `TextDelta` events synthesized from the -//! parsed response so reasoning-capable Chat Completions providers -//! (Moonshot Kimi, Deepseek, OpenRouter relays, …) keep their -//! thinking blocks visible in the TUI. Real per-token streaming is -//! a follow-up (would need `stream: true` on the request body and a -//! Chat Completions SSE parser). - -use anyhow::Result; +//! Streaming requests send `stream: true` and parse Chat Completions +//! SSE (`data: {...}` chunks plus `[DONE]`). Non-SSE JSON responses +//! are still accepted as a compatibility fallback and synthesize the +//! same text/thinking events the old path emitted. + +use anyhow::{bail, Context, Result}; use puffer_provider_openai::{ - build_chat_completions_request, extract_chat_completions_reasoning, + build_chat_completions_request, build_json_post_request, extract_chat_completions_reasoning, extract_chat_completions_tool_calls, extract_chat_completions_visible_text, parse_chat_completions_response, OpenAIChatCompletionTool, OpenAIChatCompletionsRequest, - OpenAIChatResponseFormat, OpenAIRequestConfig, OpenAIResponsesToolChoiceMode, + OpenAIChatMessage, OpenAIChatResponseFormat, OpenAIRequestConfig, OpenAIResponseToolCall, + OpenAIResponsesToolChoiceMode, }; use puffer_provider_registry::{ AuthStore, OpenAiCompletionsCompat, ProviderDescriptor, ThinkingFormat, }; use puffer_resources::LoadedResources; use puffer_tools::ToolRegistry; +use reqwest::blocking::Response; use serde_json::{json, Value}; -use std::collections::HashSet; +use std::collections::{BTreeMap, HashSet}; +use std::io::{BufRead, Read}; use super::conversation::{ build_system_reminder, generate_openai_summary, items_to_chat_messages, @@ -32,7 +29,7 @@ use super::conversation::{ }; use super::{ parse_openai_text, parse_openai_text_fallback, send_openai_request_with_refresh, - OpenAIExecutionConfig, + send_openai_request_with_refresh_streaming_using_parser, OpenAIExecutionConfig, }; use crate::permissions::{load_runtime_permission_context_with_inputs, RuntimePermissionInputs}; use crate::runtime::agent_loop::{AssistantTurn, TurnSession}; @@ -72,24 +69,7 @@ impl TurnSession for OpenAICompletionsTurnSession { items: &mut Vec, on_event: &mut dyn FnMut(TurnStreamEvent), ) -> Result { - // Use the rich `send_and_parse` (not `one_turn_blocking`) so - // we keep `reasoning_chain` after parsing. Synthesize streaming - // events from the (already-final) response so the TUI's - // thinking + assistant cards stay populated. Real per-token - // streaming is a follow-up — needs `stream: true` on the wire - // body and a Chat Completions SSE parser. For reasoning-capable - // providers this is the difference between "thinking block - // visible" and "thinking block missing" (issue raised against - // `kimi-coding/k2p5` with `effort: xhigh`). - let result = self.send_and_parse(state, auth_store, items)?; - if let Some(reasoning) = result.reasoning_chain.as_deref() { - if !reasoning.is_empty() { - on_event(TurnStreamEvent::ThinkingDelta(reasoning.to_string())); - } - } - if !result.assistant_text.is_empty() { - on_event(TurnStreamEvent::TextDelta(result.assistant_text.clone())); - } + let result = self.send_streaming_and_parse(state, auth_store, items, on_event)?; Ok(result.into_assistant_turn()) } @@ -129,6 +109,7 @@ struct CompletionsTurnResult { tool_calls: Vec, assistant_text: String, reasoning_chain: Option, + emitted_tool_call_ids: HashSet, } impl CompletionsTurnResult { @@ -138,7 +119,7 @@ impl CompletionsTurnResult { tool_calls: self.tool_calls, assistant_text: self.assistant_text, input_tokens_hint: None, - emitted_tool_call_ids: HashSet::new(), + emitted_tool_call_ids: self.emitted_tool_call_ids, usage_report: None, } } @@ -156,6 +137,52 @@ impl OpenAICompletionsTurnSession { auth_store: &mut AuthStore, items: &mut Vec, ) -> Result { + let prepared = self.prepare_request(state, items); + + let body_for_each_attempt = move |request_config: &OpenAIRequestConfig| { + build_prepared_chat_completions_request(request_config, &prepared, false) + }; + + let response: Value = send_openai_request_with_refresh( + auth_store, + &mut self.execution, + &state.config.network.proxy, + body_for_each_attempt, + )?; + + Self::result_from_response_value(&response, state) + } + + fn send_streaming_and_parse( + &mut self, + state: &mut AppState, + auth_store: &mut AuthStore, + items: &mut Vec, + on_event: &mut dyn FnMut(TurnStreamEvent), + ) -> Result { + let prepared = self.prepare_request(state, items); + + let body_for_each_attempt = move |request_config: &OpenAIRequestConfig| { + build_prepared_chat_completions_request(request_config, &prepared, true) + }; + + let streamed = send_openai_request_with_refresh_streaming_using_parser( + auth_store, + &mut self.execution, + &state.config.network.proxy, + body_for_each_attempt, + on_event, + parse_chat_completions_stream_response, + )?; + + Ok(Self::result_from_stream(streamed, state)) + } + + fn prepare_request( + &self, + state: &AppState, + items: &[ConversationItem], + ) -> PreparedCompletionsRequest { let messages = items_to_chat_messages( items, Some(&self.system_prompt), @@ -164,15 +191,6 @@ impl OpenAICompletionsTurnSession { Some(&self.system_reminder), ); - let model_id = self.model_id.clone(); - let tools = self.tools.clone(); - let response_format = self.response_format.clone(); - - // Resolve effort + thinking params per the model's compat. When - // `requires_reasoning_content_on_assistant_messages` is set, also - // patch every prior assistant message to carry an empty - // `reasoning_content` so DeepSeek-style relays don't reject the - // replay. let reasoning_fields = resolve_reasoning_fields( self.compat.as_ref(), self.model_supports_reasoning, @@ -188,40 +206,30 @@ impl OpenAICompletionsTurnSession { } } } - let messages_for_attempt = messages.clone(); - let body_for_each_attempt = move |request_config: &OpenAIRequestConfig| { - build_chat_completions_request( - request_config, - &OpenAIChatCompletionsRequest { - model: model_id.clone(), - messages: messages_for_attempt.clone(), - tools: tools.clone(), - tool_choice: if tools.is_empty() { - None - } else { - Some(OpenAIResponsesToolChoiceMode::Auto) - }, - response_format: response_format.clone(), - reasoning_effort: reasoning_fields.reasoning_effort.clone(), - reasoning: reasoning_fields.reasoning.clone(), - thinking: reasoning_fields.thinking.clone(), - enable_thinking: reasoning_fields.enable_thinking, - chat_template_kwargs: reasoning_fields.chat_template_kwargs.clone(), - }, - ) - }; + PreparedCompletionsRequest { + model_id: self.model_id.clone(), + messages, + tools: self.tools.clone(), + response_format: self.response_format.clone(), + reasoning_fields, + } + } - let response: Value = send_openai_request_with_refresh( - auth_store, - &mut self.execution, - &state.config.network.proxy, - body_for_each_attempt, - )?; + fn result_from_response_value( + response: &Value, + state: &AppState, + ) -> Result { + let parsed = chat_completions_result_from_json_value(response)?; + Ok(Self::result_from_stream(parsed, state)) + } - let parsed = parse_chat_completions_response(&serde_json::to_string(&response)?)?; - let tool_calls_vendor = extract_chat_completions_tool_calls(&parsed)?; - let tool_calls: Vec = tool_calls_vendor + fn result_from_stream( + streamed: ChatCompletionsStreamResult, + state: &AppState, + ) -> CompletionsTurnResult { + let tool_calls: Vec = streamed + .tool_calls .iter() .map(|tc| ToolCallRequest { call_id: tc.call_id.clone(), @@ -229,20 +237,13 @@ impl OpenAICompletionsTurnSession { input: serde_json::to_string(&tc.arguments).unwrap_or_default(), }) .collect(); - - // Strip any inline `…` block from the visible - // text so it doesn't double-render alongside the thinking card - // emitted from `extract_chat_completions_reasoning`. - let assistant_text_from_msg = extract_chat_completions_visible_text(&parsed); - let reasoning_chain = extract_chat_completions_reasoning(&parsed); - let mut pre_tool_items: Vec = Vec::new(); - if !assistant_text_from_msg.trim().is_empty() { + if !streamed.assistant_text.trim().is_empty() { pre_tool_items.push(ConversationItem::assistant_message( - &assistant_text_from_msg, + &streamed.assistant_text, )); } - for tc in &tool_calls_vendor { + for tc in &streamed.tool_calls { pre_tool_items.push(ConversationItem::FunctionCall { call_id: tc.call_id.clone(), name: tc.name.clone(), @@ -251,26 +252,468 @@ impl OpenAICompletionsTurnSession { } let final_assistant_text = if tool_calls.is_empty() { - if assistant_text_from_msg.trim().is_empty() { - parse_openai_text(&response) - .or_else(|_| parse_openai_text_fallback(&response, state)) + if streamed.assistant_text.trim().is_empty() { + parse_openai_text(&streamed.raw_response) + .or_else(|_| parse_openai_text_fallback(&streamed.raw_response, state)) .unwrap_or_default() } else { - assistant_text_from_msg + streamed.assistant_text } } else { String::new() }; - Ok(CompletionsTurnResult { + CompletionsTurnResult { pre_tool_items, tool_calls, assistant_text: final_assistant_text, - reasoning_chain, + reasoning_chain: streamed.reasoning_chain, + emitted_tool_call_ids: streamed.emitted_tool_call_ids, + } + } +} + +struct PreparedCompletionsRequest { + model_id: String, + messages: Vec, + tools: Vec, + response_format: Option, + reasoning_fields: ReasoningFields, +} + +fn build_prepared_chat_completions_request( + config: &OpenAIRequestConfig, + prepared: &PreparedCompletionsRequest, + stream: bool, +) -> Result { + let request = OpenAIChatCompletionsRequest { + model: prepared.model_id.clone(), + messages: prepared.messages.clone(), + tools: prepared.tools.clone(), + tool_choice: if prepared.tools.is_empty() { + None + } else { + Some(OpenAIResponsesToolChoiceMode::Auto) + }, + response_format: prepared.response_format.clone(), + reasoning_effort: prepared.reasoning_fields.reasoning_effort.clone(), + reasoning: prepared.reasoning_fields.reasoning.clone(), + thinking: prepared.reasoning_fields.thinking.clone(), + enable_thinking: prepared.reasoning_fields.enable_thinking, + chat_template_kwargs: prepared.reasoning_fields.chat_template_kwargs.clone(), + }; + + if !stream { + return build_chat_completions_request(config, &request); + } + + let path = config + .chat_completions_path + .as_deref() + .unwrap_or("/v1/chat/completions"); + let mut body = serde_json::to_value(&request)?; + body["stream"] = Value::Bool(true); + build_json_post_request(config, path, &body) +} + +struct ChatCompletionsStreamResult { + assistant_text: String, + reasoning_chain: Option, + tool_calls: Vec, + emitted_tool_call_ids: HashSet, + raw_response: Value, +} + +fn parse_chat_completions_stream_response( + url: &str, + response: Response, + on_event: &mut G, +) -> Result +where + G: FnMut(TurnStreamEvent) + ?Sized, +{ + let status = response.status(); + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string); + if !status.is_success() { + let text = response.text().unwrap_or_default(); + if let Some(quota) = + super::super::quota::classify_response("openai", status.as_u16(), &text) + { + return Err(anyhow::Error::new(quota)); + } + bail!("request failed with status {}: {}", status, text); + } + + let mut reader = std::io::BufReader::new(response); + let looks_like_sse = if is_chat_completions_event_stream(content_type.as_deref(), "") { + true + } else { + let prefix = reader.fill_buf()?; + let prefix = std::str::from_utf8(prefix).unwrap_or_default(); + is_chat_completions_event_stream(content_type.as_deref(), prefix) + }; + + if looks_like_sse { + return parse_chat_completions_sse_reader(reader, on_event) + .with_context(|| format!("failed to parse Chat Completions SSE response from {url}")); + } + + let mut text = String::new(); + reader.read_to_string(&mut text)?; + let raw: Value = serde_json::from_str(&text) + .with_context(|| format!("response from {url} was not valid JSON"))?; + let result = chat_completions_result_from_json_value(&raw) + .with_context(|| format!("response from {url} was not a valid Chat Completions payload"))?; + if let Some(reasoning) = result.reasoning_chain.as_deref() { + if !reasoning.is_empty() { + on_event(TurnStreamEvent::ThinkingDelta(reasoning.to_string())); + } + } + if !result.assistant_text.is_empty() { + on_event(TurnStreamEvent::TextDelta(result.assistant_text.clone())); + } + Ok(result) +} + +fn chat_completions_result_from_json_value( + response: &Value, +) -> Result { + let parsed = parse_chat_completions_response(&serde_json::to_string(response)?)?; + let tool_calls = extract_chat_completions_tool_calls(&parsed)?; + Ok(ChatCompletionsStreamResult { + assistant_text: extract_chat_completions_visible_text(&parsed), + reasoning_chain: extract_chat_completions_reasoning(&parsed), + tool_calls, + emitted_tool_call_ids: HashSet::new(), + raw_response: response.clone(), + }) +} + +fn is_chat_completions_event_stream(content_type: Option<&str>, text: &str) -> bool { + content_type.is_some_and(|value| value.starts_with("text/event-stream")) + || text.trim_start().starts_with("data:") + || text.trim_start().starts_with("event:") +} + +fn parse_chat_completions_sse_reader( + mut reader: R, + on_event: &mut G, +) -> Result +where + R: BufRead, + G: FnMut(TurnStreamEvent) + ?Sized, +{ + let mut state = ChatCompletionsSseState::default(); + let mut line = String::new(); + let mut data_lines = Vec::new(); + + loop { + line.clear(); + let read = reader.read_line(&mut line)?; + if read == 0 { + flush_chat_completions_sse_event(&data_lines, &mut state, on_event)?; + break; + } + + let trimmed = line.trim_end_matches(['\r', '\n']); + if trimmed.is_empty() { + if flush_chat_completions_sse_event(&data_lines, &mut state, on_event)? { + data_lines.clear(); + break; + } + data_lines.clear(); + continue; + } + + if let Some(data) = trimmed.strip_prefix("data:") { + data_lines.push(data.trim_start().to_string()); + } + } + + if !state.terminal { + bail!("stream closed before Chat Completions [DONE]"); + } + Ok(state.into_result()) +} + +fn flush_chat_completions_sse_event( + data_lines: &[String], + state: &mut ChatCompletionsSseState, + on_event: &mut G, +) -> Result +where + G: FnMut(TurnStreamEvent) + ?Sized, +{ + let data = data_lines.join("\n"); + if data.is_empty() { + return Ok(false); + } + if data == "[DONE]" { + state.emit_complete_tool_calls(true, on_event); + state.terminal = true; + return Ok(true); + } + + let event: Value = serde_json::from_str(&data) + .with_context(|| format!("invalid Chat Completions SSE payload: {data}"))?; + state.process_event(&event, on_event)?; + Ok(false) +} + +#[derive(Default)] +struct ChatCompletionsSseState { + id: Option, + finish_reason: Option, + terminal: bool, + assistant_text: String, + reasoning_chain: String, + tool_call_deltas: BTreeMap, + emitted_tool_call_ids: HashSet, +} + +impl ChatCompletionsSseState { + fn process_event(&mut self, event: &Value, on_event: &mut G) -> Result<()> + where + G: FnMut(TurnStreamEvent) + ?Sized, + { + if let Some(error) = event.get("error") { + let message = error + .get("message") + .and_then(Value::as_str) + .or_else(|| event.get("message").and_then(Value::as_str)) + .unwrap_or("Chat Completions stream failed"); + bail!("{message}"); + } + + if self.id.is_none() { + if let Some(id) = event.get("id").and_then(Value::as_str) { + self.id = Some(id.to_string()); + } + } + + for choice in event + .get("choices") + .and_then(Value::as_array) + .into_iter() + .flatten() + { + if let Some(delta) = choice.get("delta") { + self.process_delta(delta, on_event)?; + } + if let Some(reason) = choice.get("finish_reason").and_then(Value::as_str) { + if !reason.is_empty() { + self.finish_reason = Some(reason.to_string()); + if reason == "tool_calls" { + self.emit_complete_tool_calls(true, on_event); + } + } + } + } + + Ok(()) + } + + fn process_delta(&mut self, delta: &Value, on_event: &mut G) -> Result<()> + where + G: FnMut(TurnStreamEvent) + ?Sized, + { + if let Some(content) = delta.get("content").and_then(Value::as_str) { + self.assistant_text.push_str(content); + on_event(TurnStreamEvent::TextDelta(content.to_string())); + } + + for key in ["reasoning_content", "reasoning"] { + if let Some(reasoning) = delta.get(key).and_then(Value::as_str) { + self.reasoning_chain.push_str(reasoning); + on_event(TurnStreamEvent::ThinkingDelta(reasoning.to_string())); + } + } + + if let Some(tool_calls) = delta.get("tool_calls").and_then(Value::as_array) { + for (position, tool_call) in tool_calls.iter().enumerate() { + let index = tool_call + .get("index") + .and_then(Value::as_u64) + .map(|value| value as usize) + .unwrap_or(position); + { + let entry = self.tool_call_deltas.entry(index).or_default(); + if let Some(id) = tool_call.get("id").and_then(Value::as_str) { + if !id.is_empty() { + entry.call_id = id.to_string(); + } + } + if let Some(kind) = tool_call.get("type").and_then(Value::as_str) { + if !kind.is_empty() { + entry.kind = kind.to_string(); + } + } + if let Some(function) = tool_call.get("function") { + if let Some(name) = function.get("name").and_then(Value::as_str) { + if !name.is_empty() { + entry.name = name.to_string(); + } + } + if let Some(arguments) = function.get("arguments").and_then(Value::as_str) { + entry.arguments.push_str(arguments); + } + } + } + self.maybe_emit_tool_call(index, false, on_event); + } + } + + Ok(()) + } + + fn maybe_emit_tool_call(&mut self, index: usize, allow_raw: bool, on_event: &mut G) + where + G: FnMut(TurnStreamEvent) + ?Sized, + { + if self + .tool_call_deltas + .get(&index) + .map(|entry| entry.emitted) + .unwrap_or(true) + { + return; + } + let Some((tool_call, raw_arguments)) = self.completed_tool_call(index, allow_raw) else { + return; + }; + on_event(TurnStreamEvent::ToolCallsRequested(vec![ToolCallRequest { + call_id: tool_call.call_id.clone(), + tool_id: tool_call.name.clone(), + input: raw_arguments, + }])); + self.emitted_tool_call_ids.insert(tool_call.call_id.clone()); + if let Some(entry) = self.tool_call_deltas.get_mut(&index) { + entry.emitted = true; + } + } + + fn emit_complete_tool_calls(&mut self, allow_raw: bool, on_event: &mut G) + where + G: FnMut(TurnStreamEvent) + ?Sized, + { + let indexes: Vec = self.tool_call_deltas.keys().copied().collect(); + for index in indexes { + self.maybe_emit_tool_call(index, allow_raw, on_event); + } + } + + fn completed_tool_call( + &self, + index: usize, + allow_raw: bool, + ) -> Option<(OpenAIResponseToolCall, String)> { + let entry = self.tool_call_deltas.get(&index)?; + if entry.call_id.is_empty() || entry.name.is_empty() { + return None; + } + let parsed_arguments = match serde_json::from_str::(&entry.arguments) { + Ok(value) => value, + Err(_) if allow_raw => Value::String(entry.arguments.clone()), + Err(_) => return None, + }; + Some(( + OpenAIResponseToolCall { + item_id: None, + status: None, + call_id: entry.call_id.clone(), + name: entry.name.clone(), + arguments: parsed_arguments, + }, + entry.arguments.clone(), + )) + } + + fn into_result(mut self) -> ChatCompletionsStreamResult { + let mut noop = |_| {}; + self.emit_complete_tool_calls(true, &mut noop); + let tool_calls = self + .tool_call_deltas + .keys() + .filter_map(|index| self.completed_tool_call(*index, true).map(|(call, _)| call)) + .collect::>(); + let raw_response = self.build_raw_response(); + ChatCompletionsStreamResult { + assistant_text: self.assistant_text, + reasoning_chain: (!self.reasoning_chain.is_empty()).then_some(self.reasoning_chain), + tool_calls, + emitted_tool_call_ids: self.emitted_tool_call_ids, + raw_response, + } + } + + fn build_raw_response(&self) -> Value { + let tool_calls = self + .tool_call_deltas + .values() + .filter(|entry| !entry.call_id.is_empty() && !entry.name.is_empty()) + .map(|entry| { + json!({ + "id": entry.call_id, + "type": if entry.kind.is_empty() { "function" } else { entry.kind.as_str() }, + "function": { + "name": entry.name, + "arguments": entry.arguments, + } + }) + }) + .collect::>(); + + let mut message = json!({ + "role": "assistant", + "content": if self.assistant_text.is_empty() { + Value::Null + } else { + Value::String(self.assistant_text.clone()) + }, + }); + if !self.reasoning_chain.is_empty() { + message["reasoning_content"] = Value::String(self.reasoning_chain.clone()); + } + if !tool_calls.is_empty() { + message["tool_calls"] = Value::Array(tool_calls); + } + + json!({ + "id": self.id.clone().unwrap_or_default(), + "object": "chat.completion", + "choices": [{ + "index": 0, + "message": message, + "finish_reason": self.finish_reason.clone().unwrap_or_else(|| "stop".to_string()), + }], }) } } +#[derive(Default)] +struct ChatCompletionsToolCallDelta { + call_id: String, + kind: String, + name: String, + arguments: String, + emitted: bool, +} + +#[cfg(test)] +fn parse_chat_completions_sse_response_for_tests( + stream: &str, + on_event: &mut G, +) -> Result +where + G: FnMut(TurnStreamEvent), +{ + parse_chat_completions_sse_reader(std::io::BufReader::new(stream.as_bytes()), on_event) +} + pub(super) fn setup_completions_session( state: &mut AppState, resources: &LoadedResources, @@ -584,3 +1027,66 @@ mod reasoning_fields_tests { assert!(f.reasoning_effort.is_none()); } } + +#[cfg(test)] +mod streaming_tests { + use super::*; + use crate::runtime::TurnStreamEvent; + + #[test] + fn parses_gemini_complete_tool_call_chunk() { + let stream = concat!( + "data: {\"id\":\"chatcmpl-gemini\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"tool_calls\":[{\"index\":0,\"id\":\"call_gemini_1\",\"type\":\"function\",\"function\":{\"name\":\"read_file\",\"arguments\":\"{\\\"path\\\":\\\"Cargo.toml\\\"}\"}}]},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-gemini\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":\"tool_calls\"}]}\n\n", + "data: [DONE]\n\n" + ); + let mut events = Vec::new(); + let parsed = + parse_chat_completions_sse_response_for_tests(stream, &mut |event| events.push(event)) + .unwrap(); + + assert_eq!(parsed.assistant_text, ""); + assert_eq!(parsed.tool_calls.len(), 1); + assert_eq!(parsed.tool_calls[0].call_id, "call_gemini_1"); + assert_eq!(parsed.tool_calls[0].name, "read_file"); + assert_eq!( + parsed.tool_calls[0].arguments, + json!({ "path": "Cargo.toml" }) + ); + assert!(events.iter().any(|event| matches!( + event, + TurnStreamEvent::ToolCallsRequested(calls) + if calls.len() == 1 + && calls[0].call_id == "call_gemini_1" + && calls[0].input == "{\"path\":\"Cargo.toml\"}" + ))); + } + + #[test] + fn parses_openai_fragmented_tool_call_deltas() { + let stream = concat!( + "data: {\"id\":\"chatcmpl-openai\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"I'll check. \"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-openai\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"tool_calls\":[{\"index\":0,\"id\":\"call_openai_1\",\"type\":\"function\",\"function\":{\"name\":\"read_file\",\"arguments\":\"{\\\"pa\"}}]},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-openai\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"th\\\":\\\"Cargo\"}}]},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-openai\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\".toml\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n", + "data: [DONE]\n\n" + ); + let mut text_deltas = Vec::new(); + let parsed = parse_chat_completions_sse_response_for_tests(stream, &mut |event| { + if let TurnStreamEvent::TextDelta(delta) = event { + text_deltas.push(delta); + } + }) + .unwrap(); + + assert_eq!(text_deltas, vec!["I'll check. ".to_string()]); + assert_eq!(parsed.assistant_text, "I'll check. "); + assert_eq!(parsed.tool_calls.len(), 1); + assert_eq!(parsed.tool_calls[0].call_id, "call_openai_1"); + assert_eq!(parsed.tool_calls[0].name, "read_file"); + assert_eq!( + parsed.tool_calls[0].arguments, + json!({ "path": "Cargo.toml" }) + ); + } +} diff --git a/crates/puffer-core/runtime/tests.rs b/crates/puffer-core/runtime/tests.rs index f80c3922a..d7139bf9f 100644 --- a/crates/puffer-core/runtime/tests.rs +++ b/crates/puffer-core/runtime/tests.rs @@ -1,7 +1,7 @@ use super::*; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use base64::Engine as _; -use puffer_config::{ensure_workspace_dirs, ConfigPaths, PufferConfig}; +use puffer_config::{ensure_workspace_dirs, ConfigPaths, ProxyConfig, PufferConfig}; use puffer_provider_openai::{ OpenAIAuth, OpenAIRequestConfig, OpenAIResponseToolCall, OpenAIResponsesTextConfig, OpenAIResponsesTextFormat, OpenAIResponsesTool, @@ -1241,6 +1241,258 @@ fn parse_openai_sse_response_streaming_emits_text_deltas() { ); } +fn copilot_test_execution(base_url: String, token: &str) -> super::openai::OpenAIExecutionConfig { + super::openai::OpenAIExecutionConfig { + provider_id: "github-copilot".to_string(), + request_config: OpenAIRequestConfig { + base_url, + version: "test".to_string(), + auth: OpenAIAuth::ApiKey(token.to_string()), + originator: "codex_cli_rs".to_string(), + session_id: Some("session-test".to_string()), + account_id: None, + custom_headers: Vec::new(), + query_params: Vec::new(), + chat_completions_path: Some("/chat/completions".to_string()), + responses_path: None, + }, + refresh_token: None, + codex_style: false, + } +} + +fn copilot_test_auth_store(github_token: &str) -> AuthStore { + let mut auth_store = AuthStore::default(); + auth_store.set_oauth( + "github-copilot", + OAuthCredential { + access_token: github_token.to_string(), + refresh_token: String::new(), + expires_at_ms: 0, + account_id: None, + organization_id: None, + email: None, + plan_type: None, + rate_limit_tier: None, + scopes: Vec::new(), + organization_name: None, + organization_role: None, + workspace_role: None, + }, + ); + auth_store +} + +fn install_copilot_exchange_for_test( + github_token: &'static str, + base_url: String, + token: &'static str, +) -> super::copilot::TestBearerExchangeGuard { + super::copilot::install_test_bearer_exchange(move |received| { + assert_eq!(received, github_token); + Ok(super::copilot::CopilotAuth { + token: token.to_string(), + api_url: base_url.clone(), + session: Some(super::copilot::CopilotSession { + token: format!("{token}-session"), + }), + }) + }) +} + +fn spawn_copilot_retry_server( + expected_requests: usize, + response_for_index: F, +) -> (String, Arc>>, thread::JoinHandle<()>) +where + F: Fn(usize) -> (u16, &'static str, String) + Send + 'static, +{ + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let address = listener.local_addr().unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let request_log = Arc::clone(&requests); + let server = thread::spawn(move || { + let deadline = Instant::now() + Duration::from_secs(2); + let mut handled = 0_usize; + while handled < expected_requests && Instant::now() < deadline { + match listener.accept() { + Ok((mut stream, _)) => { + stream.set_nonblocking(false).unwrap(); + let mut buffer = vec![0_u8; 65_536]; + let bytes = stream.read(&mut buffer).unwrap(); + let request = String::from_utf8_lossy(&buffer[..bytes]).to_string(); + request_log.lock().unwrap().push(request); + let (status, content_type, body) = response_for_index(handled); + let reason = if status == 200 { "OK" } else { "Unauthorized" }; + let response = format!( + "HTTP/1.1 {status} {reason}\r\ncontent-type: {content_type}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + stream.write_all(response.as_bytes()).unwrap(); + handled += 1; + } + Err(error) if error.kind() == ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(10)); + } + Err(error) => panic!("listener accept failed: {error}"), + } + } + }); + (format!("http://{address}"), requests, server) +} + +#[test] +fn copilot_non_streaming_401_reexchanges_bearer_and_retries_once() { + let _guard = env_lock(); + let (base_url, requests, server) = spawn_copilot_retry_server(2, |index| { + if index == 0 { + ( + 401, + "application/json", + json!({ "error": "expired bearer" }).to_string(), + ) + } else { + (200, "application/json", json!({ "ok": true }).to_string()) + } + }); + let _exchange = + install_copilot_exchange_for_test("gho-test-token", base_url.clone(), "fresh-copilot"); + let mut auth_store = copilot_test_auth_store("gho-test-token"); + let mut execution = copilot_test_execution(base_url, "stale-copilot"); + + let response = super::openai::send_openai_request_with_refresh( + &mut auth_store, + &mut execution, + &ProxyConfig::default(), + |config| { + puffer_provider_openai::build_json_post_request( + config, + "/chat/completions", + &json!({ "model": "gpt-4o", "messages": [] }), + ) + }, + ); + server.join().unwrap(); + let response = response.unwrap(); + + assert_eq!(response["ok"], json!(true)); + assert_eq!( + execution.request_config.auth, + OpenAIAuth::ApiKey("fresh-copilot".to_string()) + ); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer stale-copilot")); + let second = requests[1].to_ascii_lowercase(); + assert!(second.contains("authorization: bearer fresh-copilot")); + assert!(second.contains("copilot-session-token: fresh-copilot-session")); +} + +#[test] +fn copilot_streaming_401_reexchanges_bearer_and_retries_once() { + let _guard = env_lock(); + let (base_url, requests, server) = spawn_copilot_retry_server(2, |index| { + if index == 0 { + ( + 401, + "application/json", + json!({ "error": "expired bearer" }).to_string(), + ) + } else { + ( + 200, + "text/event-stream", + concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_copilot\"}}\n\n", + "event: response.output_text.delta\n", + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"stream \"}\n\n", + "event: response.output_text.delta\n", + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_copilot\",\"status\":\"completed\"}}\n\n" + ) + .to_string(), + ) + } + }); + let _exchange = + install_copilot_exchange_for_test("gho-test-token", base_url.clone(), "fresh-copilot"); + let mut auth_store = copilot_test_auth_store("gho-test-token"); + let mut execution = copilot_test_execution(base_url, "stale-copilot"); + let mut deltas = Vec::new(); + + let response = super::openai::send_openai_request_with_refresh_streaming( + &mut auth_store, + &mut execution, + &ProxyConfig::default(), + |config| { + puffer_provider_openai::build_json_post_request( + config, + "/chat/completions", + &json!({ "model": "gpt-4o", "messages": [], "stream": true }), + ) + }, + &mut |event| { + if let TurnStreamEvent::TextDelta(delta) = event { + deltas.push(delta); + } + }, + ); + server.join().unwrap(); + let response = response.unwrap(); + + assert_eq!(response.assistant_text, "stream ok"); + assert_eq!(deltas, vec!["stream ".to_string(), "ok".to_string()]); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[1] + .to_ascii_lowercase() + .contains("authorization: bearer fresh-copilot")); +} + +#[test] +fn copilot_401_retry_stops_after_second_unauthorized() { + let _guard = env_lock(); + let (base_url, requests, server) = spawn_copilot_retry_server(2, |_| { + ( + 401, + "application/json", + json!({ "error": "still unauthorized" }).to_string(), + ) + }); + let _exchange = + install_copilot_exchange_for_test("gho-test-token", base_url.clone(), "fresh-copilot"); + let mut auth_store = copilot_test_auth_store("gho-test-token"); + let mut execution = copilot_test_execution(base_url, "stale-copilot"); + + let response = super::openai::send_openai_request_with_refresh( + &mut auth_store, + &mut execution, + &ProxyConfig::default(), + |config| { + puffer_provider_openai::build_json_post_request( + config, + "/chat/completions", + &json!({ "model": "gpt-4o", "messages": [] }), + ) + }, + ); + server.join().unwrap(); + let error = response.unwrap_err().to_string(); + + assert!(error.contains("request failed with status 401")); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2, "must retry once, not loop forever"); + assert!(requests[1] + .to_ascii_lowercase() + .contains("authorization: bearer fresh-copilot")); +} + #[test] fn execute_user_prompt_refreshes_openai_oauth_after_401() { let _guard = refresh_env_lock() diff --git a/crates/puffer-core/runtime/tests/agent_loop_e2e.rs b/crates/puffer-core/runtime/tests/agent_loop_e2e.rs index a616aae71..c3826163a 100644 --- a/crates/puffer-core/runtime/tests/agent_loop_e2e.rs +++ b/crates/puffer-core/runtime/tests/agent_loop_e2e.rs @@ -596,6 +596,58 @@ fn openai_completions_agent_loop_runs_tool_then_text() { ); } +#[test] +fn openai_completions_agent_loop_uses_real_sse_streaming() { + let temp = tempfile::tempdir().unwrap(); + let (base_url, requests, server) = spawn_server("text/event-stream", 1, |_| { + concat!( + "data: {\"id\":\"chatcmpl-sse\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hello \"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-sse\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"world\"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"chatcmpl-sse\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n" + ) + .to_string() + }); + + let mut registry = ProviderRegistry::new(); + registry.register(openai_completions_provider(base_url)); + let mut auth_store = AuthStore::default(); + auth_store.set_api_key("openai-completions-test", "sk-openai"); + + let mut state = AppState::new( + PufferConfig::default(), + temp.path().to_path_buf(), + session_for(temp.path()), + ); + state.current_provider = Some("openai-completions-test".to_string()); + state.current_model = Some("openai-completions-test/gpt-5".to_string()); + + let mut text_deltas = Vec::new(); + let turn = execute_user_prompt_streaming( + &mut state, + &LoadedResources::default(), + ®istry, + &mut auth_store, + "say hello", + |event| { + if let TurnStreamEvent::TextDelta(delta) = event { + text_deltas.push(delta); + } + }, + ) + .unwrap(); + + server.join().unwrap(); + + assert_eq!(turn.assistant_text, "hello world"); + assert_eq!(text_deltas, vec!["hello ".to_string(), "world".to_string()]); + let captured = requests.lock().unwrap(); + assert_eq!(captured.len(), 1); + let body = extract_request_body(&captured[0]); + let body_json: Value = serde_json::from_str(body).unwrap(); + assert_eq!(body_json["stream"], json!(true), "request body: {body}"); +} + // --------------------------------------------------------------------------- // Cross-provider behavior: same prompt + same tool, both Anthropic and // OpenAI Responses produce semantically equivalent end states (one tool From cd717f783dad14ca4deda079f150f4a34365b8d6 Mon Sep 17 00:00:00 2001 From: Milhous Date: Tue, 7 Jul 2026 13:48:44 +0800 Subject: [PATCH 2/4] fix(copilot): surface device polling network errors (#501) --- crates/puffer-cli/src/copilot_login.rs | 135 ++++++++++++++++++------- 1 file changed, 101 insertions(+), 34 deletions(-) diff --git a/crates/puffer-cli/src/copilot_login.rs b/crates/puffer-cli/src/copilot_login.rs index 27c49a33f..1a58132dc 100644 --- a/crates/puffer-cli/src/copilot_login.rs +++ b/crates/puffer-cli/src/copilot_login.rs @@ -7,7 +7,7 @@ //! provider credential; the runtime later exchanges it for a short-lived //! Copilot bearer (see `puffer-core/runtime/copilot.rs`). -use anyhow::{bail, Context, Result}; +use anyhow::{anyhow, bail, Context, Result}; use puffer_provider_registry::COPILOT_USER_AGENT; use serde::Deserialize; use std::time::Duration; @@ -76,6 +76,7 @@ pub(crate) fn start_device_flow() -> Result { } /// Outcome of a single poll of the device-flow token endpoint. +#[derive(Debug)] pub(crate) enum DeviceFlowPoll { /// User has not authorized yet — keep polling. Pending, @@ -87,39 +88,32 @@ pub(crate) enum DeviceFlowPoll { Failed(String), } -/// Polls the token endpoint once with the device code. -pub(crate) fn poll_device_flow(device_code: &str) -> Result { - #[derive(Deserialize)] - struct Resp { - #[serde(default)] - access_token: Option, - #[serde(default)] - error: Option, +struct DevicePollHttpResponse { + status: reqwest::StatusCode, + body: String, +} + +#[derive(Deserialize)] +struct DevicePollResponse { + #[serde(default)] + access_token: Option, + #[serde(default)] + error: Option, +} + +fn classify_device_flow_poll_response( + response: Result, +) -> Result { + let response = response?; + if !response.status.is_success() { + bail!( + "GitHub device-flow token poll failed ({}): {}", + response.status, + response.body + ); } - let client = http_client()?; - // A single poll must not abort the whole login on a transient blip. Network - // errors and non-JSON/unknown bodies (e.g. a 5xx HTML error page from an - // infra hiccup) are treated as Pending so the caller keeps polling until the - // device code genuinely expires; only GitHub's documented terminal device- - // flow errors end the flow. - let response = match client - .post(ACCESS_TOKEN_URL) - .header("Accept", "application/json") - .header("User-Agent", COPILOT_USER_AGENT) - .form(&[ - ("client_id", COPILOT_CLIENT_ID), - ("device_code", device_code), - ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"), - ]) - .send() - { - Ok(response) => response, - Err(_) => return Ok(DeviceFlowPoll::Pending), - }; - let body = response.text().unwrap_or_default(); - let Ok(parsed) = serde_json::from_str::(&body) else { - return Ok(DeviceFlowPoll::Pending); - }; + let parsed: DevicePollResponse = + serde_json::from_str(&response.body).context("parsing GitHub device-flow poll response")?; if let Some(token) = parsed.access_token.filter(|t| !t.is_empty()) { return Ok(DeviceFlowPoll::Done(token)); } @@ -136,7 +130,80 @@ pub(crate) fn poll_device_flow(device_code: &str) -> Result { | "incorrect_device_code" | "device_flow_disabled"), ) => Ok(DeviceFlowPoll::Failed(err.to_string())), - // Unknown error code — treat as transient rather than aborting. + // Unknown GitHub error code — treat as transient rather than aborting. Some(_) => Ok(DeviceFlowPoll::Pending), } } + +/// Polls the token endpoint once with the device code. +pub(crate) fn poll_device_flow(device_code: &str) -> Result { + let client = http_client()?; + // GitHub's device-flow protocol has explicit non-terminal states + // (`authorization_pending`, `slow_down`). Transport failures are not one of + // them: surface those as poll errors so desktop callers can use their + // consecutive-error guard instead of waiting until device-code expiry. + let response = match client + .post(ACCESS_TOKEN_URL) + .header("Accept", "application/json") + .header("User-Agent", COPILOT_USER_AGENT) + .form(&[ + ("client_id", COPILOT_CLIENT_ID), + ("device_code", device_code), + ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"), + ]) + .send() + { + Ok(response) => { + let status = response.status(); + let body = response.text().unwrap_or_default(); + Ok(DevicePollHttpResponse { status, body }) + } + Err(error) => Err(anyhow!("GitHub device-flow poll network error: {error}")), + }; + classify_device_flow_poll_response(response) +} + +#[cfg(test)] +mod tests { + use super::*; + use anyhow::anyhow; + use reqwest::StatusCode; + + fn classify_ok(body: &str) -> Result { + classify_device_flow_poll_response(Ok(DevicePollHttpResponse { + status: StatusCode::OK, + body: body.to_string(), + })) + } + + #[test] + fn authorization_pending_remains_pending() { + let result = classify_ok(r#"{"error":"authorization_pending"}"#).unwrap(); + assert!(matches!(result, DeviceFlowPoll::Pending)); + } + + #[test] + fn slow_down_remains_slow_down() { + let result = classify_ok(r#"{"error":"slow_down"}"#).unwrap(); + assert!(matches!(result, DeviceFlowPoll::SlowDown)); + } + + #[test] + fn transport_errors_are_not_mapped_to_pending() { + let error = classify_device_flow_poll_response(Err(anyhow!("connection refused"))) + .expect_err("transport failures must reject the poll RPC"); + assert!(error.to_string().contains("connection refused")); + } + + #[test] + fn malformed_poll_response_is_not_mapped_to_pending() { + let error = classify_ok("bad gateway") + .expect_err("malformed poll responses are not protocol pending states"); + assert!( + error + .to_string() + .contains("parsing GitHub device-flow poll response"), + "{error:#}" + ); + } +} From 25c2a064490596de5e938ee6efa92d1810c5301c Mon Sep 17 00:00:00 2001 From: Milhous Date: Tue, 7 Jul 2026 13:49:03 +0800 Subject: [PATCH 3/4] fix(auth): hydrate Gemini env keys for google (#501) --- crates/puffer-cli/src/non_interactive.rs | 101 +++++++++++++++++++++-- 1 file changed, 93 insertions(+), 8 deletions(-) diff --git a/crates/puffer-cli/src/non_interactive.rs b/crates/puffer-cli/src/non_interactive.rs index 868abc4be..5f0d1049e 100644 --- a/crates/puffer-cli/src/non_interactive.rs +++ b/crates/puffer-cli/src/non_interactive.rs @@ -600,14 +600,20 @@ fn compose_prompt( } fn hydrate_env_auth(auth_store: &mut AuthStore) { - for (provider, env_name) in [ - ("openai", "OPENAI_API_KEY"), - ("anthropic", "ANTHROPIC_API_KEY"), - ] { - if let Ok(value) = std::env::var(env_name) { - let trimmed = value.trim(); - if !trimmed.is_empty() { - auth_store.set_api_key(provider, trimmed.to_string()); + let mappings: &[(&str, &[&str])] = &[ + ("openai", &["OPENAI_API_KEY"]), + ("anthropic", &["ANTHROPIC_API_KEY"]), + ("google", &["GEMINI_API_KEY", "GOOGLE_API_KEY"]), + ]; + + for (provider, env_names) in mappings { + for env_name in *env_names { + if let Ok(value) = std::env::var(env_name) { + let trimmed = value.trim(); + if !trimmed.is_empty() { + auth_store.set_api_key(*provider, trimmed.to_string()); + break; + } } } } @@ -798,6 +804,47 @@ impl ReplayArtifact { mod tests { use super::*; use puffer_provider_registry::ProviderDescriptor; + use std::sync::{Mutex, OnceLock}; + + fn env_lock() -> &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) + } + + struct EnvGuard { + name: &'static str, + prior: Option, + } + + impl EnvGuard { + fn set(name: &'static str, value: &str) -> Self { + let prior = std::env::var(name).ok(); + std::env::set_var(name, value); + Self { name, prior } + } + + fn remove(name: &'static str) -> Self { + let prior = std::env::var(name).ok(); + std::env::remove_var(name); + Self { name, prior } + } + } + + impl Drop for EnvGuard { + fn drop(&mut self) { + match self.prior.take() { + Some(value) => std::env::set_var(self.name, value), + None => std::env::remove_var(self.name), + } + } + } + + fn api_key(auth_store: &AuthStore, provider: &str) -> Option { + match auth_store.get(provider) { + Some(StoredCredential::ApiKey { key }) => Some(key.clone()), + _ => None, + } + } #[test] fn compose_prompt_includes_transcript_and_skill() { @@ -833,6 +880,44 @@ mod tests { assert!(!looks_like_jsonl_transcript(r#"{"role":"user"}"#)); } + #[test] + fn hydrate_env_auth_uses_gemini_api_key_for_google() { + let _lock = env_lock() + .lock() + .unwrap_or_else(|poison| poison.into_inner()); + let _openai = EnvGuard::remove("OPENAI_API_KEY"); + let _anthropic = EnvGuard::remove("ANTHROPIC_API_KEY"); + let _google = EnvGuard::remove("GOOGLE_API_KEY"); + let _gemini = EnvGuard::set("GEMINI_API_KEY", " gemini-key "); + let mut auth_store = AuthStore::default(); + + hydrate_env_auth(&mut auth_store); + + assert_eq!( + api_key(&auth_store, "google").as_deref(), + Some("gemini-key") + ); + } + + #[test] + fn hydrate_env_auth_uses_google_api_key_alias_for_google() { + let _lock = env_lock() + .lock() + .unwrap_or_else(|poison| poison.into_inner()); + let _openai = EnvGuard::remove("OPENAI_API_KEY"); + let _anthropic = EnvGuard::remove("ANTHROPIC_API_KEY"); + let _gemini = EnvGuard::remove("GEMINI_API_KEY"); + let _google = EnvGuard::set("GOOGLE_API_KEY", " google-key "); + let mut auth_store = AuthStore::default(); + + hydrate_env_auth(&mut auth_store); + + assert_eq!( + api_key(&auth_store, "google").as_deref(), + Some("google-key") + ); + } + #[test] fn custom_model_selector_registers_unknown_provider_model() { let mut providers = ProviderRegistry::new(); From 5022e650e546cbf2c0b546fa51101c195607bf05 Mon Sep 17 00:00:00 2001 From: Milhous Date: Tue, 7 Jul 2026 13:49:19 +0800 Subject: [PATCH 4/4] test(cli): set puffer home for tmux agent loop (#501) --- crates/puffer-cli/tests/tmux_agent_loop.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/puffer-cli/tests/tmux_agent_loop.rs b/crates/puffer-cli/tests/tmux_agent_loop.rs index e8418960e..20493fcd9 100644 --- a/crates/puffer-cli/tests/tmux_agent_loop.rs +++ b/crates/puffer-cli/tests/tmux_agent_loop.rs @@ -366,11 +366,11 @@ fn tmux_agent_loop_renders_assistant_reply_from_mock_anthropic() { "sh", &[ "-lc", - // HOME=workspace makes `.puffer/` in the workspace be the - // workspace config dir. Tracing PUFFER_HTTP_TRACE_PATH lets + // PUFFER_HOME=workspace makes `.puffer/` in the workspace be the + // user config dir. Tracing PUFFER_HTTP_TRACE_PATH lets // post-mortem inspection see the wire bytes if the test fails. &format!( - "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", + "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", ws = workspace.display(), bin = binary ), @@ -440,7 +440,7 @@ fn tmux_agent_loop_accepts_codex_default_provider_alias() { &[ "-lc", &format!( - "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", + "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", ws = workspace.display(), bin = binary ), @@ -548,7 +548,7 @@ fn tmux_agent_loop_drives_tool_round_trip_in_tui() { &[ "-lc", &format!( - "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", + "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", ws = workspace.display(), bin = binary ), @@ -641,7 +641,7 @@ fn tmux_agent_loop_validates_workflow_shorthand_in_tui() { &[ "-lc", &format!( - "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", + "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", ws = workspace.display(), bin = binary ), @@ -743,7 +743,7 @@ fn tmux_agent_loop_answers_ask_user_question_with_keyboard_selection() { &[ "-lc", &format!( - "HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", + "PUFFER_HOME='{ws}' HOME='{ws}' PUFFER_HTTP_TRACE_PATH='{ws}/wire.log' '{bin}'", ws = workspace.display(), bin = binary ),