From 7c4390c20994f8ad02aca33f62839f355cdbd382 Mon Sep 17 00:00:00 2001 From: Rayeed Zaman Date: Sun, 8 Nov 2026 22:13:12 -0500 Subject: [PATCH] Throttle output webhooks and skip building throttled payloads Coglet sent an `output` webhook for every yielded chunk: should_send let Output events skip the throttle. Each payload is a full snapshot with all logs and outputs so far, so a model that logs 1 MB and yields 1,000 chunks sent about 1,000 webhooks and 1 GB. The snapshot was also built before the filter and throttle check, so events that were then dropped still cloned all logs and outputs. `output` events now share the throttle with `logs`, as docs/http.md and COG_THROTTLE_RESPONSE_INTERVAL describe and as the Python and Go runtimes did. Prediction checks should_send before building the payload. A dropped event loses no data: the next webhook, or the terminal one, has the full state. Fixes #3219 Co-Authored-By: Claude Opus 5.5 --- crates/coglet/src/prediction.rs | 80 ++++++++++++++++++++++++++++++++- crates/coglet/src/webhook.rs | 62 ++++++++++++++++++------- 2 files changed, 126 insertions(+), 16 deletions(-) diff --git a/crates/coglet/src/prediction.rs b/crates/coglet/src/prediction.rs index 3e4ada3248..7274673a59 100644 --- a/crates/coglet/src/prediction.rs +++ b/crates/coglet/src/prediction.rs @@ -553,8 +553,13 @@ impl Prediction { /// /// Builds the current state as a JSON payload and sends it via the /// stored WebhookSender. Spawns a tokio task — does not block. + /// + /// The payload holds all logs and outputs so far, so it's only built if + /// the event passes the filter and throttle. fn fire_webhook(&self, event: WebhookEventType) { - if let Some(ref webhook) = self.webhook { + if let Some(ref webhook) = self.webhook + && webhook.should_send(event) + { let payload = self.build_webhook_payload(); webhook.send(event, &payload); } @@ -730,6 +735,79 @@ mod tests { assert_eq!(pred.outputs().len(), 2); } + /// Regression test for #3219: a streaming prediction used to send one + /// `output` webhook per chunk, each carrying the full logs. + #[tokio::test] + async fn output_webhooks_are_throttled_and_terminal_has_everything() { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + for status in ["succeeded", "failed", "canceled"] { + crate::install_crypto_provider(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + let config = crate::webhook::WebhookConfig { + response_interval: std::time::Duration::from_secs(10), + ..Default::default() + }; + let webhook = WebhookSender::new(format!("{}/webhook", server.uri()), config).unwrap(); + let mut pred = Prediction::new("pred_3219".to_string(), Some(webhook)); + + pred.set_processing(); + let big_log = "A".repeat(100_000); + pred.append_log(&big_log); + for _ in 0..100 { + pred.append_output(serde_json::json!("tok ")); + } + match status { + "succeeded" => { + let outputs = pred.take_outputs(); + pred.set_succeeded(PredictionOutput::Stream(outputs)); + } + "failed" => pred.set_failed("boom".to_string()), + _ => pred.set_canceled(), + } + + let payloads = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let payloads: Vec = server + .received_requests() + .await + .unwrap() + .iter() + .map(|r| serde_json::from_slice(&r.body).unwrap()) + .collect(); + if payloads.iter().any(|p| p["status"] == status) { + return payloads; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{status}: terminal webhook was not received")); + + // The start webhook opens the throttle window; every output and + // log event inside it is dropped. Only the terminal one follows. + assert_eq!( + payloads.len(), + 2, + "{status}: got {} webhooks", + payloads.len() + ); + assert!(payloads.iter().any(|p| p["status"] == "processing")); + let terminal = payloads.iter().find(|p| p["status"] == status).unwrap(); + assert_eq!(terminal["logs"], big_log, "{status}"); + assert_eq!( + terminal["output"].as_array().map(Vec::len), + Some(100), + "{status}" + ); + } + } + #[tokio::test] async fn prediction_stream_emits_start_output_log_and_completed() { let mut prediction = Prediction::new("pred_stream".to_string(), None); diff --git a/crates/coglet/src/webhook.rs b/crates/coglet/src/webhook.rs index a736607b51..24383ca1f6 100644 --- a/crates/coglet/src/webhook.rs +++ b/crates/coglet/src/webhook.rs @@ -137,7 +137,16 @@ impl WebhookSender { &self.url } - fn should_send(&self, event: WebhookEventType) -> bool { + /// Whether a non-terminal `send` for this event would go out now: the event + /// passes the filter and the throttle interval has elapsed. + /// + /// Callers can check this before building the payload, which holds all + /// logs and outputs so far. + /// + /// `output` and `logs` events share the throttle, as in the Python and Go + /// runtimes. Payloads are full snapshots, so a dropped event loses no data: + /// the next webhook, or the terminal one, has it. + pub(crate) fn should_send(&self, event: WebhookEventType) -> bool { if !self.config.events_filter.contains(&event) { return false; } @@ -146,15 +155,6 @@ impl WebhookSender { return true; } - // Output events are never throttled: they are high-value (contain actual - // prediction results), relatively infrequent (one per output chunk/file), - // and in the old Python runtime were effectively unthrottled because file - // uploads were synchronous. Throttling them causes the director to miss - // intermediate output data. - if matches!(event, WebhookEventType::Output) { - return true; - } - // Recover from poison - losing throttle state is acceptable let last = self.last_sent.lock().unwrap_or_else(|e| e.into_inner()); last.elapsed() >= self.config.response_interval @@ -569,13 +569,13 @@ mod tests { } #[tokio::test] - async fn send_output_not_throttled() { + async fn send_output_throttled() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/webhook")) .respond_with(ResponseTemplate::new(200)) - .expect(2) + .expect(1) .mount(&server) .await; @@ -586,19 +586,51 @@ mod tests { }; let sender = WebhookSender::new(url, config).unwrap(); - // Output events bypass throttling — both should be sent sender.send( WebhookEventType::Output, - &serde_json::json!({"output": "1"}), + &serde_json::json!({"output": ["1"]}), ); + // Second send should be throttled sender.send( WebhookEventType::Output, - &serde_json::json!({"output": "2"}), + &serde_json::json!({"output": ["1", "2"]}), ); tokio::time::sleep(Duration::from_millis(50)).await; } + #[test] + fn should_send_throttles_output_and_logs_together() { + let config = WebhookConfig { + response_interval: Duration::from_secs(10), + ..test_config() + }; + let sender = WebhookSender::new("http://localhost/webhook".to_string(), config).unwrap(); + + assert!(sender.should_send(WebhookEventType::Output)); + sender.update_last_sent(); + + assert!(!sender.should_send(WebhookEventType::Output)); + assert!(!sender.should_send(WebhookEventType::Logs)); + assert!(sender.should_send(WebhookEventType::Completed)); + } + + #[test] + fn should_send_respects_filter() { + let config = WebhookConfig { + events_filter: [WebhookEventType::Start, WebhookEventType::Completed] + .into_iter() + .collect(), + ..test_config() + }; + let sender = WebhookSender::new("http://localhost/webhook".to_string(), config).unwrap(); + + assert!(sender.should_send(WebhookEventType::Start)); + assert!(!sender.should_send(WebhookEventType::Output)); + assert!(!sender.should_send(WebhookEventType::Logs)); + assert!(sender.should_send(WebhookEventType::Completed)); + } + #[tokio::test] async fn send_terminal_sync_posts_json() { let server = MockServer::start().await;