feat: retry LLM API calls once after 401 by force-refreshing the OAuth token

The Client trait's default chat_completions, chat_completions_streaming,
and embeddings methods now classify failures via ApiStatusError: on a
401 with a cached OAuth token, the token is distrusted (identity-aware
marker) and the call retried exactly once — the retry's prepare step
sees the marker and force-refreshes. Streaming retries only while the
SSE handler has received no content, preventing duplicate rendering.
A second 401 propagates the original error; other retry errors
propagate as-is. API-key clients never retry. No backoff by design:
cost is bounded to one refresh + one retry per failing request.
This commit is contained in:
2026-08-14 13:17:56 -06:00
parent 684f19250a
commit 2596194417
3 changed files with 202 additions and 9 deletions
+48
View File
@@ -176,6 +176,14 @@ impl SseHandler {
self.thinking.push(block);
}
/// Whether any output (text, tool calls, or thinking blocks) has been
/// accumulated. `Client::chat_completions_streaming` gates its 401 retry
/// on this: content already streamed to the user would be rendered a
/// second time by a retry, so partial responses are never retried.
pub fn has_received_content(&self) -> bool {
!self.buffer.is_empty() || !self.tool_calls.is_empty() || !self.thinking.is_empty()
}
pub fn abort(&self) -> AbortSignal {
self.abort_signal.clone()
}
@@ -422,6 +430,46 @@ mod tests {
assert!(error_message.contains("test_function_loop"));
}
fn new_handler() -> (
SseHandler,
tokio::sync::mpsc::UnboundedReceiver<SseEvent>,
) {
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
let abort_signal = crate::utils::create_abort_signal();
(SseHandler::new(sender, abort_signal), receiver)
}
#[test]
fn test_has_received_content_text() {
let (mut handler, _rx) = new_handler();
assert!(!handler.has_received_content());
handler.text("hello").unwrap();
assert!(handler.has_received_content());
}
#[test]
fn test_has_received_content_tool_call() {
let (mut handler, _rx) = new_handler();
assert!(!handler.has_received_content());
let call = ToolCall::new("test_function".to_string(), json!({"param": 1}), None);
handler.tool_call(call).unwrap();
assert!(handler.has_received_content());
}
#[test]
fn test_has_received_content_thinking() {
let (mut handler, _rx) = new_handler();
assert!(!handler.has_received_content());
handler.thinking_block(ThinkingBlock::Thinking {
thinking: "hmm".to_string(),
signature: "sig".to_string(),
});
assert!(handler.has_received_content());
}
fn split_chunks(text: &str) -> Vec<Vec<u8>> {
let len = text.len();
let cut1 = random_range(1..len - 1);