From a1748b552c52fd5f68de6ba852421fd778aec010 Mon Sep 17 00:00:00 2001
From: Hare <kei.hiracchi.0928@gmail.com>
Date: Wed, 21 Oct 2026 05:52:48 +0900
Subject: [PATCH] fix: validate closed-session diagnostics and fault quotas

---
 crates/tools/src/bash_output.rs               | 37 +++++++++
 .../src/local/command_transport/tests.rs      | 83 +++++++++++++++++++
 crates/worker-runtime/src/http_server.rs      |  6 +-
 3 files changed, 125 insertions(+), 1 deletion(-)

diff --git a/crates/tools/src/bash_output.rs b/crates/tools/src/bash_output.rs
index fad3b02e9..6e9db7783 100644
--- a/crates/tools/src/bash_output.rs
+++ b/crates/tools/src/bash_output.rs
@@ -612,4 +612,41 @@ mod tests {
             Some(workdir::CommandOutputTruncation::ByteBudget)
         );
     }
+
+    #[test]
+    fn incomplete_collection_with_quota_retains_both_causes_without_claiming_complete_output() {
+        let root = BashOutputView::new();
+        let mut writer = root.begin().unwrap();
+        writer.append(CommandStream::Stdout, b"abc");
+        let terminal: workdir::CommandTerminal = serde_json::from_value(serde_json::json!({
+            "status":"completed", "exit_code":0, "process_failure":null,
+            "output_collection":{"state":"incomplete","failure":{"stage":"stdout_drain","reason":"other","os_code":null}},
+            "stdout_bytes":9, "stderr_bytes":0, "stdout_delivered_bytes":3, "stderr_delivered_bytes":0,
+            "truncation_reason":"byte_budget"
+        })).unwrap();
+        let metadata = writer.finish_transport(&terminal).unwrap();
+        assert!(!metadata.terminal.transport_complete);
+        assert_eq!(metadata.terminal.exit_code, Some(0));
+        assert_eq!(
+            metadata.terminal.output_collection,
+            Some(terminal.output_collection)
+        );
+        assert_eq!(
+            metadata.transport_truncation,
+            Some(workdir::CommandOutputTruncation::ByteBudget)
+        );
+        assert_eq!(metadata.stdout_seen, 9);
+        assert_eq!(metadata.stdout_saved, 3);
+        assert!(metadata.storage_truncated);
+        assert!(metadata.retention_limit.is_none());
+        assert_eq!(
+            root.read(&metadata.id, CommandStream::Stdout, 0, 100)
+                .unwrap(),
+            b"abc"
+        );
+        assert_eq!(
+            serde_json::to_value(root.metadata(&metadata.id).unwrap()).unwrap(),
+            serde_json::to_value(metadata).unwrap()
+        );
+    }
 }
diff --git a/crates/workdir/src/local/command_transport/tests.rs b/crates/workdir/src/local/command_transport/tests.rs
index 324ce614e..647e1db16 100644
--- a/crates/workdir/src/local/command_transport/tests.rs
+++ b/crates/workdir/src/local/command_transport/tests.rs
@@ -783,3 +783,86 @@ async fn spawn_fault_retains_safe_reason_without_claiming_output_collection() {
 
 #[cfg(feature = "http-client")]
 mod http;
+
+#[tokio::test]
+async fn drain_failure_and_byte_budget_keep_independent_loss_diagnostics() {
+    for stream in [CommandStream::Stdout, CommandStream::Stderr] {
+        let ready = Gate::default();
+        let exited = Gate::default();
+        let process = FaultProcess {
+            ready: ready.clone(),
+            exited: exited.clone(),
+            wait_error: None,
+        };
+        let reader = FaultReader {
+            prefix: b"partial",
+            prefix_read: ready,
+            fail_ready: exited,
+        };
+        let (sender, mut receiver) = mpsc::channel(QUEUE_CHUNKS);
+        let (_cancel, cancel_rx) = watch::channel(false);
+        let telemetry = CommandTelemetry::new();
+        let terminal = tokio::time::timeout(Duration::from_secs(5), async {
+            match stream {
+                CommandStream::Stdout => {
+                    supervise(
+                        process,
+                        reader,
+                        &b""[..],
+                        10,
+                        3,
+                        "quota-fault",
+                        &telemetry,
+                        cancel_rx,
+                        &sender,
+                    )
+                    .await
+                }
+                CommandStream::Stderr => {
+                    supervise(
+                        process,
+                        &b""[..],
+                        reader,
+                        10,
+                        3,
+                        "quota-fault",
+                        &telemetry,
+                        cancel_rx,
+                        &sender,
+                    )
+                    .await
+                }
+            }
+        })
+        .await
+        .unwrap();
+        let Message::Bytes(chunk) = receiver.recv().await.unwrap();
+        assert_eq!(chunk.stream, stream);
+        assert_eq!(chunk.bytes, b"par");
+        assert!(receiver.try_recv().is_err());
+        assert_eq!(terminal.exit_code, Some(0));
+        assert_eq!(terminal.stdout_bytes + terminal.stderr_bytes, 7);
+        assert_eq!(
+            terminal.stdout_delivered_bytes + terminal.stderr_delivered_bytes,
+            3
+        );
+        assert_eq!(
+            terminal.truncation_reason,
+            Some(crate::CommandOutputTruncation::ByteBudget)
+        );
+        assert_eq!(
+            terminal.output_collection,
+            CommandOutputCollection::Incomplete {
+                failure: CommandFailure {
+                    stage: match stream {
+                        CommandStream::Stdout => CommandFailureStage::StdoutDrain,
+                        CommandStream::Stderr => CommandFailureStage::StderrDrain,
+                    },
+                    reason: CommandIoErrorKind::PermissionDenied,
+                    os_code: None,
+                },
+            }
+        );
+        assert!(terminal.metadata_valid());
+    }
+}
diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs
index a1faefc73..aefc66973 100644
--- a/crates/worker-runtime/src/http_server.rs
+++ b/crates/worker-runtime/src/http_server.rs
@@ -3206,9 +3206,13 @@ mod tests {
             .await
             .expect_err("close must be terminal");
         assert!(matches!(
-            closed_error,
+            closed_error.classification_source(),
             workdir::WorkdirError::Unavailable(_)
         ));
+        assert_eq!(
+            closed_error.failure_reason(),
+            Some(workdir::WorkdirFailureReason::SessionClosed)
+        );
         assert!(
             state
                 .workdir_sessions