diff --git a/Cargo.lock b/Cargo.lock index d1e5cb9..06185f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -81,8 +81,7 @@ dependencies = [ [[package]] name = "agent-client-protocol-derive" version = "2.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4893bfb11e0e193aa5a64611a126ee87567acb1879556519c902722899229de" +source = "git+https://github.com/danielkov/rust-sdk?rev=03eab519aa5ecfc97a43bf17f9721c38c065c47b#03eab519aa5ecfc97a43bf17f9721c38c065c47b" dependencies = [ "quote", "syn 3.0.4", @@ -109,8 +108,7 @@ dependencies = [ [[package]] name = "agentkit-acp-http" version = "2.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dac1fe9230b43784cb66312e93eb72d6420c0a027029879feb02361435856b66" +source = "git+https://github.com/danielkov/rust-sdk?rev=03eab519aa5ecfc97a43bf17f9721c38c065c47b#03eab519aa5ecfc97a43bf17f9721c38c065c47b" dependencies = [ "agentkit-acp-sdk", "async-stream", @@ -147,8 +145,7 @@ dependencies = [ [[package]] name = "agentkit-acp-sdk" version = "2.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b93c10bb458dd1c2c822275a7e1c390b46846a6bcd3cbc84aafa312563e36ad" +source = "git+https://github.com/danielkov/rust-sdk?rev=03eab519aa5ecfc97a43bf17f9721c38c065c47b#03eab519aa5ecfc97a43bf17f9721c38c065c47b" dependencies = [ "agent-client-protocol-derive", "agentkit-acp-schema", diff --git a/Cargo.toml b/Cargo.toml index 4e607a8..e76d718 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -139,3 +139,8 @@ todo = "deny" unimplemented = "deny" disallowed_methods = "deny" disallowed_macros = "deny" + +# Bounded SSE replay fairness; keep HTTP and core channel types on one revision. +[patch.crates-io] +agentkit-acp-http = { git = "https://github.com/danielkov/rust-sdk", rev = "03eab519aa5ecfc97a43bf17f9721c38c065c47b" } +agentkit-acp-sdk = { git = "https://github.com/danielkov/rust-sdk", rev = "03eab519aa5ecfc97a43bf17f9721c38c065c47b" } diff --git a/src/gateway/http_client/recovery.rs b/src/gateway/http_client/recovery.rs index 810ea57..b9eb499 100644 --- a/src/gateway/http_client/recovery.rs +++ b/src/gateway/http_client/recovery.rs @@ -316,8 +316,9 @@ pub(super) async fn run( Timeout, } let event = { - // select polls its left branch first. Keep timeout > driver > - // inbound frame > stdin priority; disabled branches never poll + // select polls its left branch first. Drain admitted inbound + // frames before driving more HTTP work: timeout > inbound frame + // > driver > stdin. Disabled branches never poll // their underlying receiver or deadline. let timeout = std::pin::pin!(async { if phase != "live" { @@ -340,11 +341,11 @@ pub(super) async fn run( std::future::pending().await } }); - match select(timeout, select(&mut transport, select(frame, line))).await { + match select(timeout, select(frame, select(&mut transport, line))).await { Either::Left(_) => Event::Timeout, - Either::Right((Either::Left((result, _)), _)) => Event::Driver(result), - Either::Right((Either::Right((Either::Left((frame, _)), _)), _)) => { - Event::Frame(frame) + Either::Right((Either::Left((frame, _)), _)) => Event::Frame(frame), + Either::Right((Either::Right((Either::Left((result, _)), _)), _)) => { + Event::Driver(result) } Either::Right((Either::Right((Either::Right((line, _)), _)), _)) => { Event::Line(line) diff --git a/src/gateway/http_server.rs b/src/gateway/http_server.rs index 211bf55..7ee3e14 100644 --- a/src/gateway/http_server.rs +++ b/src/gateway/http_server.rs @@ -5,6 +5,10 @@ use agent_client_protocol::{ }; use futures_util::StreamExt; +// One budget for a response and all controller polling/output in a drain. +// Progress does not restart it or renew leases while another controller waits. +const OUTPUT_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); + pub(super) struct Connection { gateway: Arc, initialize: Option<(Value, Arc)>, @@ -374,12 +378,41 @@ impl Connection { } async fn flush(&mut self, channel: &BoundedChannel) -> agent_client_protocol::Result<()> { + self.flush_with_response(channel, None).await + } + + async fn flush_with_response( + &mut self, + channel: &BoundedChannel, + response: Option, + ) -> agent_client_protocol::Result<()> { + let deadline = tokio::time::Instant::now() + OUTPUT_DRAIN_TIMEOUT; + let drain = async { + if let Some(response) = response { + send(channel, response, deadline).await?; + } + self.flush_before(channel, deadline).await + }; + let timeout = std::pin::pin!(tokio::time::sleep_until(deadline)); + let drain = std::pin::pin!(drain); + // Poll the deadline first, including when actor output is already ready. + match select(timeout, drain).await { + Either::Left(_) => Err(channel.tx.fail("gateway HTTP output admission stalled")), + Either::Right((result, _)) => result, + } + } + + async fn flush_before( + &mut self, + channel: &BoundedChannel, + deadline: tokio::time::Instant, + ) -> agent_client_protocol::Result<()> { let mut failed = Vec::new(); for (id, control) in &mut self.controls { match control.poll().await { Ok(messages) => { for message in messages { - send(channel, message)?; + send(channel, message, deadline).await?; } } Err(error) => failed.push((id.clone(), error)), @@ -408,7 +441,9 @@ impl Connection { ]), ), ]), - )?; + deadline, + ) + .await?; for (_, request_id) in control.pending.drain() { send( channel, @@ -424,7 +459,9 @@ impl Connection { ]), ), ]), - )?; + deadline, + ) + .await?; } } } @@ -493,10 +530,7 @@ impl Connection { ])) } }; - if let Some(response) = response { - send(&channel, response)?; - } - self.flush(&channel).await?; + self.flush_with_response(&channel, response).await?; } drop(charged_frame); } @@ -510,14 +544,25 @@ impl Connection { fn sdk_error(error: impl std::fmt::Display) -> agent_client_protocol::Error { agent_client_protocol::Error::internal_error().data(error.to_string()) } -fn send(channel: &BoundedChannel, message: Value) -> agent_client_protocol::Result<()> { +async fn send( + channel: &BoundedChannel, + message: Value, + deadline: tokio::time::Instant, +) -> agent_client_protocol::Result<()> { + // Ready sends must also respect the deadline; timeout_at alone can finish + // a ready future even when its timer has already elapsed. + if tokio::time::Instant::now() >= deadline { + return Err(channel.tx.fail("gateway HTTP output admission stalled")); + } let message = serde_json::from_value(message).map_err(|error| channel.tx.fail(&error.to_string()))?; // Wake the terminal failure observer even if request setup converts this // SDK error to Failure. Saturation never silently drops an accepted frame. - channel - .tx - .try_send(TransportFrame::Single(message)) + // Admission stays charged through SDK routing and HTTP body handoff, not + // merely dequeue. Wait for actual release instead of relying on scheduling. + tokio::time::timeout_at(deadline, channel.tx.send(TransportFrame::Single(message))) + .await + .map_err(|_| channel.tx.fail("gateway HTTP output admission stalled"))? .map_err(|error| channel.tx.fail(&error.to_string())) } @@ -547,6 +592,17 @@ mod tests { use super::*; use serde_json::json; + // Single-send tests use the same production deadline path. Connection tests + // below exercise sharing that deadline across complete output drains. + async fn send(channel: &BoundedChannel, message: Value) -> agent_client_protocol::Result<()> { + super::send( + channel, + message, + tokio::time::Instant::now() + OUTPUT_DRAIN_TIMEOUT, + ) + .await + } + fn limits() -> ChannelLimits { ChannelLimits { max_frame_bytes: 1024, @@ -583,33 +639,67 @@ mod tests { assert!(error.to_string().contains("initialize is required")); } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn dequeued_frame_retains_admission_until_dropped() { let (channel, mut peer) = BoundedChannel::duplex(limits()).unwrap(); let message = json!({"jsonrpc":"2.0","id":1,"result":{}}); - send(&channel, message.clone()).unwrap(); + send(&channel, message.clone()).await.unwrap(); let charged = peer.rx.next().await.unwrap(); let decoded = charged.decode(); assert!(matches!(decoded, TransportFrame::Single(_))); // Removing a frame from the queue does not release its reservation. - let rejected = send(&channel, message.clone()).unwrap_err(); + let rejected = send(&channel, message.clone()).await.unwrap_err(); + assert!(rejected.to_string().contains("output admission stalled")); assert_eq!(channel.tx.failure().await.to_string(), rejected.to_string()); drop(decoded); drop(charged); assert!( - send(&channel, message).is_err(), + send(&channel, message).await.is_err(), "saturation stays terminal" ); } + #[tokio::test(start_paused = true)] + async fn output_waits_for_http_body_charge_release_before_stall_deadline() { + let (channel, mut peer) = BoundedChannel::duplex(limits()).unwrap(); + let message = json!({"jsonrpc":"2.0","id":1,"result":{}}); + send(&channel, message.clone()).await.unwrap(); + // HTTP routing/body ownership can retain a dequeued frame. Progress is + // tied to that ownership ending, not a producer yield or queue dequeue. + let held = Arc::new(peer.rx.next().await.unwrap()); + let body = held.clone(); + drop(held); + let mut sending = Box::pin(send(&channel, message)); + assert!(futures_util::poll!(&mut sending).is_pending()); + tokio::time::advance(OUTPUT_DRAIN_TIMEOUT / 2).await; + assert!(futures_util::poll!(&mut sending).is_pending()); + drop(body); + sending.await.unwrap(); + assert!(peer.rx.next().await.is_some()); + } + + #[tokio::test] + async fn cancelled_output_wait_releases_registration_without_failing_channel() { + let (channel, mut peer) = BoundedChannel::duplex(limits()).unwrap(); + let message = json!({"jsonrpc":"2.0","id":1,"result":{}}); + send(&channel, message.clone()).await.unwrap(); + let held = peer.rx.next().await.unwrap(); + let mut sending = Box::pin(send(&channel, message.clone())); + assert!(futures_util::poll!(&mut sending).is_pending()); + drop(sending); + drop(held); + send(&channel, message).await.unwrap(); + assert!(peer.rx.next().await.is_some()); + } + #[tokio::test] async fn dropping_consumed_frame_releases_current_admission() { let (channel, mut peer) = BoundedChannel::duplex(limits()).unwrap(); let message = json!({"jsonrpc":"2.0","id":1,"result":{}}); - send(&channel, message.clone()).unwrap(); + send(&channel, message.clone()).await.unwrap(); let charged = peer.rx.next().await.unwrap(); drop(charged); - send(&channel, message).unwrap(); + send(&channel, message).await.unwrap(); assert!(peer.rx.next().await.is_some()); } @@ -619,12 +709,146 @@ mod tests { let message = json!({"jsonrpc":"2.0","id":1,"result":{"text":"x".repeat(64 * 1024)}}); let mut total = 0; while total <= MAX_REPLAY { - send(&channel, message.clone()).unwrap(); + send(&channel, message.clone()).await.unwrap(); let charged = peer.rx.next().await.unwrap(); total += charged.as_bytes().len(); drop(charged); } - send(&channel, message).unwrap(); + send(&channel, message).await.unwrap(); + } + + async fn assert_slow_progress_has_one_drain_deadline( + batches: &[usize], + poll_delay: Duration, + with_response: bool, + ) { + let directory = tempfile::tempdir().unwrap(); + let token = directory.path().join("token"); + std::fs::write(&token, "test-token").unwrap(); + let gateway = Arc::new(Gateway { + stopping: AtomicBool::new(false), + token: BearerToken::load(&token).unwrap(), + roots: Vec::new(), + sessions: Mutex::new(HashMap::new()), + }); + let mut connection = Connection::new(gateway); + let mut actors = Vec::new(); + for (index, count) in batches.iter().copied().enumerate() { + let session = format!("session-{index}"); + let (sender, mut receiver) = mpsc::channel::(2); + let (_stopped, stopped) = watch::channel(()); + connection.controls.insert( + session.clone(), + Control { + session, + attachment: "current".into(), + entry: Entry { + root: PathBuf::new(), + sender, + stopped, + }, + cursor: 0, + pending: HashMap::new(), + }, + ); + // The resident actor mailbox is the real adapter boundary. Each + // controller contributes a batch to the same HTTP output channel. + actors.push(tokio::spawn(async move { + let envelope = receiver.recv().await.unwrap(); + assert!(matches!(envelope.request, Request::Poll { .. })); + tokio::time::sleep(poll_delay).await; + let events = (1..=count) + .map(|cursor| { + json!({ + "cursor":cursor, + "message":{"jsonrpc":"2.0","id":cursor,"result":{}} + }) + }) + .collect::>(); + let _ = envelope.reply.send(Ok(json!({"events":events}))); + })); + } + let (channel, mut peer) = BoundedChannel::duplex(limits()).unwrap(); + let response = with_response.then(|| json!({"jsonrpc":"2.0","id":0,"result":{}})); + if with_response { + // A prior HTTP body owns the only slot when the response arrives. + send(&channel, json!({"jsonrpc":"2.0","id":-1,"result":{}})) + .await + .unwrap(); + } + let started = tokio::time::Instant::now(); + let output = async { + let result = if response.is_some() { + connection.flush_with_response(&channel, response).await + } else { + connection.flush(&channel).await + }; + channel.tx.close_channel(); + result + }; + let consumer = async { + while let Some(charged) = peer.rx.next().await { + // Every individual frame makes progress well before five + // seconds, but the complete batch exceeds the drain budget. + tokio::time::sleep(Duration::from_secs(1)).await; + drop(charged); + } + }; + let (result, ()) = tokio::join!(output, consumer); + for actor in actors { + actor.abort(); + let _ = actor.await; + } + let error = result.expect_err("slow progress reset the output drain deadline"); + assert!( + error.to_string().contains("output admission stalled"), + "{error}" + ); + assert!(started.elapsed() <= OUTPUT_DRAIN_TIMEOUT + Duration::from_secs(1)); + assert!(started.elapsed() < LEASE); + assert_eq!( + connection.controls.len(), + batches.len(), + "drain timeout must not be reported as controller replacement" + ); + assert_eq!(channel.tx.failure().await.to_string(), error.to_string()); + } + + #[tokio::test(start_paused = true)] + async fn slow_progress_cannot_extend_one_controller_output_drain() { + assert_slow_progress_has_one_drain_deadline(&[128], Duration::ZERO, false).await; + } + + #[tokio::test(start_paused = true)] + async fn slow_progress_cannot_reset_output_drain_between_controllers() { + // Each controller fits individually inside the budget. Only a shared + // deadline across the multiplexed connection bounds the whole drain. + assert_slow_progress_has_one_drain_deadline(&[4, 4], Duration::ZERO, false).await; + } + + #[tokio::test(start_paused = true)] + async fn slow_progress_response_and_replay_share_output_drain_deadline() { + assert_slow_progress_has_one_drain_deadline(&[4], Duration::ZERO, true).await; + } + + #[tokio::test(start_paused = true)] + async fn slow_actor_polling_consumes_the_shared_output_drain_budget() { + assert_slow_progress_has_one_drain_deadline(&[0, 0, 0], Duration::from_secs(2), false) + .await; + } + + #[tokio::test(start_paused = true)] + async fn ready_output_cannot_bypass_expired_drain_deadline() { + let (channel, mut peer) = BoundedChannel::duplex(limits()).unwrap(); + let error = super::send( + &channel, + json!({"jsonrpc":"2.0","id":1,"result":{}}), + tokio::time::Instant::now(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("output admission stalled")); + assert!(peer.rx.next().await.is_none()); } #[tokio::test] @@ -686,7 +910,7 @@ mod tests { .unwrap(); } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn pending_request_capacity_preserves_response_lane() { let (sender, mut receiver) = mpsc::channel::(1); let (_stopped, stopped) = watch::channel(()); @@ -704,7 +928,9 @@ mod tests { .collect(), }; let (channel, mut peer) = BoundedChannel::duplex(limits()).unwrap(); - send(&channel, json!({"jsonrpc":"2.0","id":1,"result":{}})).unwrap(); + send(&channel, json!({"jsonrpc":"2.0","id":1,"result":{}})) + .await + .unwrap(); let charge = Arc::new(peer.rx.next().await.unwrap()); let pending = control.pending.len(); assert!( @@ -740,7 +966,11 @@ mod tests { drop(charge); // Actor acceptance does not release SDK admission while the accepted // child-bound request still owns decoded data. - assert!(send(&channel, json!({"jsonrpc":"2.0","id":2,"result":{}})).is_err()); + assert!( + send(&channel, json!({"jsonrpc":"2.0","id":2,"result":{}})) + .await + .is_err() + ); drop(accepted); } } diff --git a/tests/gateway.rs b/tests/gateway.rs index 11328ef..564259b 100644 --- a/tests/gateway.rs +++ b/tests/gateway.rs @@ -272,6 +272,7 @@ struct Bridge { impl Bridge { async fn request(&mut self, id: u64, method: &str, params: Value) -> Value { + let started = std::time::Instant::now(); let mut bytes = serde_json::to_vec( &json!({"jsonrpc":"2.0", "id":id, "method":method, "params":params}), ) @@ -291,7 +292,12 @@ impl Bridge { self.updates.push(message["params"]["update"].clone()); } if message["id"] == id { - assert!(message.get("error").is_none(), "{message}"); + assert!( + message.get("error").is_none(), + "{message}; {method} failed after {:?}, received {} updates", + started.elapsed(), + self.updates.len() + ); return message["result"].clone(); } } @@ -2565,3 +2571,71 @@ async fn assert_no_replay_snapshot(bridge: &mut Bridge, id: &str, state: &str) { #[path = "gateway_background_cancel/mod.rs"] mod background_cancel; + +// Seed durable history only after every writer has exited. No prompt or model +// request is needed: replay must work even when inference is unavailable. +async fn assert_bridge_replays_ready_history(records: usize) { + let fixture = Fixture::new(); + let (mut gateway, url) = fixture.gateway().await; + let mut bridge = fixture.bridge(&url, None); + let id = bridge.handshake().await; + bridge.detach().await; + stop_gateway(&mut gateway).await; + let path = fs::read_dir(fixture.home.path().join(".kit/sessions")) + .unwrap() + .flat_map(|workspace| fs::read_dir(workspace.unwrap().path()).unwrap()) + .filter_map(Result::ok) + .find(|entry| entry.file_name() == format!("{id}.jsonl").as_str()) + .unwrap() + .path(); + let mut transcript = fs::read_to_string(&path).unwrap(); + let mut record: Value = serde_json::from_str(transcript.lines().next().unwrap()).unwrap(); + let generation = transcript + .lines() + .map(|line| { + serde_json::from_str::(line).unwrap()["generation"] + .as_u64() + .unwrap() + }) + .max() + .unwrap(); + for index in 0..records { + record["generation"] = json!(generation + index as u64 + 1); + record["item"] = json!({ + "kind":"Assistant", "id":format!("replay-{index}"), + "created_at":1_700_000_000_000u64 + index as u64, + "finish_reason":"Completed", "metadata":{}, "usage":null, + "parts":[{"Text":{"text":format!("History entry {index:04}"),"metadata":{}}}] + }); + transcript.push_str(&serde_json::to_string(&record).unwrap()); + transcript.push('\n'); + } + fs::write(path, transcript).unwrap(); + let (mut gateway, url) = fixture.gateway().await; + let mut bridge = fixture.bridge(&url, Some(&id)); + assert_eq!(bridge.handshake().await, id); + let expected = (0..records) + .map(|index| format!("History entry {index:04}")) + .collect::(); + assert_replayed_answer(&bridge.updates, &expected); + // Each assistant record expands into multiple wire events; this exceeds + // the incoming channel's fixed frame bound without increasing that bound. + assert!(bridge.updates.len() >= records); + bridge.detach().await; + stop_gateway(&mut gateway).await; +} + +#[tokio::test] +async fn bridge_replays_thirteen_short_history_records() { + assert_bridge_replays_ready_history(13).await; +} + +#[tokio::test] +async fn bridge_replays_160_short_history_records() { + assert_bridge_replays_ready_history(160).await; +} + +#[tokio::test] +async fn bridge_replays_1000_short_history_records() { + assert_bridge_replays_ready_history(1000).await; +}