fix: validate live fixture snapshots and shutdown boundaries

crates/worker-runtime/src/lib.rs

@@ -32,10 +32,10 @@ pub mod workspace_issuer;
 #[cfg(feature = "http-server")]
 pub mod workspace_request;
 
-#[cfg(feature = "fs-store")]
-pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions};
 #[cfg(feature = "test-support")]
 pub use fs_store::test_support;
+#[cfg(feature = "fs-store")]
+pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions};
 pub use management::RuntimeOptions;
 pub use runtime::{
     Runtime, RuntimeEventSelectorSubscription, RuntimeSubscriptionRecvError,

tests/e2e/src/worker_environment.rs

@@ -1247,7 +1247,7 @@ impl ManagedProcess {
     fn shutdown(&mut self, timeout: Duration) -> Result<ExitStatus> {
         if let Some(status) = self.child.try_wait()? {
             self.join_readers();
-            return Ok(status);
+            return self.require_successful_exit(status);
         }
         let signal_result = unsafe { libc::kill(self.child.id() as i32, libc::SIGTERM) };
         if signal_result != 0 {
@@ -1257,13 +1257,7 @@ impl ManagedProcess {
         loop {
             if let Some(status) = self.child.try_wait()? {
                 self.join_readers();
-                if status.success() {
+                return self.require_successful_exit(status);
-                    return Ok(status);
-                }
-                return Err(HarnessError::Protocol(format!(
-                    "{} did not complete graceful shutdown: {status}",
-                    self.name
-                )));
             }
             if start.elapsed() >= timeout {
                 let _ = self.child.kill();
@@ -1278,6 +1272,17 @@ impl ManagedProcess {
         }
     }
 
+    fn require_successful_exit(&self, status: ExitStatus) -> Result<ExitStatus> {
+        if status.success() {
+            Ok(status)
+        } else {
+            Err(HarnessError::Protocol(format!(
+                "{} did not complete graceful shutdown: {status}",
+                self.name
+            )))
+        }
+    }
+
     fn join_readers(&mut self) {
         for reader in self.readers.drain(..) {
             let _ = reader.join();
@@ -1398,6 +1403,34 @@ fn io_error(operation: &str, error: std::io::Error) -> HarnessError {
 mod startup_address_tests {
     use super::parse_address_after;
 
+    #[test]
+    fn shutdown_rejects_nonzero_exit_even_when_process_already_exited() -> super::Result<()> {
+        let root = tempfile::tempdir()?;
+        for code in [0, 7] {
+            let mut command = std::process::Command::new("/bin/sh");
+            command
+                .env_clear()
+                .current_dir(root.path())
+                .args(["-c", &format!("exit {code}")]);
+            let mut process = super::ManagedProcess::spawn(
+                "already-exited",
+                command,
+                &root.path().join(format!("{code}.stdout")),
+                &root.path().join(format!("{code}.stderr")),
+            )?;
+            // Reap first to deterministically exercise the initial try_wait branch.
+            assert_eq!(process.child.wait()?.code(), Some(code));
+            let result = process.shutdown(std::time::Duration::from_secs(1));
+            if code == 0 {
+                assert!(result?.success());
+            } else {
+                assert!(matches!(result, Err(super::HarnessError::Protocol(message))
+                    if message.contains("did not complete graceful shutdown")));
+            }
+        }
+        Ok(())
+    }
+
     #[test]
     fn startup_address_waits_for_complete_valid_socket_record() {
         let record = "yoi-server listening on http://127.0.0.1:49123\n";

tests/e2e/src/worker_session.rs

@@ -129,6 +129,14 @@ fn error(message: &str) -> HarnessError {
 
 pub fn validate_full_snapshot(raw: &serde_json::Value, expected: &SessionSnapshot) -> Result<()> {
     let event: Event = serde_json::from_value(raw.clone())?;
+    // Product DTOs intentionally tolerate unknown fields for compatibility. This
+    // fixture boundary does not: anything discarded by typed decoding must never
+    // be persisted to an artifact, including fields outside the Session payload.
+    if serde_json::to_value(&event)? != *raw {
+        return Err(error(
+            "noncanonical or unknown fields in full protocol Snapshot envelope",
+        ));
+    }
     let Event::Snapshot {
         session,
         state,
@@ -235,6 +243,36 @@ mod tests {
         assert!(validate_full_snapshot(&extra, &expected).is_err());
     }
 
+    #[test]
+    fn full_snapshot_rejects_unknown_fields_outside_session_before_artifact_write() {
+        let (raw, expected) = sample();
+        for pointer in [
+            "",
+            "/data",
+            "/data/greeting",
+            "/data/state",
+            "/data/state/state",
+        ] {
+            let mut changed = raw.clone();
+            changed
+                .pointer_mut(pointer)
+                .unwrap()
+                .as_object_mut()
+                .unwrap()
+                .insert(
+                    "hidden_reasoning".into(),
+                    serde_json::json!("must not reach artifact"),
+                );
+            assert!(
+                validate_full_snapshot(&changed, &expected).is_err(),
+                "{pointer}"
+            );
+        }
+        let mut changed = raw;
+        changed["data"]["greeting"]["private_key"] = serde_json::json!("must not reach artifact");
+        assert!(validate_full_snapshot(&changed, &expected).is_err());
+    }
+
     #[test]
     fn non_snapshot_event_cannot_stand_in_for_live_snapshot() {
         let (_, expected) = sample();

tests/e2e/tests/worker_restore.rs

@@ -21,7 +21,7 @@ fn supplied_stopped_worker_restores_exact_live_session_and_cleans_up() -> Result
     let result = (|| {
         let installed = seed::start(&fixture)?;
         let subscription = restore(&installed, &fixture, "fixed-fixture-restore-single")?;
-        stop_and_cleanup(installed, subscription, &fixture)
+        stop_and_cleanup(installed, vec![subscription], &fixture)
     })();
     fixture.source.assert_unchanged()?;
     result
@@ -112,7 +112,7 @@ fn concurrent_fixed_id_restores_keep_sessions_storage_and_cleanup_isolated() ->
         let second_session = FixtureSource::capture(&second_session_root)?;
         let first_root = first.environment.root.clone();
         let first_artifacts = first.environment.artifacts_dir.clone();
-        stop_and_cleanup(first, first_subscription, &fixture)?;
+        stop_and_cleanup(first, vec![first_subscription], &fixture)?;
         assert!(!first_root.exists());
         assert!(first_artifacts.join("cleanup.json").is_file());
         assert!(second.environment.root.is_dir());
@@ -136,8 +136,7 @@ fn concurrent_fixed_id_restores_keep_sessions_storage_and_cleanup_isolated() ->
         )?;
         reconnected.assert_alive()?;
         second_session.assert_unchanged()?;
-        stop_and_cleanup(second, second_subscription, &fixture)?;
+        stop_and_cleanup(second, vec![second_subscription, reconnected], &fixture)?;
-        reconnected.expect_stopped()?;
         Ok(())
     })();
     fixture.source.assert_unchanged()?;
@@ -236,7 +235,7 @@ fn assert_single_state(
 
 fn stop_and_cleanup(
     installed: Installed,
-    subscription: WorkerSessionSubscription,
+    subscriptions: Vec<WorkerSessionSubscription>,
     fixture: &Fixture,
 ) -> Result<()> {
     phase(&installed.environment, "stop")?;
@@ -248,7 +247,12 @@ fn stop_and_cleanup(
     if result.state != WorkerOperationState::Accepted {
         return Err(failure("Worker stop was not accepted"));
     }
-    subscription.expect_stopped()?;
+    installed.environment.assert_connected()?;
+    for subscription in subscriptions {
+        subscription.expect_stopped()?;
+        // A transport close caused by Server/Runtime shutdown is not stop evidence.
+        installed.environment.assert_connected()?;
+    }
     assert_single_state(&installed.environment, fixture, "stopped")?;
     phase(&installed.environment, "shutdown-cleanup")?;
     let report = installed.environment.cleanup()?;

tests/e2e/tests/worker_restore/fixture.rs

@@ -73,6 +73,18 @@ impl Fixture {
         // never skips, and no test generates or repairs a fixture.
         let root = std::env::var_os("YOI_E2E_RESTORE_FIXTURE")
             .map(PathBuf::from)
+            .map(|path| {
+                if path.is_absolute() {
+                    path
+                } else {
+                    Path::new(env!("CARGO_MANIFEST_DIR"))
+                        .parent()
+                        .unwrap()
+                        .parent()
+                        .unwrap()
+                        .join(path)
+                }
+            })
             .unwrap_or_else(|| {
                 PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("fixtures/worker-codex-v1")
             });

tests/e2e/tests/worker_restore/seed.rs

@@ -73,6 +73,9 @@ pub fn start(fixture: &Fixture) -> Result<Installed> {
         &format!("http://{}", reserved_endpoint.local_addr()?),
         &template.worker_name(),
     )?;
+    // Close the placeholder before Backend probes it: an accepting TCP listener
+    // with no HTTP service would stall discovery until its transport timeout.
+    drop(reserved_endpoint);
     super::phase(&environment, "start-seeded-backend")?;
     environment.start_seeded_backend(
         fixture.manifest.workspace_id.clone(),
@@ -122,10 +125,9 @@ pub fn start(fixture: &Fixture) -> Result<Installed> {
         ))?;
         let projected = session_store::project_session_snapshot(template.session_id(), &logs);
         fixture.verify_conversation(&projected)?;
-        snapshot = Some(projected);
+        snapshot = Some(full_restored_snapshot(&sessions, &logs));
         Ok(())
     })?;
-    drop(reserved_endpoint);
     fixture.source.assert_unchanged()?;
     Ok(Installed {
         environment,
@@ -135,6 +137,120 @@ pub fn start(fixture: &Fixture) -> Result<Installed> {
     })
 }
 
+fn full_restored_snapshot(
+    store: &dyn session_store::Store,
+    logs: &[session_store::LogEntry],
+) -> SessionSnapshot {
+    session_store::public_snapshot::project_current_session_snapshot_with_accounting(store, logs)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use protocol::{
+        Event, Greeting, InFlightSnapshot, InvokeKind, SessionSnapshotEntryData,
+        WorkerStateSnapshot, WorkerStatus,
+    };
+    use session_store::{
+        LogEntry, LoggedContentPart, LoggedHistoryEntry, LoggedItem, LoggedRole,
+        LoggedSessionHistoryEntryId, LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin,
+    };
+
+    #[test]
+    fn completed_run_full_snapshot_retains_invoke_and_usage_accounting() -> Result<()> {
+        // Projection regression only: no saved Worker or substitute E2E fixture.
+        let temp = tempfile::tempdir()?;
+        let store = support(session_store::WorkerSessionStore::new(temp.path()))?;
+        let session_id = session_store::new_session_id();
+        let logs = vec![
+            LogEntry::AnnotatedSegmentStart {
+                ts: 1,
+                session_id,
+                system_prompt: None,
+                config: Default::default(),
+                history: vec![],
+                forked_from: None,
+                compacted_from: None,
+            },
+            LogEntry::Invoke {
+                ts: 10,
+                trigger: InvokeKind::UserSend,
+            },
+            LogEntry::AnnotatedAssistantItem {
+                ts: 20,
+                entry: LoggedHistoryEntry {
+                    item: LoggedItem::Message {
+                        role: LoggedRole::Assistant,
+                        content: vec![LoggedContentPart::Text {
+                            text: "unit response".into(),
+                        }],
+                    },
+                    metadata: LoggedSessionHistoryMetadata {
+                        entry_id: LoggedSessionHistoryEntryId("unit-response".into()),
+                        origin: LoggedSessionHistoryOrigin::LegacyUnknown,
+                        derivation: None,
+                        input_metadata: None,
+                    },
+                },
+            },
+            LogEntry::LlmUsage {
+                ts: 30,
+                history_len: 1,
+                input_total_tokens: 100,
+                cache_read_tokens: 20,
+                cache_write_tokens: 0,
+                output_tokens: 5,
+            },
+            serde_json::from_value(serde_json::json!({"kind":"run_completed", "ts":40,
+                "interrupted":false,"result":"finished"}))?,
+        ];
+        let conversation = session_store::project_session_snapshot(session_id, &logs);
+        let expected = full_restored_snapshot(&store, &logs);
+        assert_eq!(expected.entries.len(), conversation.entries.len() + 2);
+        assert!(matches!(
+            expected.entries[expected.entries.len() - 2].data,
+            SessionSnapshotEntryData::Invoke {
+                trigger: InvokeKind::UserSend
+            }
+        ));
+        assert_eq!(expected.entries[expected.entries.len() - 2].timestamp, 10);
+        assert!(matches!(
+            expected.entries.last().unwrap().data,
+            SessionSnapshotEntryData::Usage {
+                input_tokens: 100,
+                cache_read_input_tokens: 20,
+                output_tokens: 5
+            }
+        ));
+        assert_eq!(expected.entries.last().unwrap().timestamp, 30);
+        let live = session_store::public_snapshot::project_current_session_snapshot_with_accounting(
+            &store, &logs,
+        );
+        let event = Event::Snapshot {
+            session: live,
+            state: WorkerStateSnapshot::from(WorkerStatus::Idle),
+            in_flight: InFlightSnapshot::default(),
+            internal_workers: vec![],
+            greeting: Greeting {
+                worker_name: "unit".into(),
+                cwd: ".".into(),
+                provider: "unit".into(),
+                model: "unit".into(),
+                reasoning: None,
+                scope_summary: "none".into(),
+                tools: vec![],
+                context_window: 0,
+                context_tokens: 0,
+                context_usage: None,
+            },
+        };
+        let raw = serde_json::to_value(event)?;
+        yoi_e2e::validate_full_snapshot(&raw, &expected)?;
+        assert!(yoi_e2e::validate_full_snapshot(&raw, &conversation).is_err());
+        Ok(())
+    }
+}
+
 struct Bootstrap {
     token: String,
     workspace_public_key: String,