diff --git a/crates/openshell-server/src/grpc/policy.rs b/crates/openshell-server/src/grpc/policy.rs index b09a6df79..279329c66 100644 --- a/crates/openshell-server/src/grpc/policy.rs +++ b/crates/openshell-server/src/grpc/policy.rs @@ -5412,6 +5412,8 @@ pub(super) async fn handle_push_sandbox_logs( state: &Arc, request: Request>, ) -> Result, Status> { + const MAX_BATCH_LINES: usize = 100; + let principal = request .extensions() .get::() @@ -5437,9 +5439,20 @@ pub(super) async fn handle_push_sandbox_logs( ) .await?; - let logs: Vec<_> = batch.logs.into_iter().take(100).collect(); + let received_lines = batch.logs.len(); + let logs: Vec<_> = batch.logs.into_iter().take(MAX_BATCH_LINES).collect(); validate_sandbox_log_timestamps(&logs)?; + if received_lines > logs.len() { + warn!( + sandbox_id = %batch.sandbox_id, + received_lines, + accepted_lines = logs.len(), + dropped_lines = received_lines - logs.len(), + "Truncated sandbox log batch" + ); + } + for log in logs { let mut log = log; log.source = "sandbox".to_string(); @@ -8265,6 +8278,148 @@ mod tests { request } + #[tokio::test] + async fn push_sandbox_logs_accepts_batched_replay_and_warns_on_truncation() { + use openshell_core::proto::open_shell_client::OpenShellClient; + use openshell_core::proto::open_shell_server::OpenShellServer; + use tokio_stream::wrappers::TcpListenerStream; + + let state = test_server_state().await; + let sandbox_id = "sb-log-replay"; + state + .store + .put_message(&test_sandbox( + sandbox_id, + "log-replay", + ProtoSandboxPolicy::default(), + vec![], + )) + .await + .unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let principal = crate::auth::principal::test_support::sandbox(sandbox_id); + let service = OpenShellServer::with_interceptor( + crate::grpc::OpenShellService::new(state.clone()), + move |mut request: Request<()>| { + request.extensions_mut().insert(principal.clone()); + Ok(request) + }, + ); + let server = tokio::spawn( + tonic::transport::Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)), + ); + let mut client = OpenShellClient::connect(format!("http://{address}")) + .await + .unwrap(); + + let diagnostics = tempfile::NamedTempFile::new().unwrap(); + let subscriber = tracing_subscriber::fmt() + .json() + .with_env_filter("off,openshell_server::grpc::policy=warn") + .with_writer(std::sync::Mutex::new(diagnostics.reopen().unwrap())) + .finish(); + // This test uses a current-thread runtime, including the gateway tasks. + let _capture = tracing::subscriber::set_default(subscriber); + let mut seen_warnings = 0; + + for (batch_size, count, accepted_count) in [ + (200, 200, 100), + (1, 0, 0), + (1, 1, 1), + (50, 50, 50), + (100, 100, 100), + (101, 101, 100), + (199, 199, 100), + (201, 201, 100), + (50, 200, 200), + (50, 201, 201), + ] { + state.tracing_log_bus.remove(sandbox_id); + let logs: Vec<_> = (0..count) + .map(|index| SandboxLogLine { + sandbox_id: "untrusted-id".to_string(), + event_time: openshell_core::time::timestamp_from_millis(index).ok(), + level: "INFO".to_string(), + target: "reconnect-test".to_string(), + message: format!("line {index}"), + source: "gateway".to_string(), + fields: HashMap::from([("index".to_string(), index.to_string())]), + }) + .collect(); + let expected: Vec<_> = logs + .iter() + .take(accepted_count) + .cloned() + .map(|mut log| { + log.sandbox_id = sandbox_id.to_string(); + log.source = "sandbox".to_string(); + log + }) + .collect(); + let received_lines = logs.len(); + let requests: Vec<_> = logs + .chunks(batch_size) + .map(|chunk| PushSandboxLogsRequest { + sandbox_id: sandbox_id.to_string(), + logs: chunk.to_vec(), + }) + .collect(); + client + .push_sandbox_logs(tokio_stream::iter(requests)) + .await + .unwrap(); + + let received: Vec<_> = state + .tracing_log_bus + .tail(sandbox_id, 201) + .into_iter() + .map(|entry| { + let Some(openshell_core::proto::sandbox_stream_event::Payload::Log(log)) = + entry.event.payload + else { + panic!("expected a log event"); + }; + log + }) + .collect(); + assert_eq!( + received.len(), + expected.len(), + "gateway ingestion of {count} log lines in batches of {batch_size}" + ); + assert_eq!( + received, expected, + "{count} log lines, batches of {batch_size}" + ); + + let warnings: Vec = std::fs::read_to_string(diagnostics.path()) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + let emitted = &warnings[seen_warnings..]; + assert_eq!(emitted.len(), usize::from(received_lines > accepted_count)); + if let Some(warning) = emitted.first() { + assert_eq!(warning["level"], "WARN"); + assert_eq!( + warning["fields"], + serde_json::json!({ + "message": "Truncated sandbox log batch", + "sandbox_id": sandbox_id, + "received_lines": received_lines, + "accepted_lines": accepted_count, + "dropped_lines": received_lines - accepted_count, + }) + ); + } + seen_warnings = warnings.len(); + } + server.abort(); + } + #[test] fn sandbox_log_timestamp_validation_rejects_malformed_batches_atomically() { let logs = vec![ diff --git a/crates/openshell-supervisor-process/src/log_push.rs b/crates/openshell-supervisor-process/src/log_push.rs index c2ee97760..2d1c8dc70 100644 --- a/crates/openshell-supervisor-process/src/log_push.rs +++ b/crates/openshell-supervisor-process/src/log_push.rs @@ -180,13 +180,15 @@ pub fn start_log_push( const MAX_BACKOFF: tokio::time::Duration = tokio::time::Duration::from_secs(30); /// Initial backoff delay after a connection failure. const INITIAL_BACKOFF: tokio::time::Duration = tokio::time::Duration::from_secs(1); +/// Lines per regular or reconnect-replay batch, below the gateway's ingest cap. +const LOG_PUSH_BATCH_LINES: usize = 50; async fn run_push_loop( gateway: GatewayEndpoint, sandbox_id: String, mut rx: mpsc::Receiver, ) { - let mut batch = Vec::with_capacity(50); + let mut batch = Vec::with_capacity(LOG_PUSH_BATCH_LINES); let mut backoff = INITIAL_BACKOFF; let mut attempt: u64 = 0; @@ -236,20 +238,13 @@ async fn run_push_loop( }); // --- Flush any lines buffered during reconnect --- - if !batch.is_empty() { - let lines = std::mem::take(&mut batch); - if push_tx - .send(PushSandboxLogsRequest { - sandbox_id: sandbox_id.clone(), - logs: lines, - }) - .await - .is_err() - { - // RPC died immediately — go back to reconnect. - backoff = INITIAL_BACKOFF; - continue; - } + if flush_reconnect_buffer(&push_tx, &sandbox_id, &mut batch) + .await + .is_err() + { + // RPC died immediately — go back to reconnect. + backoff = INITIAL_BACKOFF; + continue; } // --- Batch and send loop (runs until stream breaks) --- @@ -274,7 +269,7 @@ async fn run_push_loop( return; }; batch.push(line); - if batch.len() >= 50 { + if batch.len() >= LOG_PUSH_BATCH_LINES { let lines = std::mem::take(&mut batch); if push_tx.send(PushSandboxLogsRequest { sandbox_id: sandbox_id.clone(), @@ -316,6 +311,24 @@ async fn run_push_loop( } } +async fn flush_reconnect_buffer( + tx: &mpsc::Sender, + sandbox_id: &str, + batch: &mut Vec, +) -> Result<(), mpsc::error::SendError> { + // The reconnect buffer can hold 200 lines, but the gateway only ingests + // the first 100 of each request. Replay using the regular batch size. + while !batch.is_empty() { + let count = batch.len().min(LOG_PUSH_BATCH_LINES); + tx.send(PushSandboxLogsRequest { + sandbox_id: sandbox_id.to_string(), + logs: batch.drain(..count).collect(), + }) + .await?; + } + Ok(()) +} + /// Drain incoming log lines during a backoff delay so the tracing layer's /// `try_send` doesn't fill up. Lines received during backoff are kept in `batch` /// (up to a limit) so they can be sent after reconnecting. @@ -524,6 +537,87 @@ mod tests { } } + #[tokio::test] + async fn reconnect_replay_retains_all_buffered_lines_at_the_gateway_ingest_cap() { + let expected: Vec<_> = (0..200).map(|i| test_line(&format!("line {i}"))).collect(); + let (tx, mut rx) = mpsc::channel(LOG_PUSH_QUEUE_LINES); + for line in &expected { + tx.try_send(line.clone()).unwrap(); + } + drop(tx); + + let mut batch = Vec::new(); + drain_during_backoff(&mut rx, &mut batch, INITIAL_BACKOFF).await; + assert_eq!(batch, expected); + + let (push_tx, mut push_rx) = mpsc::channel(32); + flush_reconnect_buffer(&push_tx, "sb-test", &mut batch) + .await + .unwrap(); + drop(push_tx); + + let mut received = Vec::new(); + let mut batch_sizes = Vec::new(); + while let Some(request) = push_rx.recv().await { + assert_eq!(request.sandbox_id, "sb-test"); + batch_sizes.push(request.logs.len()); + // Match handle_push_sandbox_logs in openshell-server/src/grpc/policy.rs. + received.extend(request.logs.into_iter().take(100)); + } + eprintln!( + "reconnect replay: batches={batch_sizes:?}, retained={}/{}", + received.len(), + expected.len() + ); + assert_eq!( + received.len(), + expected.len(), + "the gateway ingest cap must not truncate reconnect-buffered logs" + ); + assert_eq!(received, expected, "replay must preserve order and fields"); + assert!(batch.is_empty()); + } + + #[tokio::test] + async fn reconnect_replay_preserves_empty_and_partial_batches() { + for count in [0, 1, 49, 50, 51, 100, 101, 199] { + let expected: Vec<_> = (0..count) + .map(|i| test_line(&format!("line {i}"))) + .collect(); + let mut batch = expected.clone(); + let (tx, mut rx) = mpsc::channel(32); + flush_reconnect_buffer(&tx, "sb-test", &mut batch) + .await + .unwrap(); + drop(tx); + + let mut received = Vec::new(); + while let Some(request) = rx.recv().await { + assert_eq!(request.sandbox_id, "sb-test"); + assert!(!request.logs.is_empty()); + assert!(request.logs.len() <= LOG_PUSH_BATCH_LINES); + received.extend(request.logs); + } + assert_eq!(received, expected, "replay of {count} lines"); + assert!(batch.is_empty()); + } + } + + #[tokio::test] + async fn reconnect_replay_returns_an_error_if_the_stream_closes() { + let mut batch: Vec<_> = (0..200).map(|i| test_line(&format!("line {i}"))).collect(); + let (tx, rx) = mpsc::channel(1); + drop(rx); + + let error = flush_reconnect_buffer(&tx, "sb-test", &mut batch) + .await + .expect_err("a closed stream must return to the reconnect loop"); + + assert_eq!(error.0.logs.len(), LOG_PUSH_BATCH_LINES); + assert_eq!(batch.len(), 200 - LOG_PUSH_BATCH_LINES); + assert_eq!(batch[0].message, "line 50"); + } + #[tokio::test] async fn drain_during_backoff_buffers_up_to_the_cap_and_drops_the_rest() { let (tx, mut rx) = mpsc::channel::(1024);