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/tests/callback_test.rs
+63 -3
View File
@@ -4,17 +4,77 @@
mod common;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use async_trait::async_trait;
use common::MockLlmClient;
use llm_worker::Worker;
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent as ClientStatusEvent};
use llm_worker::llm_client::retry::RetryPolicy;
use llm_worker::llm_client::{ClientError, LlmClient, Request, ResponseStream};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
// =============================================================================
// Tests
// =============================================================================
#[derive(Clone)]
struct FailOnceClient {
calls: Arc<AtomicUsize>,
events: Vec<Event>,
}
#[async_trait]
impl LlmClient for FailOnceClient {
async fn stream(&self, _request: Request) -> Result<ResponseStream, ClientError> {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
return Err(ClientError::Api {
status: Some(504),
code: None,
message: "gateway timeout".into(),
retry_after: None,
});
}
Ok(Box::pin(futures::stream::iter(
self.events.clone().into_iter().map(Ok),
)))
}
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
}
#[tokio::test]
async fn test_callback_llm_retry_event() {
let events = vec![Event::Status(ClientStatusEvent {
status: ResponseStatus::Completed,
})];
let client = FailOnceClient {
calls: Arc::new(AtomicUsize::new(0)),
events,
};
let mut worker = Worker::new(client).with_retry_policy(RetryPolicy {
base: Duration::from_millis(1),
cap: Duration::from_millis(1),
max_attempts: 2,
total_timeout: Duration::from_secs(1),
});
let notices = Arc::new(Mutex::new(Vec::new()));
let sink = notices.clone();
worker.on_llm_retry(move |llm_call, notice| {
sink.lock().unwrap().push((llm_call, notice.clone()));
});
let result = worker.run("retry once").await;
assert!(result.is_ok(), "worker should succeed after one retry");
let notices = notices.lock().unwrap();
assert_eq!(notices.len(), 1);
assert_eq!(notices[0].0, 0);
assert_eq!(notices[0].1.failed_attempt, 1);
assert_eq!(notices[0].1.max_attempts, 2);
assert_eq!(notices[0].1.status, Some(504));
}
/// Verify that on_text_block correctly receives delta and stop events
#[tokio::test]
crates/llm-worker/tests/common/mod.rs
+1
View File
@@ -59,6 +59,7 @@ impl LlmClient for MockLlmClient {
status: Some(500),
code: Some("mock_error".to_string()),
message: "No more mock responses".to_string(),
retry_after: None,
});
}
let events = self.responses[count].clone();
crates/llm-worker/tests/transport_retry_test.rs
+56 -122
View File
@@ -1,12 +1,7 @@
//! HTTP transport の transient エラーリトライ挙動の integration テスト。
//! HTTP transport の単発 request / error classification テスト。
//!
//! 対応チケット: `tickets/llm-worker-transient-retry.md`。
//! - 503 / 529 / connect refused でリトライ発火
//! - max_attempts 上限到達でエラー
//! - `Retry-After` ヘッダで指数バックオフを上書き
//! - `parse_sse` 由来の `ClientError::Sse`（mid-stream 想定）はリトライしない
use std::time::{Duration, Instant};
//! Retry/backoff は Worker の lifecycle 管理に属するため、transport は 1 回だけ
//! request を送り、HTTP status / Retry-After を `ClientError` に載せて返す。
use futures::StreamExt;
use llm_worker::llm_client::LlmClient;
@@ -14,16 +9,16 @@ use llm_worker::llm_client::auth::AuthRequirement;
use llm_worker::llm_client::capability::ModelCapability;
use llm_worker::llm_client::error::ClientError;
use llm_worker::llm_client::event::Event;
use llm_worker::llm_client::retry::RetryPolicy;
use llm_worker::llm_client::scheme::Scheme;
use llm_worker::llm_client::transport::{HttpTransport, ResolvedAuth};
use llm_worker::llm_client::types::Request;
use serde_json::Value;
use std::time::Duration;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
/// SSE 本体は触らないテスト用 scheme。`parse_fail` を立てると
/// stream 消費中（= retry loop の外）で `ClientError::Sse` を返す。
/// stream 消費中で `ClientError::Sse` を返す。
#[derive(Clone)]
struct DummyScheme {
parse_fail: bool,
@@ -31,18 +26,23 @@ struct DummyScheme {
impl Scheme for DummyScheme {
type State = ();
fn default_base_url(&self) -> &'static str {
""
}
fn path(&self, _: &str) -> String {
"/v1/chat".into()
}
fn required_auth(&self) -> AuthRequirement {
AuthRequirement::None
}
fn build_request_body(&self, _: &str, _: &Request, _: &ModelCapability) -> Value {
serde_json::json!({})
}
fn parse_sse(&self, _: &str, _: &str, _: &mut ()) -> Result<Vec<Event>, ClientError> {
if self.parse_fail {
Err(ClientError::Sse(
@@ -52,25 +52,13 @@ impl Scheme for DummyScheme {
Ok(vec![])
}
}
fn default_capability(&self) -> ModelCapability {
ModelCapability::minimal()
}
}
fn fast_policy(max_attempts: u32) -> RetryPolicy {
RetryPolicy {
base: Duration::from_millis(1),
cap: Duration::from_millis(1),
max_attempts,
total_timeout: Duration::from_secs(60),
}
}
fn build_transport(
base_url: impl Into<String>,
parse_fail: bool,
policy: RetryPolicy,
) -> HttpTransport<DummyScheme> {
fn build_transport(base_url: impl Into<String>, parse_fail: bool) -> HttpTransport<DummyScheme> {
HttpTransport::new(
DummyScheme { parse_fail },
"test-model",
@@ -78,7 +66,6 @@ fn build_transport(
ResolvedAuth::None,
ModelCapability::minimal(),
)
.with_retry_policy(policy)
}
fn ok_sse() -> ResponseTemplate {
@@ -88,78 +75,11 @@ fn ok_sse() -> ResponseTemplate {
}
#[tokio::test]
async fn retries_503_then_succeeds() {
async fn retryable_status_returns_api_error_without_retrying() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/chat"))
.respond_with(ResponseTemplate::new(503).set_body_string("upstream connect error"))
.up_to_n_times(2)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/v1/chat"))
.respond_with(ok_sse())
.mount(&server)
.await;
let transport = build_transport(server.uri(), false, fast_policy(5));
let mut stream = transport
.stream(Request::default())
.await
.expect("stream should succeed after retries");
while stream.next().await.is_some() {}
let received = server.received_requests().await.unwrap();
assert_eq!(received.len(), 3, "two failures plus one success expected");
}
#[tokio::test]
async fn retries_529_then_exhausts() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/chat"))
.respond_with(ResponseTemplate::new(529).set_body_string("overloaded"))
.mount(&server)
.await;
let transport = build_transport(server.uri(), false, fast_policy(3));
match transport.stream(Request::default()).await {
Err(ClientError::Api {
status: Some(529), ..
}) => {}
Err(other) => panic!("expected Api(529), got {other:?}"),
Ok(_) => panic!("expected error after exhausting retries"),
}
let received = server.received_requests().await.unwrap();
assert_eq!(received.len(), 3, "should hit max_attempts and stop");
}
#[tokio::test]
async fn connect_refused_retries_then_fails() {
// 接続不能なローカルアドレスを使う。Linux では `Connection refused` で
// 即時失敗するため、`fast_policy` ならテストが秒以下で終わる。
let unreachable = "http://127.0.0.1:1";
let transport = build_transport(unreachable, false, fast_policy(3));
match transport.stream(Request::default()).await {
Err(ClientError::Http(e)) => {
assert!(
e.is_connect() || e.is_timeout(),
"expected connect/timeout, got {e:?}"
);
}
Err(other) => panic!("expected Http error, got {other:?}"),
Ok(_) => panic!("expected error connecting to closed port"),
}
}
#[tokio::test]
async fn retry_after_header_overrides_backoff() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/chat"))
.respond_with(ResponseTemplate::new(503).insert_header("retry-after", "1"))
.up_to_n_times(1)
.mount(&server)
.await;
@@ -169,34 +89,48 @@ async fn retry_after_header_overrides_backoff() {
.mount(&server)
.await;
// base/cap を 1ms に絞った policy で `Retry-After: 1` を観察すると、
// 指数バックオフ単独なら 1ms 程度で終わるはずが Retry-After に従って
// 1 秒待つ → 経過時間で override を検証できる。
let policy = RetryPolicy {
base: Duration::from_millis(1),
cap: Duration::from_millis(1),
max_attempts: 3,
total_timeout: Duration::from_secs(10),
};
let transport = build_transport(server.uri(), false, policy);
let transport = build_transport(server.uri(), false);
match transport.stream(Request::default()).await {
Err(ClientError::Api {
status: Some(503), ..
}) => {}
Err(other) => panic!("expected Api(503), got {other:?}"),
Ok(_) => panic!("transport must not retry internally"),
}
let start = Instant::now();
let mut stream = transport.stream(Request::default()).await.expect("ok");
while stream.next().await.is_some() {}
let elapsed = start.elapsed();
assert!(
elapsed >= Duration::from_secs(1),
"Retry-After=1 should make us wait >=1s, elapsed={elapsed:?}"
);
assert!(
elapsed < Duration::from_secs(3),
"Retry-After=1 should not balloon, elapsed={elapsed:?}"
let received = server.received_requests().await.unwrap();
assert_eq!(
received.len(),
1,
"transport should send exactly one request"
);
}
#[tokio::test]
async fn mid_stream_sse_error_does_not_retry() {
async fn retry_after_header_is_preserved_on_api_error() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/chat"))
.respond_with(ResponseTemplate::new(503).insert_header("retry-after", "1"))
.mount(&server)
.await;
let transport = build_transport(server.uri(), false);
match transport.stream(Request::default()).await {
Err(
err @ ClientError::Api {
status: Some(503), ..
},
) => {
assert_eq!(err.retry_after(), Some(Duration::from_secs(1)));
}
Err(other) => panic!("expected Api(503), got {other:?}"),
Ok(_) => panic!("expected error"),
}
}
#[tokio::test]
async fn mid_stream_sse_error_is_stream_item_error() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/chat"))
@@ -211,11 +145,11 @@ async fn mid_stream_sse_error_does_not_retry() {
.mount(&server)
.await;
let transport = build_transport(server.uri(), true, fast_policy(5));
let transport = build_transport(server.uri(), true);
let mut stream = transport
.stream(Request::default())
.await
.expect("status 200 should bypass retry loop");
.expect("status 200 should open stream");
let mut saw_sse_err = false;
while let Some(item) = stream.next().await {
if matches!(item, Err(ClientError::Sse(_))) {
@@ -225,11 +159,11 @@ async fn mid_stream_sse_error_does_not_retry() {
assert!(saw_sse_err, "expected Sse error from stream consumer");
let received = server.received_requests().await.unwrap();
assert_eq!(received.len(), 1, "mid-stream Sse must not retry");
assert_eq!(received.len(), 1, "mid-stream Sse must not reopen stream");
}
#[tokio::test]
async fn non_retryable_status_returns_immediately() {
async fn non_retryable_status_returns_api_error() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/chat"))
@@ -237,7 +171,7 @@ async fn non_retryable_status_returns_immediately() {
.mount(&server)
.await;
let transport = build_transport(server.uri(), false, fast_policy(5));
let transport = build_transport(server.uri(), false);
match transport.stream(Request::default()).await {
Err(ClientError::Api {
status: Some(401), ..
@@ -247,5 +181,5 @@ async fn non_retryable_status_returns_immediately() {
}
let received = server.received_requests().await.unwrap();
assert_eq!(received.len(), 1, "401 must not retry");
assert_eq!(received.len(), 1);
}