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:
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user