Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: surface llm retry and continuation state

Browse Source
This commit is contained in:
Hare
2026-05-26 07:13:59 +09:00
parent 41402c0951
commit be753099ae
20 changed files with 626 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/llm-worker/src/lib.rs
+3 -1
View File
@@ -59,4 +59,6 @@ pub use interceptor::Interceptor;
pub use message::{ContentPart, Item, Message, Role};
pub use tool::{ToolCall, ToolOutputLimits, ToolResult};
pub use usage_record::UsageRecord;
pub use worker::{RunOutput, ToolRegistryError, Worker, WorkerConfig, WorkerError, WorkerResult};
pub use worker::{
LlmRetryNotice, RunOutput, ToolRegistryError, Worker, WorkerConfig, WorkerError, WorkerResult,
};
crates/llm-worker/src/llm_client/client.rs
+4 -8
View File
@@ -36,6 +36,8 @@ impl std::fmt::Display for ConfigWarning {
}
}
pub type ResponseStream = Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>;
/// LLMクライアントのtrait
///
/// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。
@@ -49,10 +51,7 @@ pub trait LlmClient: Send + Sync {
/// # Returns
/// * `Ok(Stream)` - イベントストリーム
/// * `Err(ClientError)` - エラー
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError>;
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError>;
/// Clone this client into a new `Box<dyn LlmClient>`.
///
@@ -85,10 +84,7 @@ impl Clone for Box<dyn LlmClient> {
/// これにより、動的ディスパッチを使用するクライアントも `Worker` で利用可能になる。
#[async_trait]
impl LlmClient for Box<dyn LlmClient> {
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
(**self).stream(request).await
}
crates/llm-worker/src/llm_client/error.rs
+20 -1
View File
@@ -1,6 +1,6 @@
//! LLMクライアントエラー型
use std::fmt;
use std::{fmt, time::Duration};
/// LLMクライアントのエラー
#[derive(Debug)]
@@ -16,6 +16,7 @@ pub enum ClientError {
status: Option<u16>,
code: Option<String>,
message: String,
retry_after: Option<Duration>,
},
/// 設定エラー
Config(String),
@@ -31,6 +32,7 @@ impl fmt::Display for ClientError {
status,
code,
message,
..
} => {
write!(f, "API error")?;
if let Some(s) = status {
@@ -68,6 +70,22 @@ impl From<serde_json::Error> for ClientError {
}
}
impl ClientError {
pub fn status(&self) -> Option<u16> {
match self {
ClientError::Api { status, .. } => *status,
_ => None,
}
}
pub fn retry_after(&self) -> Option<Duration> {
match self {
ClientError::Api { retry_after, .. } => *retry_after,
_ => None,
}
}
}
/// transient な失敗としてリトライ対象になるかを判定する。
///
/// 対象:
@@ -97,6 +115,7 @@ mod tests {
status,
code: None,
message: String::new(),
retry_after: None,
}
}
crates/llm-worker/src/llm_client/retry.rs
+4 -4
View File
@@ -1,8 +1,8 @@
//! HTTP transient エラー向けリトライポリシー。
//! LLM response stream を開く前の transient error 向けリトライポリシー。
//!
//! `transport.rs` の HTTP 送信〜ステータスチェック区間で `is_retryable`
//! が true を返した失敗をリトライする際に、待ち時間と打ち切り条件を
//! 提供する。SSE 読み出し開始後の失敗は対象外。
//! Worker が `LlmClient::stream` の open error に対して `is_retryable` を見て
//! retry / backoff / TUI event / cancellation をまとめて管理する。
//! SSE 読み出し開始後の失敗は対象外。
use std::time::Duration;
crates/llm-worker/src/llm_client/scheme/gemini/events.rs
+1
View File
@@ -131,6 +131,7 @@ impl GeminiScheme {
status: None,
code: Some("parse_error".to_string()),
message: format!("Failed to parse Gemini SSE data: {} -> {}", e, data),
retry_after: None,
})?;
let mut events = Vec::new();
crates/llm-worker/src/llm_client/scheme/openai_chat/events.rs
+1
View File
@@ -75,6 +75,7 @@ impl OpenAIScheme {
status: None,
code: Some("parse_error".to_string()),
message: format!("Failed to parse SSE data: {} -> {}", e, data),
retry_after: None,
})?;
let mut events = Vec::new();
crates/llm-worker/src/llm_client/scheme/openai_responses/events.rs
+1
View File
@@ -597,6 +597,7 @@ fn from_json<T: for<'de> Deserialize<'de>>(data: &str) -> Result<T, ClientError>
status: None,
code: Some("parse_error".to_string()),
message: format!("Failed to parse SSE data: {e}"),
retry_after: None,
})
}
crates/llm-worker/src/llm_client/transport.rs
+20 -58
View File
@@ -12,15 +12,12 @@ use async_trait::async_trait;
use eventsource_stream::Eventsource;
use futures::{Stream, StreamExt, TryStreamExt};
use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue, RETRY_AFTER};
use tokio::time::Instant;
use tracing::warn;
use super::auth::{AuthProvider, AuthRequirement};
use super::capability::ModelCapability;
use super::client::{ConfigWarning, LlmClient};
use super::error::{ClientError, is_retryable};
use super::client::{ConfigWarning, LlmClient, ResponseStream};
use super::error::ClientError;
use super::event::Event;
use super::retry::RetryPolicy;
use super::scheme::Scheme;
use super::types::{Request, RequestConfig};
@@ -67,7 +64,6 @@ pub struct HttpTransport<S: Scheme> {
base_url: String,
auth: ResolvedAuth,
capability: ModelCapability,
retry_policy: RetryPolicy,
}
impl<S: Scheme> HttpTransport<S> {
@@ -89,7 +85,6 @@ impl<S: Scheme> HttpTransport<S> {
base_url,
auth,
capability,
retry_policy: RetryPolicy::default(),
}
}
@@ -99,12 +94,6 @@ impl<S: Scheme> HttpTransport<S> {
self
}
/// リトライポリシーを差し替える（テスト用 / 将来の manifest 化フック）。
pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
self.retry_policy = policy;
self
}
fn build_url(&self) -> String {
let path = self.scheme.path(&self.model_id);
let url = format!("{}{}", self.base_url, path);
@@ -171,14 +160,12 @@ impl<S: Scheme + Clone> Clone for HttpTransport<S> {
base_url: self.base_url.clone(),
auth: self.auth.clone(),
capability: self.capability.clone(),
retry_policy: self.retry_policy.clone(),
}
}
}
/// エラーレスポンスを `ClientError::Api` に変換し、`Retry-After` の秒数を
/// 同時に取り出す。リトライループで wait の上書きに使う。
async fn classify_error_response(resp: reqwest::Response) -> (ClientError, Option<Duration>) {
/// エラーレスポンスを `ClientError::Api` に変換する。
async fn classify_error_response(resp: reqwest::Response) -> ClientError {
let status = resp.status().as_u16();
let retry_after = resp
.headers()
@@ -187,7 +174,7 @@ async fn classify_error_response(resp: reqwest::Response) -> (ClientError, Optio
.and_then(|s| s.trim().parse::<u64>().ok())
.map(Duration::from_secs);
let text = resp.text().await.unwrap_or_default();
let err = if let Ok(json) = serde_json::from_str::<serde_json::Value>(&text) {
if let Ok(json) = serde_json::from_str::<serde_json::Value>(&text) {
let error = json.get("error").unwrap_or(&json);
let code = error.get("type").and_then(|v| v.as_str()).map(String::from);
let message = error
@@ -199,15 +186,16 @@ async fn classify_error_response(resp: reqwest::Response) -> (ClientError, Optio
status: Some(status),
code,
message,
retry_after,
}
} else {
ClientError::Api {
status: Some(status),
code: None,
message: text,
retry_after,
}
};
(err, retry_after)
}
}
#[async_trait]
@@ -220,51 +208,25 @@ impl<S: Scheme + Clone + 'static> LlmClient for HttpTransport<S> {
self.scheme.validate_config(config)
}
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
let url = self.build_url();
let headers = self.build_headers().await?;
let body = self
.scheme
.build_request_body(&self.model_id, &request, &self.capability);
let policy = &self.retry_policy;
let started = Instant::now();
let mut attempt: u32 = 0;
let response = loop {
let send_result = self
.http_client
.post(&url)
.headers(headers.clone())
.json(&body)
.send()
.await;
let response = self
.http_client
.post(&url)
.headers(headers)
.json(&body)
.send()
.await
.map_err(ClientError::Http)?;
let (err, retry_after) = match send_result {
Ok(resp) if resp.status().is_success() => break resp,
Ok(resp) => classify_error_response(resp).await,
Err(e) => (ClientError::Http(e), None),
};
let next_attempt = attempt + 1;
if next_attempt >= policy.max_attempts || !is_retryable(&err) {
return Err(err);
}
let wait = retry_after.unwrap_or_else(|| policy.backoff(attempt));
if started.elapsed() + wait > policy.total_timeout {
return Err(err);
}
warn!(
error = %err,
attempt = next_attempt,
wait_ms = wait.as_millis() as u64,
"transient HTTP error, retrying"
);
tokio::time::sleep(wait).await;
attempt = next_attempt;
};
if !response.status().is_success() {
return Err(classify_error_response(response).await);
}
let scheme = self.scheme.clone();
let byte_stream = response.bytes_stream().map_err(std::io::Error::other);
crates/llm-worker/src/worker.rs
+207 -24
View File
@@ -1,5 +1,5 @@
use std::collections::HashMap;
use std::marker::PhantomData;
use std::{marker::PhantomData, time::Instant};
use futures::StreamExt;
use tokio::sync::mpsc;
@@ -17,8 +17,8 @@ use crate::{
PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction,
},
llm_client::{
ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ToolDefinition,
types::parse_tool_arguments,
ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream,
ToolDefinition, error::is_retryable, retry::RetryPolicy, types::parse_tool_arguments,
},
state::{Locked, Mutable, WorkerState},
timeline::event::{ErrorEvent, StatusEvent, UsageEvent},
@@ -99,6 +99,8 @@ enum ToolExecutionResult {
Paused,
}
const MAX_STREAM_CONTINUATIONS: u32 = 3;
/// Central component for managing LLM interactions
///
/// Receives input from the user, sends requests to the LLM, and
@@ -131,9 +133,28 @@ enum ToolExecutionResult {
/// let out = worker.run("Continue").await?;
/// let mut worker = out.worker;
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LlmRetryNotice {
/// 直近で失敗した attempt 番号。1 origin。
pub failed_attempt: u32,
pub max_attempts: u32,
pub wait: std::time::Duration,
pub elapsed: std::time::Duration,
pub status: Option<u16>,
pub error: String,
}
#[derive(Debug)]
enum StreamCompletion {
Complete,
Interrupted { reason: String },
}
pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
/// LLM client
client: C,
/// Retry policy for opening an LLM response stream.
retry_policy: RetryPolicy,
/// Event timeline
timeline: Timeline,
/// Text block collector (Timeline handler)
@@ -175,6 +196,10 @@ pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
llm_call_start_cbs: Vec<Box<dyn Fn(usize) + Send + Sync>>,
/// LlmCall-end callbacks
llm_call_end_cbs: Vec<Box<dyn Fn(usize) + Send + Sync>>,
/// Transport-level retry callbacks for a specific LlmCall.
llm_retry_cbs: Vec<Box<dyn Fn(usize, &LlmRetryNotice) + Send + Sync>>,
/// Stream continuation callbacks for a specific LlmCall.
llm_continuation_cbs: Vec<Box<dyn Fn(usize, u32, u32, &str) + Send + Sync>>,
/// Non-fatal warning callbacks. Invoked when the Worker wants to
/// surface an advisory message to the upper layer (e.g. Pod) so it
/// can be forwarded to the user — distinct from `tracing::warn!`,
@@ -355,6 +380,34 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
self.llm_call_end_cbs.push(Box::new(callback));
}
/// Register a transport-level retry callback.
pub fn on_llm_retry(
&mut self,
callback: impl Fn(usize, &LlmRetryNotice) + Send + Sync + 'static,
) {
self.llm_retry_cbs.push(Box::new(callback));
}
/// Register a stream continuation callback.
pub fn on_llm_continuation(
&mut self,
callback: impl Fn(usize, u32, u32, &str) + Send + Sync + 'static,
) {
self.llm_continuation_cbs.push(Box::new(callback));
}
fn emit_llm_continuation(
&self,
llm_call: usize,
attempt: u32,
max_attempts: u32,
reason: &str,
) {
for cb in &self.llm_continuation_cbs {
cb(llm_call, attempt, max_attempts, reason);
}
}
/// Register a non-fatal warning callback.
///
/// The callback is invoked with a short human-readable message
@@ -964,6 +1017,8 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
}
let mut stream_continuations: u32 = 0;
let mut continuing_stream = false;
loop {
if self.try_cancelled() {
info!("Execution cancelled");
@@ -973,9 +1028,11 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
let current_turn = self.turn_count;
debug!(turn = current_turn, "Turn start");
for cb in &self.turn_start_cbs {
cb(current_turn);
if !continuing_stream {
debug!(turn = current_turn, "Turn start");
for cb in &self.turn_start_cbs {
cb(current_turn);
}
}
// Drain interceptor-side inputs that are meant to land in
@@ -1080,13 +1137,50 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
// Stream LLM response
let request = self.build_request(&tool_definitions, &request_context);
self.stream_response(request).await?;
let stream_outcome = self.stream_response(request, current_llm_call).await?;
for cb in &self.llm_call_end_cbs {
cb(current_llm_call);
}
self.llm_call_count += 1;
if let StreamCompletion::Interrupted { reason } = stream_outcome {
stream_continuations += 1;
if stream_continuations > MAX_STREAM_CONTINUATIONS {
self.last_run_interrupted = true;
return Err(WorkerError::Client(ClientError::Api {
status: None,
code: None,
message: format!("LLM stream interrupted too many times: {reason}"),
retry_after: None,
}));
}
self.timeline.abort_current_block();
self.timeline.flush_usage();
let reasoning_items = self.reasoning_item_collector.take_collected();
let text_blocks = self.text_block_collector.take_collected();
// Do not recover tool calls from an interrupted stream. A completed
// tool_use is executable only when the provider finishes the stream.
let _dropped_tool_calls = self.tool_call_collector.take_collected();
let assistant_items =
self.build_assistant_items(&reasoning_items, &text_blocks, &[]);
if !assistant_items.is_empty() {
self.append_history_items(assistant_items);
}
self.emit_llm_continuation(
current_llm_call,
stream_continuations,
MAX_STREAM_CONTINUATIONS,
&reason,
);
continuing_stream = true;
continue;
}
stream_continuations = 0;
continuing_stream = false;
for cb in &self.turn_end_cbs {
cb(current_turn);
}
@@ -1138,8 +1232,88 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
}
async fn open_stream_with_retry(
&mut self,
request: Request,
llm_call: usize,
) -> Result<ResponseStream, WorkerError> {
let policy = self.retry_policy.clone();
let started = Instant::now();
let mut failed_attempt: u32 = 0;
loop {
let stream_result = tokio::select! {
stream_result = self.client.stream(request.clone()) => stream_result,
cancel = self.cancel_rx.recv() => {
if cancel.is_some() {
info!("Cancelled before stream started");
}
self.timeline.abort_current_block();
self.last_run_interrupted = true;
return Err(WorkerError::Cancelled);
}
};
match stream_result {
Ok(stream) => return Ok(stream),
Err(err) => {
let next_failed_attempt = failed_attempt + 1;
if next_failed_attempt >= policy.max_attempts || !is_retryable(&err) {
self.last_run_interrupted = true;
return Err(WorkerError::Client(err));
}
let wait = err
.retry_after()
.unwrap_or_else(|| policy.backoff(failed_attempt));
let elapsed = started.elapsed();
if elapsed + wait > policy.total_timeout {
self.last_run_interrupted = true;
return Err(WorkerError::Client(err));
}
warn!(
error = %err,
failed_attempt = next_failed_attempt,
wait_ms = wait.as_millis() as u64,
"transient LLM request error, retrying"
);
let notice = LlmRetryNotice {
failed_attempt: next_failed_attempt,
max_attempts: policy.max_attempts,
wait,
elapsed,
status: err.status(),
error: err.to_string(),
};
for cb in &self.llm_retry_cbs {
cb(llm_call, &notice);
}
tokio::select! {
_ = tokio::time::sleep(wait) => {}
cancel = self.cancel_rx.recv() => {
if cancel.is_some() {
info!("Cancelled during LLM retry backoff");
}
self.timeline.abort_current_block();
self.last_run_interrupted = true;
return Err(WorkerError::Cancelled);
}
}
failed_attempt = next_failed_attempt;
}
}
}
}
/// Open a stream, dispatch all events to the timeline, handle cancellation.
async fn stream_response(&mut self, request: Request) -> Result<(), WorkerError> {
async fn stream_response(
&mut self,
request: Request,
llm_call: usize,
) -> Result<StreamCompletion, WorkerError> {
debug!(
item_count = request.items.len(),
tool_count = request.tools.len(),
@@ -1147,18 +1321,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
"Sending request to LLM"
);
let mut stream = tokio::select! {
stream_result = self.client.stream(request) => stream_result
.inspect_err(|_| self.last_run_interrupted = true)?,
cancel = self.cancel_rx.recv() => {
if cancel.is_some() {
info!("Cancelled before stream started");
}
self.timeline.abort_current_block();
self.last_run_interrupted = true;
return Err(WorkerError::Cancelled);
}
};
let mut stream = self.open_stream_with_retry(request, llm_call).await?;
let mut event_count: usize = 0;
loop {
@@ -1175,12 +1338,17 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
warn!(error = %e, "Stream error");
}
}
let event = result
.inspect_err(|_| {
let event = match result {
Ok(event) => event,
Err(err) => {
self.last_run_interrupted = true;
// 部分情報でも発火しておく（料金会計用）
self.timeline.flush_usage();
})?;
return Ok(StreamCompletion::Interrupted {
reason: err.to_string(),
});
}
};
self.timeline.dispatch(&event);
}
None => break,
@@ -1200,7 +1368,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
// ストリーム完了時に集約済み Usage を 1 度だけ発火
self.timeline.flush_usage();
debug!(event_count = event_count, "Stream completed");
Ok(())
Ok(StreamCompletion::Complete)
}
/// Execute tools and push results to history.
@@ -1254,6 +1422,7 @@ impl<C: LlmClient> Worker<C, Mutable> {
Self {
client,
retry_policy: RetryPolicy::default(),
timeline,
text_block_collector,
tool_call_collector,
@@ -1270,6 +1439,8 @@ impl<C: LlmClient> Worker<C, Mutable> {
turn_end_cbs: Vec::new(),
llm_call_start_cbs: Vec::new(),
llm_call_end_cbs: Vec::new(),
llm_retry_cbs: Vec::new(),
llm_continuation_cbs: Vec::new(),
warning_cbs: Vec::new(),
tool_result_cbs: Vec::new(),
history_append_cbs: Vec::new(),
@@ -1385,6 +1556,12 @@ impl<C: LlmClient> Worker<C, Mutable> {
self
}
/// Set the retry policy used when opening an LLM response stream.
pub fn with_retry_policy(mut self, retry_policy: RetryPolicy) -> Self {
self.retry_policy = retry_policy;
self
}
/// Validate current configuration against the provider
///
/// Returns an error if there are unsupported settings.
@@ -1507,6 +1684,7 @@ impl<C: LlmClient> Worker<C, Mutable> {
let locked_prefix_len = self.history.len();
Worker {
client: self.client,
retry_policy: self.retry_policy,
timeline: self.timeline,
text_block_collector: self.text_block_collector,
tool_call_collector: self.tool_call_collector,
@@ -1523,6 +1701,8 @@ impl<C: LlmClient> Worker<C, Mutable> {
turn_end_cbs: self.turn_end_cbs,
llm_call_start_cbs: self.llm_call_start_cbs,
llm_call_end_cbs: self.llm_call_end_cbs,
llm_retry_cbs: self.llm_retry_cbs,
llm_continuation_cbs: self.llm_continuation_cbs,
warning_cbs: self.warning_cbs,
tool_result_cbs: self.tool_result_cbs,
history_append_cbs: self.history_append_cbs,
@@ -1594,6 +1774,7 @@ impl<C: LlmClient> Worker<C, Locked> {
pub fn unlock(self) -> Worker<C, Mutable> {
Worker {
client: self.client,
retry_policy: self.retry_policy,
timeline: self.timeline,
text_block_collector: self.text_block_collector,
tool_call_collector: self.tool_call_collector,
@@ -1610,6 +1791,8 @@ impl<C: LlmClient> Worker<C, Locked> {
turn_end_cbs: self.turn_end_cbs,
llm_call_start_cbs: self.llm_call_start_cbs,
llm_call_end_cbs: self.llm_call_end_cbs,
llm_retry_cbs: self.llm_retry_cbs,
llm_continuation_cbs: self.llm_continuation_cbs,
warning_cbs: self.warning_cbs,
tool_result_cbs: self.tool_result_cbs,
history_append_cbs: self.history_append_cbs,