From 9668728ba0d491531fba5cea09b303c571b870b0 Mon Sep 17 00:00:00 2001 From: davidabram Date: Sun, 11 Oct 2026 13:40:15 +0200 Subject: [PATCH 1/2] sync: Add tests for first-terminal-error precedence across streams Cover two concurrent trace-sync termination scenarios that were not exercised before. The first test runs three streams where the messages stream fails with Forbidden while a sibling stream also fails with AuthenticationFailed and a third stream still has backlog. It asserts that the first terminal error is reported and that neither sibling reads or ingests past its in-flight batch once termination is signalled. The second test lets a sibling batch fail ambiguously with a server error after termination, using a counting empty credential store. It asserts that the first Forbidden error is kept, the sibling stops after its in-flight batch, and credential state is loaded exactly once. Co-authored-by: SCE --- cli/src/services/sync/sync.rs | 216 +++++++++++++++++++++++++++++++++- 1 file changed, 215 insertions(+), 1 deletion(-) diff --git a/cli/src/services/sync/sync.rs b/cli/src/services/sync/sync.rs index a97e82681..ef03b382e 100644 --- a/cli/src/services/sync/sync.rs +++ b/cli/src/services/sync/sync.rs @@ -538,6 +538,8 @@ fn cursor_for_stream(cursors: &AgentTraceCursors, stream: IngestionStream) -> i6 mod tests { use std::cell::{Cell, RefCell}; use std::rc::Rc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; use std::time::Duration; use tokio::sync::oneshot; @@ -547,12 +549,15 @@ mod tests { use crate::services::agent_trace_export::AgentTraceMessageExportRow; use crate::services::agent_trace_sync::control_plane::{ AgentTraceIngestionBatchResponse, AuthenticatedControlPlaneClient, ControlPlaneError, - IngestionStream, + CredentialStore, IngestionStream, }; use crate::services::agent_trace_sync::StreamSyncError; + use crate::services::auth::TokenResponse; use crate::services::sync::progress::NoopProgressReporter; + use crate::services::token_storage::StoredTokens; const NOT_COMPLETED_PROBE: Duration = Duration::from_millis(50); + const BACKLOG_ROWS: i64 = 5; #[tokio::test] async fn sibling_failure_does_not_cancel_in_flight_credential_save() { @@ -767,6 +772,215 @@ mod tests { assert!(terminated.get()); assert_eq!(sibling_ingests.get(), 0); } + + fn unreachable_client() -> AuthenticatedControlPlaneClient { + AuthenticatedControlPlaneClient::new( + reqwest::Client::builder() + .tls_certs_only(Vec::new()) + .build() + .expect("test http client"), + "http://127.0.0.1:1", + "http://127.0.0.1:1", + "client", + ) + } + + #[tokio::test] + async fn first_terminal_error_wins_when_sibling_also_fails_and_third_stream_stops_backlog() { + let client = unreachable_client(); + let mut reporter = NoopProgressReporter; + let progress = Rc::new(RefCell::new(&mut reporter)); + let terminated = Cell::new(false); + let in_flight = Cell::new(0usize); + let second_reads = Cell::new(0usize); + let second_ingests = Cell::new(0usize); + let third_reads = Cell::new(0usize); + let third_ingests = Cell::new(0usize); + + let failing = sync_one_stream( + &client, + "repo", + "source", + IngestionStream::Messages, + 0, + "messages", + async |cursor, _limit| Ok(vec![message_row(cursor + 1)]), + async |_request| { + while in_flight.get() < 2 { + tokio::task::yield_now().await; + } + Err(ControlPlaneError::Forbidden("denied".to_string())) + }, + Rc::clone(&progress), + &terminated, + ); + let second_failing = sync_one_stream( + &client, + "repo", + "source", + IngestionStream::Parts, + 0, + "parts", + async |cursor, _limit| { + second_reads.set(second_reads.get() + 1); + Ok((cursor < BACKLOG_ROWS) + .then(|| message_row(cursor + 1)) + .into_iter() + .collect()) + }, + async |_request| { + second_ingests.set(second_ingests.get() + 1); + in_flight.set(in_flight.get() + 1); + while !terminated.get() { + tokio::task::yield_now().await; + } + Err(ControlPlaneError::AuthenticationFailed( + "expired".to_string(), + )) + }, + Rc::clone(&progress), + &terminated, + ); + let backlogged = sync_one_stream( + &client, + "repo", + "source", + IngestionStream::AgentTraces, + 0, + "agent_traces", + async |cursor, _limit| { + third_reads.set(third_reads.get() + 1); + Ok((cursor < BACKLOG_ROWS) + .then(|| message_row(cursor + 1)) + .into_iter() + .collect()) + }, + async |request| { + third_ingests.set(third_ingests.get() + 1); + in_flight.set(in_flight.get() + 1); + while !terminated.get() { + tokio::task::yield_now().await; + } + Ok(AgentTraceIngestionBatchResponse { + accepted: request.rows.len(), + cursor: request.rows[0].source_row_id, + }) + }, + Rc::clone(&progress), + &terminated, + ); + + let result = join_three_to_completion(failing, second_failing, backlogged).await; + + assert!(matches!( + result, + Err(TraceSyncError::Stream { + stream: "messages", + source: StreamSyncError::Terminal(ControlPlaneError::Forbidden(_)), + }) + )); + assert_eq!(second_reads.get(), 1); + assert_eq!(second_ingests.get(), 1); + assert_eq!(third_reads.get(), 1); + assert_eq!(third_ingests.get(), 1); + } + + struct CountingEmptyCredentialStore { + loads: Arc, + } + + impl CredentialStore for CountingEmptyCredentialStore { + async fn load(&self) -> Result, ControlPlaneError> { + self.loads.fetch_add(1, Ordering::SeqCst); + Ok(None) + } + + async fn save(&self, _token: &TokenResponse) -> Result { + unreachable!("no refresh token is ever loaded") + } + } + + #[tokio::test] + async fn ambiguous_sibling_batch_after_termination_keeps_first_error_after_one_state_attempt() { + let loads = Arc::new(AtomicUsize::new(0)); + let client = AuthenticatedControlPlaneClient::with_credential_store( + reqwest::Client::builder() + .tls_certs_only(Vec::new()) + .build() + .expect("test http client"), + "http://127.0.0.1:1", + "http://127.0.0.1:1", + "client", + CountingEmptyCredentialStore { + loads: Arc::clone(&loads), + }, + ); + let mut reporter = NoopProgressReporter; + let progress = Rc::new(RefCell::new(&mut reporter)); + let terminated = Cell::new(false); + let sibling_ingest_started = Cell::new(false); + let sibling_reads = Cell::new(0usize); + let sibling_ingests = Cell::new(0usize); + + let failing = sync_one_stream( + &client, + "repo", + "source", + IngestionStream::Messages, + 0, + "messages", + async |cursor, _limit| Ok(vec![message_row(cursor + 1)]), + async |_request| { + while !sibling_ingest_started.get() { + tokio::task::yield_now().await; + } + Err(ControlPlaneError::Forbidden("denied".to_string())) + }, + Rc::clone(&progress), + &terminated, + ); + let sibling = sync_one_stream( + &client, + "repo", + "source", + IngestionStream::Parts, + 0, + "parts", + async |cursor, _limit| { + sibling_reads.set(sibling_reads.get() + 1); + Ok((cursor < BACKLOG_ROWS) + .then(|| message_row(cursor + 1)) + .into_iter() + .collect()) + }, + async |_request| { + sibling_ingests.set(sibling_ingests.get() + 1); + sibling_ingest_started.set(true); + while !terminated.get() { + tokio::task::yield_now().await; + } + Err::(ControlPlaneError::ServerError( + "unavailable".to_string(), + )) + }, + Rc::clone(&progress), + &terminated, + ); + let idle = async { Ok::<_, TraceSyncError>(()) }; + + let result = join_three_to_completion(failing, sibling, idle).await; + + assert!(matches!( + result, + Err(TraceSyncError::Stream { + stream: "messages", + source: StreamSyncError::Terminal(ControlPlaneError::Forbidden(_)), + }) + )); + assert_eq!(sibling_reads.get(), 1); + assert_eq!(sibling_ingests.get(), 1); + assert_eq!(loads.load(Ordering::SeqCst), 1); + } } #[cfg(test)] From 4b1ab3ca124d7d0f0b4fa616340caf9beaa584b6 Mon Sep 17 00:00:00 2001 From: davidabram Date: Sun, 11 Oct 2026 15:10:53 +0200 Subject: [PATCH 2/2] agent_trace_sync: Bound ingestion batch requests with a timeout Apply a one-minute timeout to ingestion batch requests so stalled requests and response bodies cannot indefinitely delay sync completion. Classify response-body timeouts as transport failures for cursor reconciliation. Add scripted HTTP regressions for stalled requests and bodies, retrying uncommitted batches, reconciling committed batches without resending, and reporting terminal failures after a sibling batch times out. Remove the standalone test HTTP server helper. Co-authored-by: SCE --- .gitignore | 2 +- .../agent_trace_sync/control_plane.rs | 30 +- .../agent_trace_sync/test_http_server.rs | 462 ----------------- cli/src/services/sync/sync.rs | 490 ++++++++++++++---- 4 files changed, 428 insertions(+), 556 deletions(-) delete mode 100644 cli/src/services/agent_trace_sync/test_http_server.rs diff --git a/.gitignore b/.gitignore index e84719570..123b8b2cf 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,6 @@ .direnv/ node_modules/ -cli/target/ +cli/target*/ cli/package-fallback/ integrations/**/target/ result* diff --git a/cli/src/services/agent_trace_sync/control_plane.rs b/cli/src/services/agent_trace_sync/control_plane.rs index ddc7d623c..fdd6d95ea 100644 --- a/cli/src/services/agent_trace_sync/control_plane.rs +++ b/cli/src/services/agent_trace_sync/control_plane.rs @@ -8,6 +8,7 @@ use std::fmt; use std::future::Future; use std::sync::Arc; +use std::time::Duration; use anyhow::anyhow; use reqwest::StatusCode; @@ -109,6 +110,8 @@ const STATE_PATH: &str = "agent-trace/ingestion/state"; const BATCH_PATH: &str = "agent-trace/ingestion/batch"; const ME_PATH: &str = "me"; +const BATCH_REQUEST_TIMEOUT: Duration = Duration::from_mins(1); + const STATE_RETRY_POLICY: RetryPolicy = RetryPolicy::builtin(1, 60_000, 250, 2_000); /// Typed failure classification for control-plane HTTP interactions, kept @@ -252,6 +255,7 @@ pub struct AuthenticatedControlPlaneClient { workos_client_id: String, credential_store: Arc, refresh_lock: Arc>, + batch_timeout: Duration, } /// Response body for `GET /me`. @@ -319,9 +323,16 @@ impl AuthenticatedControlPlaneClient { workos_client_id: workos_client_id.into(), credential_store: Arc::new(credential_store), refresh_lock: Arc::new(tokio::sync::Mutex::new(())), + batch_timeout: BATCH_REQUEST_TIMEOUT, } } + #[cfg(test)] + pub(crate) fn with_batch_timeout(mut self, batch_timeout: Duration) -> Self { + self.batch_timeout = batch_timeout; + self + } + pub(crate) async fn wait_for_pending_refresh(&self) { let _guard = self.refresh_lock.lock().await; } @@ -411,7 +422,13 @@ impl AuthenticatedControlPlaneClient { { let url = self.endpoint(BATCH_PATH); let response = self - .execute_authenticated(|token| self.http.post(&url).bearer_auth(token).json(request)) + .execute_authenticated(|token| { + self.http + .post(&url) + .bearer_auth(token) + .timeout(self.batch_timeout) + .json(request) + }) .await?; classify_response(response).await } @@ -582,10 +599,13 @@ where { let status = response.status(); if status.is_success() { - return response - .json::() - .await - .map_err(|error| ControlPlaneError::InvalidResponse(error.to_string())); + return response.json::().await.map_err(|error| { + if error.is_timeout() { + ControlPlaneError::Transport(error.to_string()) + } else { + ControlPlaneError::InvalidResponse(error.to_string()) + } + }); } let body = response diff --git a/cli/src/services/agent_trace_sync/test_http_server.rs b/cli/src/services/agent_trace_sync/test_http_server.rs deleted file mode 100644 index a5d92729e..000000000 --- a/cli/src/services/agent_trace_sync/test_http_server.rs +++ /dev/null @@ -1,462 +0,0 @@ -//! A bespoke, in-repo, test-only HTTP/1.1 server used to exercise -//! `AuthenticatedControlPlaneClient` (and later the sync engine and CLI -//! wiring) against canned responses and captured requests, without adding a -//! `wiremock`/`httptest`-style dev-dependency to a crate that currently -//! declares none. - -use std::collections::{HashMap, VecDeque}; -use std::io::{BufRead, BufReader, Read, Write}; -use std::net::TcpListener; -use std::sync::{Arc, Mutex}; -use std::thread; -use std::time::Duration; - -#[derive(Clone, Debug)] -pub struct CapturedRequest { - pub method: String, - pub path: String, - pub headers: HashMap, - pub body: String, -} - -#[derive(Clone, Debug)] -pub struct CannedResponse { - pub status: u16, - pub body: String, -} - -impl CannedResponse { - pub fn json(status: u16, body: &serde_json::Value) -> Self { - Self { - status, - body: body.to_string(), - } - } - - pub fn text(status: u16, body: impl Into) -> Self { - Self { - status, - body: body.into(), - } - } -} - -/// A single-threaded, sequential test HTTP server: each accepted connection -/// is answered with the next queued [`CannedResponse`] (or a `500` marker -/// response when the queue is empty), and every parsed request is recorded -/// for assertions. -pub struct TestHttpServer { - pub base_url: String, - requests: Arc>>, - responses: Arc>>, -} - -impl TestHttpServer { - pub fn start() -> Self { - let listener = TcpListener::bind("127.0.0.1:0").expect("bind test http server"); - let addr = listener.local_addr().expect("read test http server addr"); - - let requests: Arc>> = Arc::new(Mutex::new(Vec::new())); - let responses: Arc>> = Arc::new(Mutex::new(VecDeque::new())); - - let thread_requests = Arc::clone(&requests); - let thread_responses = Arc::clone(&responses); - thread::spawn(move || { - for stream in listener.incoming() { - let Ok(stream) = stream else { - continue; - }; - handle_connection(stream, &thread_requests, &thread_responses); - } - }); - - Self { - base_url: format!("http://{addr}"), - requests, - responses, - } - } - - pub fn queue_response(&self, response: CannedResponse) { - self.responses - .lock() - .expect("test http server responses lock") - .push_back(response); - } - - pub fn captured_requests(&self) -> Vec { - self.requests - .lock() - .expect("test http server requests lock") - .clone() - } - - pub fn call_count(&self) -> usize { - self.requests - .lock() - .expect("test http server requests lock") - .len() - } -} - -/// A concurrent test server for sync regressions. State responses are queued -/// independently from dynamically selected batch responses, allowing each -/// stream's expected cursor to determine the response it receives. -pub struct ConcurrentBatchTestServer { - pub base_url: String, - requests: Arc>>, - state_response: Arc>>, - batch_responses: Arc>>, - metrics: Arc>, -} - -#[derive(Clone, Debug, Default)] -struct ConcurrentBatchMetrics { - state_requests: usize, - in_flight: usize, - max_in_flight: usize, - in_flight_by_stream: HashMap, - max_in_flight_by_stream: HashMap, - expected_cursors_by_stream: HashMap>, -} - -impl ConcurrentBatchTestServer { - pub fn start(batch_delay: Duration) -> Self { - let listener = TcpListener::bind("127.0.0.1:0").expect("bind concurrent test http server"); - let addr = listener - .local_addr() - .expect("read concurrent test http server addr"); - - let requests: Arc>> = Arc::new(Mutex::new(Vec::new())); - let state_response = Arc::new(Mutex::new(None)); - let batch_responses = Arc::new(Mutex::new(HashMap::new())); - let metrics = Arc::new(Mutex::new(ConcurrentBatchMetrics::default())); - - let thread_requests = Arc::clone(&requests); - let thread_state_response = Arc::clone(&state_response); - let thread_batch_responses = Arc::clone(&batch_responses); - let thread_metrics = Arc::clone(&metrics); - thread::spawn(move || { - for stream in listener.incoming() { - let Ok(stream) = stream else { - continue; - }; - let requests = Arc::clone(&thread_requests); - let state_response = Arc::clone(&thread_state_response); - let batch_responses = Arc::clone(&thread_batch_responses); - let metrics = Arc::clone(&thread_metrics); - thread::spawn(move || { - handle_concurrent_connection( - stream, - &requests, - &state_response, - &batch_responses, - &metrics, - batch_delay, - ); - }); - } - }); - - Self { - base_url: format!("http://{addr}"), - requests, - state_response, - batch_responses, - metrics, - } - } - - pub fn queue_state_response(&self, response: CannedResponse) { - *self - .state_response - .lock() - .expect("concurrent test http server state response lock") = Some(response); - } - - pub fn queue_batch_response( - &self, - stream: impl Into, - expected_cursor: i64, - response: CannedResponse, - ) { - self.batch_responses - .lock() - .expect("concurrent test http server batch responses lock") - .insert((stream.into(), expected_cursor), response); - } - - pub fn captured_requests(&self) -> Vec { - self.requests - .lock() - .expect("concurrent test http server requests lock") - .clone() - } - - pub fn state_request_count(&self) -> usize { - self.metrics - .lock() - .expect("concurrent test http server metrics lock") - .state_requests - } - - pub fn max_in_flight(&self) -> usize { - self.metrics - .lock() - .expect("concurrent test http server metrics lock") - .max_in_flight - } - - pub fn max_in_flight_for(&self, stream: &str) -> usize { - self.metrics - .lock() - .expect("concurrent test http server metrics lock") - .max_in_flight_by_stream - .get(stream) - .copied() - .unwrap_or(0) - } - - pub fn expected_cursors_for(&self, stream: &str) -> Vec { - self.metrics - .lock() - .expect("concurrent test http server metrics lock") - .expected_cursors_by_stream - .get(stream) - .cloned() - .unwrap_or_default() - } -} - -fn handle_connection( - stream: std::net::TcpStream, - requests: &Arc>>, - responses: &Arc>>, -) { - let mut reader = BufReader::new(stream.try_clone().expect("clone test http stream")); - - let mut request_line = String::new(); - if reader.read_line(&mut request_line).unwrap_or(0) == 0 { - return; - } - let mut parts = request_line.split_whitespace(); - let method = parts.next().unwrap_or_default().to_string(); - let path = parts.next().unwrap_or_default().to_string(); - - let mut headers = HashMap::new(); - let mut content_length = 0usize; - loop { - let mut line = String::new(); - if reader.read_line(&mut line).unwrap_or(0) == 0 { - break; - } - let trimmed = line.trim_end_matches(['\r', '\n']); - if trimmed.is_empty() { - break; - } - if let Some((name, value)) = trimmed.split_once(':') { - let name = name.trim().to_ascii_lowercase(); - let value = value.trim().to_string(); - if name == "content-length" { - content_length = value.parse().unwrap_or(0); - } - headers.insert(name, value); - } - } - - let mut body_bytes = vec![0u8; content_length]; - if content_length > 0 && reader.read_exact(&mut body_bytes).is_err() { - return; - } - let body = String::from_utf8_lossy(&body_bytes).to_string(); - - requests - .lock() - .expect("test http server requests lock") - .push(CapturedRequest { - method, - path, - headers, - body, - }); - - let canned = responses - .lock() - .expect("test http server responses lock") - .pop_front() - .unwrap_or_else(|| CannedResponse { - status: 500, - body: r#"{"error":"no canned response queued"}"#.to_string(), - }); - - let mut stream = stream; - let response = format!( - "HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", - canned.status, - status_reason(canned.status), - canned.body.len(), - canned.body, - ); - let _ = stream.write_all(response.as_bytes()); - let _ = stream.flush(); -} - -fn handle_concurrent_connection( - stream: std::net::TcpStream, - requests: &Arc>>, - state_response: &Arc>>, - batch_responses: &Arc>>, - metrics: &Arc>, - batch_delay: Duration, -) { - let Some(request) = read_request(&stream) else { - return; - }; - requests - .lock() - .expect("concurrent test http server requests lock") - .push(request.clone()); - - let response = match request.path.as_str() { - "/agent-trace/ingestion/state" => { - metrics - .lock() - .expect("concurrent test http server metrics lock") - .state_requests += 1; - state_response - .lock() - .expect("concurrent test http server state response lock") - .clone() - .unwrap_or_else(|| CannedResponse::text(500, "no state response queued")) - } - "/agent-trace/ingestion/batch" => { - let body: serde_json::Value = - serde_json::from_str(&request.body).unwrap_or_else(|_| serde_json::json!({})); - let stream = body - .get("stream") - .and_then(serde_json::Value::as_str) - .unwrap_or("unknown") - .to_string(); - let expected_cursor = body - .get("expectedCursor") - .and_then(serde_json::Value::as_i64) - .unwrap_or(i64::MIN); - - { - let mut metrics = metrics - .lock() - .expect("concurrent test http server metrics lock"); - metrics.in_flight += 1; - metrics.max_in_flight = metrics.max_in_flight.max(metrics.in_flight); - let stream_in_flight = { - let count = metrics - .in_flight_by_stream - .entry(stream.clone()) - .or_default(); - *count += 1; - *count - }; - let stream_max_in_flight = metrics - .max_in_flight_by_stream - .entry(stream.clone()) - .or_default(); - *stream_max_in_flight = (*stream_max_in_flight).max(stream_in_flight); - metrics - .expected_cursors_by_stream - .entry(stream.clone()) - .or_default() - .push(expected_cursor); - } - - thread::sleep(batch_delay); - let response = batch_responses - .lock() - .expect("concurrent test http server batch responses lock") - .get(&(stream.clone(), expected_cursor)) - .cloned() - .unwrap_or_else(|| CannedResponse::text(500, "no batch response queued")); - let mut metrics = metrics - .lock() - .expect("concurrent test http server metrics lock"); - metrics.in_flight -= 1; - if let Some(stream_in_flight) = metrics.in_flight_by_stream.get_mut(&stream) { - *stream_in_flight -= 1; - } - drop(metrics); - response - } - _ => CannedResponse::text(404, "unknown test route"), - }; - - write_response(stream, &response); -} - -fn read_request(stream: &std::net::TcpStream) -> Option { - let mut reader = BufReader::new(stream.try_clone().ok()?); - - let mut request_line = String::new(); - if reader.read_line(&mut request_line).unwrap_or(0) == 0 { - return None; - } - let mut parts = request_line.split_whitespace(); - let method = parts.next().unwrap_or_default().to_string(); - let path = parts.next().unwrap_or_default().to_string(); - - let mut headers = HashMap::new(); - let mut content_length = 0usize; - loop { - let mut line = String::new(); - if reader.read_line(&mut line).unwrap_or(0) == 0 { - break; - } - let trimmed = line.trim_end_matches(['\r', '\n']); - if trimmed.is_empty() { - break; - } - if let Some((name, value)) = trimmed.split_once(':') { - let name = name.trim().to_ascii_lowercase(); - let value = value.trim().to_string(); - if name == "content-length" { - content_length = value.parse().unwrap_or(0); - } - headers.insert(name, value); - } - } - - let mut body_bytes = vec![0u8; content_length]; - if content_length > 0 && reader.read_exact(&mut body_bytes).is_err() { - return None; - } - - Some(CapturedRequest { - method, - path, - headers, - body: String::from_utf8_lossy(&body_bytes).to_string(), - }) -} - -fn write_response(mut stream: std::net::TcpStream, canned: &CannedResponse) { - let response = format!( - "HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", - canned.status, - status_reason(canned.status), - canned.body.len(), - canned.body, - ); - let _ = stream.write_all(response.as_bytes()); - let _ = stream.flush(); -} - -fn status_reason(status: u16) -> &'static str { - match status { - 200 => "OK", - 400 => "Bad Request", - 401 => "Unauthorized", - 403 => "Forbidden", - 409 => "Conflict", - 500 => "Internal Server Error", - 503 => "Service Unavailable", - _ => "Unknown", - } -} diff --git a/cli/src/services/sync/sync.rs b/cli/src/services/sync/sync.rs index ef03b382e..dfc1a4d6d 100644 --- a/cli/src/services/sync/sync.rs +++ b/cli/src/services/sync/sync.rs @@ -537,19 +537,24 @@ fn cursor_for_stream(cursors: &AgentTraceCursors, stream: IngestionStream) -> i6 #[cfg(test)] mod tests { use std::cell::{Cell, RefCell}; + use std::collections::{HashMap, VecDeque}; + use std::io::{BufRead, BufReader, Read, Write}; + use std::net::{TcpListener, TcpStream}; use std::rc::Rc; use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - use std::time::Duration; + use std::sync::{Arc, Mutex}; + use std::thread; + use std::time::{Duration, SystemTime, UNIX_EPOCH}; - use tokio::sync::oneshot; + use serde_json::{json, Value}; + use tokio::sync::{oneshot, Notify}; - use super::{join_three_to_completion, sync_one_stream, TraceSyncError}; + use super::{join_three_to_completion, sync_one_stream, StreamSyncReport, TraceSyncError}; use crate::services::agent_trace_db::MessageRole; use crate::services::agent_trace_export::AgentTraceMessageExportRow; use crate::services::agent_trace_sync::control_plane::{ - AgentTraceIngestionBatchResponse, AuthenticatedControlPlaneClient, ControlPlaneError, - CredentialStore, IngestionStream, + AgentTraceIngestionBatchRequest, AgentTraceIngestionBatchResponse, + AuthenticatedControlPlaneClient, ControlPlaneError, CredentialStore, IngestionStream, }; use crate::services::agent_trace_sync::StreamSyncError; use crate::services::auth::TokenResponse; @@ -558,6 +563,9 @@ mod tests { const NOT_COMPLETED_PROBE: Duration = Duration::from_millis(50); const BACKLOG_ROWS: i64 = 5; + const TEST_BATCH_TIMEOUT: Duration = Duration::from_millis(500); + const TEST_DEADLINE: Duration = Duration::from_secs(10); + const STALL_SAFETY_BOUND: Duration = Duration::from_secs(10); #[tokio::test] async fn sibling_failure_does_not_cancel_in_flight_credential_save() { @@ -620,86 +628,6 @@ mod tests { } } - #[tokio::test] - async fn terminal_stream_failure_stops_sibling_after_its_in_flight_batch() { - let client = AuthenticatedControlPlaneClient::new( - reqwest::Client::builder() - .tls_certs_only(Vec::new()) - .build() - .expect("test http client"), - "http://127.0.0.1:1", - "http://127.0.0.1:1", - "client", - ); - let mut reporter = NoopProgressReporter; - let progress = Rc::new(RefCell::new(&mut reporter)); - let terminated = Cell::new(false); - let sibling_ingest_started = Cell::new(false); - let terminal_observed = Cell::new(false); - let sibling_reads = Cell::new(0usize); - let sibling_ingests = Cell::new(0usize); - let sibling_ingest_completed = Cell::new(false); - - let failing = sync_one_stream( - &client, - "repo", - "source", - IngestionStream::Messages, - 0, - "messages", - async |cursor, _limit| Ok(vec![message_row(cursor + 1)]), - async |_request| { - while !sibling_ingest_started.get() { - tokio::task::yield_now().await; - } - terminal_observed.set(true); - Err(ControlPlaneError::Forbidden("denied".to_string())) - }, - Rc::clone(&progress), - &terminated, - ); - let sibling = sync_one_stream( - &client, - "repo", - "source", - IngestionStream::Parts, - 0, - "parts", - async |cursor, _limit| { - sibling_reads.set(sibling_reads.get() + 1); - Ok(vec![message_row(cursor + 1)]) - }, - async |request| { - sibling_ingests.set(sibling_ingests.get() + 1); - sibling_ingest_started.set(true); - while !terminal_observed.get() { - tokio::task::yield_now().await; - } - sibling_ingest_completed.set(true); - Ok(AgentTraceIngestionBatchResponse { - accepted: request.rows.len(), - cursor: request.rows[0].source_row_id, - }) - }, - Rc::clone(&progress), - &terminated, - ); - let idle = async { Ok::<_, TraceSyncError>(()) }; - - let result = join_three_to_completion(failing, sibling, idle).await; - - assert!(matches!( - result, - Err(TraceSyncError::Stream { - stream: "messages", - source: StreamSyncError::Terminal(ControlPlaneError::Forbidden(_)), - }) - )); - assert!(sibling_ingest_completed.get()); - assert_eq!(sibling_ingests.get(), 1); - assert_eq!(sibling_reads.get(), 1); - } - #[tokio::test] async fn terminal_stream_failure_discards_sibling_rows_read_after_termination() { let client = AuthenticatedControlPlaneClient::new( @@ -746,7 +674,7 @@ mod tests { while !terminated.get() { tokio::task::yield_now().await; } - Ok(vec![message_row(cursor + 1)]) + Ok(backlog_after(cursor, BACKLOG_ROWS)) }, async |request| { sibling_ingests.set(sibling_ingests.get() + 1); @@ -760,7 +688,12 @@ mod tests { ); let idle = async { Ok::<_, TraceSyncError>(()) }; - let result = join_three_to_completion(failing, sibling, idle).await; + let result = tokio::time::timeout( + TEST_DEADLINE, + join_three_to_completion(failing, sibling, idle), + ) + .await + .expect("sync finishes within the test deadline"); assert!(matches!( result, @@ -981,6 +914,387 @@ mod tests { assert_eq!(sibling_ingests.get(), 1); assert_eq!(loads.load(Ordering::SeqCst), 1); } + + struct FixedCredentialStore; + + impl CredentialStore for FixedCredentialStore { + async fn load(&self) -> Result, ControlPlaneError> { + Ok(Some(StoredTokens { + access_token: "access".to_string(), + token_type: "Bearer".to_string(), + expires_in: 3600, + refresh_token: "refresh".to_string(), + scope: None, + stored_at_unix_seconds: SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock after epoch") + .as_secs(), + })) + } + + async fn save(&self, _token: &TokenResponse) -> Result { + unreachable!("fixed credentials never refresh") + } + } + + #[derive(Clone, Copy, Debug)] + enum BatchReply { + Accept, + Stall, + CommitThenStall, + HeadersThenStall, + } + + #[derive(Default)] + struct ScriptedShared { + script: Mutex>, + committed: Mutex>, + batch_requests: Mutex>, + state_requests: AtomicUsize, + batch_stalled: Notify, + } + + struct ScriptedControlPlane { + base_url: String, + shared: Arc, + } + + impl ScriptedControlPlane { + fn start(script: &[BatchReply]) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind scripted control plane"); + let base_url = format!("http://{}", listener.local_addr().expect("server addr")); + let shared = Arc::new(ScriptedShared { + script: Mutex::new(script.iter().copied().collect()), + ..ScriptedShared::default() + }); + let thread_shared = Arc::clone(&shared); + thread::spawn(move || { + for stream in listener.incoming() { + let Ok(stream) = stream else { + continue; + }; + let shared = Arc::clone(&thread_shared); + thread::spawn(move || shared.handle(stream)); + } + }); + Self { base_url, shared } + } + + fn client(&self) -> AuthenticatedControlPlaneClient { + AuthenticatedControlPlaneClient::with_credential_store( + reqwest::Client::builder() + .tls_certs_only(Vec::new()) + .build() + .expect("test http client"), + self.base_url.clone(), + "http://127.0.0.1:1", + "client", + FixedCredentialStore, + ) + .with_batch_timeout(TEST_BATCH_TIMEOUT) + } + + fn expected_cursors(&self, stream: &str) -> Vec { + self.shared + .batch_requests + .lock() + .expect("batch requests lock") + .iter() + .filter(|(name, _)| name == stream) + .map(|(_, cursor)| *cursor) + .collect() + } + + fn state_requests(&self) -> usize { + self.shared.state_requests.load(Ordering::SeqCst) + } + } + + impl ScriptedShared { + fn handle(&self, mut stream: TcpStream) { + let Some((path, body)) = read_scripted_request(&stream) else { + return; + }; + + if path.ends_with("/state") { + self.state_requests.fetch_add(1, Ordering::SeqCst); + let payload = { + let committed = self.committed.lock().expect("committed lock"); + let cursor = |name: &str| committed.get(name).copied().unwrap_or(0); + json!({"cursors": { + "messages": cursor("messages"), + "parts": cursor("parts"), + "diffTraces": 0, + "agentTraces": cursor("agent_traces"), + }}) + }; + write_scripted_response(&mut stream, 200, &payload); + return; + } + + let stream_name = body["stream"].as_str().unwrap_or_default().to_string(); + let rows = body["rows"].as_array().cloned().unwrap_or_default(); + let last_row_id = rows + .last() + .and_then(|row| row["sourceRowId"].as_i64()) + .unwrap_or(0); + self.batch_requests + .lock() + .expect("batch requests lock") + .push(( + stream_name.clone(), + body["expectedCursor"].as_i64().unwrap_or(i64::MIN), + )); + let reply = self.script.lock().expect("script lock").pop_front(); + + match reply { + Some(BatchReply::Accept) => { + self.commit(&stream_name, last_row_id); + write_scripted_response( + &mut stream, + 200, + &json!({"accepted": rows.len(), "cursor": last_row_id}), + ); + } + Some(BatchReply::Stall) => self.stall(stream), + Some(BatchReply::CommitThenStall) => { + self.commit(&stream_name, last_row_id); + self.stall(stream); + } + Some(BatchReply::HeadersThenStall) => { + let _ = stream.write_all( + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 64\r\n\r\n{", + ); + let _ = stream.flush(); + self.stall(stream); + } + None => write_scripted_response( + &mut stream, + 500, + &json!({"error": "no scripted batch reply"}), + ), + } + } + + fn commit(&self, stream: &str, cursor: i64) { + self.committed + .lock() + .expect("committed lock") + .insert(stream.to_string(), cursor); + } + + fn stall(&self, mut stream: TcpStream) { + self.batch_stalled.notify_one(); + let _ = stream.set_read_timeout(Some(STALL_SAFETY_BOUND)); + let mut buffer = [0u8; 64]; + while matches!(stream.read(&mut buffer), Ok(read) if read > 0) {} + } + } + + fn read_scripted_request(stream: &TcpStream) -> Option<(String, Value)> { + let mut reader = BufReader::new(stream.try_clone().ok()?); + let mut request_line = String::new(); + reader.read_line(&mut request_line).ok()?; + let path = request_line.split_whitespace().nth(1)?.to_string(); + let mut content_length = 0usize; + loop { + let mut line = String::new(); + if reader.read_line(&mut line).unwrap_or(0) == 0 { + break; + } + let line = line.trim_end_matches(['\r', '\n']); + if line.is_empty() { + break; + } + if let Some((name, value)) = line.split_once(':') { + if name.eq_ignore_ascii_case("content-length") { + content_length = value.trim().parse().unwrap_or(0); + } + } + } + let mut body = vec![0u8; content_length]; + reader.read_exact(&mut body).ok()?; + Some((path, serde_json::from_slice(&body).unwrap_or(Value::Null))) + } + + fn write_scripted_response(stream: &mut TcpStream, status: u16, payload: &Value) { + let payload = payload.to_string(); + let reason = if status == 200 { + "OK" + } else { + "Internal Server Error" + }; + let response = format!( + "HTTP/1.1 {status} {reason}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{payload}", + payload.len() + ); + let _ = stream.write_all(response.as_bytes()); + let _ = stream.flush(); + } + + fn backlog_after(cursor: i64, rows: i64) -> Vec { + (cursor < rows) + .then(|| message_row(cursor + 1)) + .into_iter() + .collect() + } + + async fn ingest_single_message( + client: &AuthenticatedControlPlaneClient, + ) -> Result { + let request = AgentTraceIngestionBatchRequest { + repository_id: "repo".to_string(), + source_instance_id: "source".to_string(), + stream: IngestionStream::Messages, + expected_cursor: 0, + rows: vec![message_row(1)], + }; + tokio::time::timeout(TEST_DEADLINE, client.ingest_messages(&request)) + .await + .expect("batch request finishes within the test deadline") + } + + async fn sync_messages_backlog( + client: &AuthenticatedControlPlaneClient, + rows: i64, + ) -> Result { + let mut reporter = NoopProgressReporter; + let terminated = Cell::new(false); + let sync = sync_one_stream( + client, + "repo", + "source", + IngestionStream::Messages, + 0, + "messages", + async |cursor, _limit| Ok(backlog_after(cursor, rows)), + async |request| client.ingest_messages(&request).await, + Rc::new(RefCell::new(&mut reporter)), + &terminated, + ); + tokio::time::timeout(TEST_DEADLINE, sync) + .await + .expect("sync finishes within the test deadline") + } + + #[tokio::test] + async fn stalled_batch_request_times_out_as_transport_error() { + let server = ScriptedControlPlane::start(&[BatchReply::Stall]); + + let error = ingest_single_message(&server.client()) + .await + .expect_err("stalled batch times out"); + + assert!( + matches!(error, ControlPlaneError::Transport(_)), + "{error:?}" + ); + assert_eq!(server.expected_cursors("messages"), [0]); + } + + #[tokio::test] + async fn stalled_batch_response_body_times_out_as_transport_error() { + let server = ScriptedControlPlane::start(&[BatchReply::HeadersThenStall]); + + let error = ingest_single_message(&server.client()) + .await + .expect_err("stalled batch body times out"); + + assert!( + matches!(error, ControlPlaneError::Transport(_)), + "{error:?}" + ); + } + + #[tokio::test] + async fn timed_out_batch_committed_by_server_reconciles_without_resending() { + let server = ScriptedControlPlane::start(&[BatchReply::CommitThenStall]); + + let report = sync_messages_backlog(&server.client(), 1) + .await + .expect("sync converges from /state"); + + assert_eq!(report.final_cursor, 1); + assert_eq!(report.uploaded, 0); + assert_eq!(report.batches, 0); + assert_eq!(server.expected_cursors("messages"), [0]); + assert_eq!(server.state_requests(), 1); + } + + #[tokio::test] + async fn timed_out_batch_not_committed_by_server_is_retried() { + let server = ScriptedControlPlane::start(&[BatchReply::Stall, BatchReply::Accept]); + + let report = sync_messages_backlog(&server.client(), 1) + .await + .expect("sync converges after retry"); + + assert_eq!(report.final_cursor, 1); + assert_eq!(report.uploaded, 1); + assert_eq!(report.batches, 1); + assert_eq!(server.expected_cursors("messages"), [0, 0]); + assert_eq!(server.state_requests(), 1); + } + + #[tokio::test] + async fn terminal_failure_is_reported_after_sibling_batch_times_out() { + let server = ScriptedControlPlane::start(&[BatchReply::Stall]); + let client = server.client(); + let mut reporter = NoopProgressReporter; + let progress = Rc::new(RefCell::new(&mut reporter)); + let terminated = Cell::new(false); + let sibling_reads = Cell::new(0usize); + + let failing = sync_one_stream( + &client, + "repo", + "source", + IngestionStream::Messages, + 0, + "messages", + async |cursor, _limit| Ok(vec![message_row(cursor + 1)]), + async |_request| { + server.shared.batch_stalled.notified().await; + Err(ControlPlaneError::Forbidden("denied".to_string())) + }, + Rc::clone(&progress), + &terminated, + ); + let stalled = sync_one_stream( + &client, + "repo", + "source", + IngestionStream::Parts, + 0, + "parts", + async |cursor, _limit| { + sibling_reads.set(sibling_reads.get() + 1); + Ok(backlog_after(cursor, BACKLOG_ROWS)) + }, + async |request| client.ingest_messages(&request).await, + Rc::clone(&progress), + &terminated, + ); + let idle = async { Ok::<_, TraceSyncError>(()) }; + + let result = tokio::time::timeout( + TEST_DEADLINE, + join_three_to_completion(failing, stalled, idle), + ) + .await + .expect("sync finishes within the test deadline"); + + assert!(matches!( + result, + Err(TraceSyncError::Stream { + stream: "messages", + source: StreamSyncError::Terminal(ControlPlaneError::Forbidden(_)), + }) + )); + assert_eq!(sibling_reads.get(), 1); + assert_eq!(server.expected_cursors("parts"), [0]); + assert_eq!(server.state_requests(), 1); + } } #[cfg(test)]