feat: US-064 - Implement context restore with bridge function replacement

NathanFlurry · claude · NathanFlurry · commit bb50c0f031b8 · 2026-03-19T22:53:58.000-07:00
Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/crates/v8-runtime/src/bridge.rs b/crates/v8-runtime/src/bridge.rs
@@ -410,6 +410,25 @@ fn async_bridge_callback(
     rv.set(promise.into());
 }
 
+/// Replace stub bridge functions on a snapshot-restored context with real
+/// session-local bridge functions. Overwrites the 38 stub globals with
+/// functions backed by session-local BridgeCallContext and SessionBuffers.
+///
+/// Returns (BridgeFnStore, AsyncBridgeFnStore) that must be kept alive
+/// for the lifetime of the V8 context.
+pub fn replace_bridge_fns(
+    scope: &mut v8::HandleScope,
+    ctx: *const BridgeCallContext,
+    pending: *const PendingPromises,
+    buffers: *const RefCell<SessionBuffers>,
+    sync_fns: &[&str],
+    async_fns: &[&str],
+) -> (BridgeFnStore, AsyncBridgeFnStore) {
+    let sync_store = register_sync_bridge_fns(scope, ctx, buffers, sync_fns);
+    let async_store = register_async_bridge_fns(scope, ctx, pending, buffers, async_fns);
+    (sync_store, async_store)
+}
+
 /// Register stub bridge functions on the V8 global for snapshot creation.
 ///
 /// Uses the same sync_bridge_callback / async_bridge_callback as real
diff --git a/crates/v8-runtime/src/session.rs b/crates/v8-runtime/src/session.rs
@@ -262,6 +262,13 @@ fn session_thread(
     #[cfg(not(test))]
     let mut _v8_context: Option<v8::Global<v8::Context>> = None;
 
+    // Whether the isolate was created from a context snapshot.
+    // When true, Execute uses the snapshot's default context (bridge IIFE
+    // already executed) and skips re-running the bridge code. Bridge function
+    // stubs in the snapshot are replaced with real session-local functions.
+    #[cfg(not(test))]
+    let mut from_snapshot = false;
+
     #[cfg(not(test))]
     let pending = bridge::PendingPromises::new();
 
@@ -320,6 +327,7 @@ fn session_thread(
                             let mut iso = if !effective_bridge_code.is_empty() {
                                 match snapshot_cache.get_or_create(&effective_bridge_code) {
                                     Ok(blob) => {
+                                        from_snapshot = true;
                                         snapshot::create_isolate_from_snapshot((*blob).clone(), heap_limit_mb)
                                     }
                                     Err(e) => {
@@ -339,10 +347,13 @@ fn session_thread(
 
                         let iso = v8_isolate.as_mut().unwrap();
 
-                        // Create a fresh V8 context per execution (clean global scope)
+                        // Create execution context: Context::new on a snapshot-restored
+                        // isolate gives a fresh clone of the snapshot's default context
+                        // (bridge IIFE already executed, all infrastructure set up).
+                        // On a non-snapshot isolate, this gives a blank context.
                         let exec_context = isolate::create_context(iso);
 
-                        // Inject globals from last InjectGlobals payload into the fresh context
+                        // Inject globals from last InjectGlobals payload
                         if let Some(ref payload) = last_globals_payload {
                             let scope = &mut v8::HandleScope::new(iso);
                             let ctx = v8::Local::new(scope, &exec_context);
@@ -370,26 +381,22 @@ fn session_thread(
                             Arc::clone(&call_id_router),
                         );
 
-                        // Register sync and async bridge functions
+                        // Replace stub bridge functions with real session-local ones
+                        // (on snapshot context) or register from scratch (on fresh context).
+                        // Both paths use the same function — global.set() works for both.
                         let _sync_store;
                         let _async_store;
                         {
                             let scope = &mut v8::HandleScope::new(iso);
                             let ctx = v8::Local::new(scope, &exec_context);
                             let scope = &mut v8::ContextScope::new(scope, ctx);
 
-                            _sync_store = bridge::register_sync_bridge_fns(
-                                scope,
-                                &bridge_ctx as *const BridgeCallContext,
-                                &session_buffers as *const std::cell::RefCell<bridge::SessionBuffers>,
-                                &SYNC_BRIDGE_FNS,
-                            );
-
-                            _async_store = bridge::register_async_bridge_fns(
+                            (_sync_store, _async_store) = bridge::replace_bridge_fns(
                                 scope,
                                 &bridge_ctx as *const BridgeCallContext,
                                 &pending as *const bridge::PendingPromises,
                                 &session_buffers as *const std::cell::RefCell<bridge::SessionBuffers>,
+                                &SYNC_BRIDGE_FNS,
                                 &ASYNC_BRIDGE_FNS,
                             );
                         }
@@ -403,14 +410,17 @@ fn session_thread(
                             _ => None,
                         };
 
-                        // Execute code (fresh context per execution)
+                        // On snapshot-restored context, skip bridge IIFE (already in
+                        // snapshot) and run user code only. On fresh context, run full
+                        // bridge code + user code as before.
+                        let bridge_code_for_exec = if from_snapshot { "" } else { &effective_bridge_code };
                         let file_path_opt = if file_path.is_empty() { None } else { Some(file_path.as_str()) };
                         let (code, exports, error) = if mode == 0 {
                             let scope = &mut v8::HandleScope::new(iso);
                             let ctx = v8::Local::new(scope, &exec_context);
                             let scope = &mut v8::ContextScope::new(scope, ctx);
                             let (c, e) =
-                                execution::execute_script(scope, &effective_bridge_code, &user_code, &mut bridge_cache);
+                                execution::execute_script(scope, bridge_code_for_exec, &user_code, &mut bridge_cache);
                             (c, None, e)
                         } else {
                             let scope = &mut v8::HandleScope::new(iso);
@@ -419,7 +429,7 @@ fn session_thread(
                             execution::execute_module(
                                 scope,
                                 &bridge_ctx,
-                                &effective_bridge_code,
+                                bridge_code_for_exec,
                                 &user_code,
                                 file_path_opt,
                                 &mut bridge_cache,
diff --git a/crates/v8-runtime/src/snapshot.rs b/crates/v8-runtime/src/snapshot.rs
@@ -66,6 +66,10 @@ pub fn create_snapshot(bridge_code: &str) -> Result<v8::StartupData, String> {
 /// These are placeholder values so bridge code that reads _processConfig or
 /// _osConfig at setup time doesn't fail. They're overwritten per-session
 /// after snapshot restore via inject_globals_from_payload.
+///
+/// Properties are set as READ_ONLY (not DONT_DELETE) so they remain
+/// configurable — inject_globals_from_payload can redefine them with
+/// READ_ONLY | DONT_DELETE after restore.
 fn inject_snapshot_defaults(scope: &mut v8::HandleScope) {
     let context = scope.get_current_context();
     let global = context.global(scope);
@@ -84,7 +88,9 @@ fn inject_snapshot_defaults(scope: &mut v8::HandleScope) {
         pc_obj.set_integrity_level(scope, v8::IntegrityLevel::Frozen);
     }
     let pc_key = v8::String::new(scope, "_processConfig").unwrap();
-    let attr = v8::PropertyAttribute::READ_ONLY | v8::PropertyAttribute::DONT_DELETE;
+    // READ_ONLY only — no DONT_DELETE so the property remains configurable
+    // for override after snapshot restore
+    let attr = v8::PropertyAttribute::READ_ONLY;
     global.define_own_property(scope, pc_key.into(), pc_val, attr);
 
     // _osConfig: default placeholder (overwritten per-session)
@@ -101,7 +107,8 @@ fn inject_snapshot_defaults(scope: &mut v8::HandleScope) {
         oc_obj.set_integrity_level(scope, v8::IntegrityLevel::Frozen);
     }
     let oc_key = v8::String::new(scope, "_osConfig").unwrap();
-    let attr2 = v8::PropertyAttribute::READ_ONLY | v8::PropertyAttribute::DONT_DELETE;
+    // READ_ONLY only — no DONT_DELETE so the property remains configurable
+    let attr2 = v8::PropertyAttribute::READ_ONLY;
     global.define_own_property(scope, oc_key.into(), oc_val, attr2);
 }
 
@@ -822,5 +829,201 @@ mod tests {
             let mut isolate = create_isolate_from_snapshot((*arc1).clone(), None);
             assert_eq!(eval(&mut isolate, "1 + 1"), "2");
         }
+
+        // --- Part 18: Context restore + replace_bridge_fns dispatches correctly ---
+        // Verifies the full context snapshot restore flow: create snapshot with
+        // stubs, restore, replace stubs with real bridge functions, verify the
+        // replaced functions dispatch to the real Rust callbacks.
+        {
+            use std::cell::RefCell;
+            use crate::bridge::{
+                replace_bridge_fns, SessionBuffers, PendingPromises,
+            };
+            use crate::host_call::BridgeCallContext;
+
+            // Create snapshot with stubs + simple bridge IIFE
+            let bridge_code = r#"
+                (function() {
+                    // Getter-based facade referencing globalThis._fsReadFile
+                    var _fs = {};
+                    Object.defineProperties(_fs, {
+                        readFile: { get: function() { return globalThis._fsReadFile; }, enumerable: true },
+                    });
+                    globalThis._fs = _fs;
+                    globalThis.__bridge_ready = true;
+                })();
+            "#;
+            let blob = create_snapshot(bridge_code).expect("snapshot creation");
+            let mut isolate = create_isolate_from_snapshot(blob, None);
+            crate::execution::disable_wasm(&mut isolate);
+
+            // Create BridgeCallContext (sync calls will fail but we verify dispatch)
+            let (ipc_tx, _ipc_rx) = crossbeam_channel::unbounded::<Vec<u8>>();
+            let call_id_router: crate::host_call::CallIdRouter =
+                Arc::new(Mutex::new(std::collections::HashMap::new()));
+            let receiver = crate::host_call::ReaderResponseReceiver::new(
+                Box::new(std::io::Cursor::new(Vec::<u8>::new())),
+            );
+            let sender = crate::host_call::ChannelFrameSender::new(ipc_tx);
+            let bridge_ctx = BridgeCallContext::with_receiver(
+                Box::new(sender),
+                Box::new(receiver),
+                "test-session".to_string(),
+                call_id_router,
+            );
+            let session_buffers = RefCell::new(SessionBuffers::new());
+            let pending = PendingPromises::new();
+
+            // Restore context and replace bridge functions
+            let scope = &mut v8::HandleScope::new(&mut isolate);
+            let context = v8::Context::new(scope, Default::default());
+            let scope = &mut v8::ContextScope::new(scope, context);
+
+            let (_sync_store, _async_store) = replace_bridge_fns(
+                scope,
+                &bridge_ctx as *const BridgeCallContext,
+                &pending as *const PendingPromises,
+                &session_buffers as *const RefCell<SessionBuffers>,
+                &["_log", "_fsReadFile"],
+                &["_scheduleTimer"],
+            );
+
+            // Verify bridge infrastructure from IIFE survives restore
+            let check = v8::String::new(scope, r#"
+                (function() {
+                    var results = [];
+                    results.push('__bridge_ready=' + globalThis.__bridge_ready);
+                    results.push('_fs_exists=' + (typeof _fs === 'object'));
+                    // Getter should resolve to the REPLACED function (not stub)
+                    results.push('_fs.readFile_type=' + typeof _fs.readFile);
+                    // Direct global should also be the replaced function
+                    results.push('_log_type=' + typeof _log);
+                    results.push('_scheduleTimer_type=' + typeof _scheduleTimer);
+                    return results.join(';');
+                })()
+            "#).unwrap();
+            let script = v8::Script::compile(scope, check, None).unwrap();
+            let result = script.run(scope).unwrap();
+            assert_eq!(
+                result.to_rust_string_lossy(scope),
+                "__bridge_ready=true;_fs_exists=true;_fs.readFile_type=function;_log_type=function;_scheduleTimer_type=function",
+                "restored context should have bridge IIFE state + replaced functions"
+            );
+        }
+
+        // --- Part 19: _processConfig is overridable after restore ---
+        // Verifies that inject_snapshot_defaults uses configurable properties
+        // so inject_globals_from_payload can override them per session.
+        {
+            use crate::bridge::serialize_v8_value;
+
+            let bridge_code = r#"
+                (function() {
+                    // Verify default _processConfig from snapshot
+                    globalThis.__snapshotCwd = _processConfig.cwd;
+                })();
+            "#;
+            let blob = create_snapshot(bridge_code).expect("snapshot creation");
+            let mut isolate = create_isolate_from_snapshot(blob, None);
+
+            let scope = &mut v8::HandleScope::new(&mut isolate);
+            let context = v8::Context::new(scope, Default::default());
+            let scope = &mut v8::ContextScope::new(scope, context);
+
+            // Verify snapshot defaults are present
+            let check = v8::String::new(scope, "__snapshotCwd").unwrap();
+            let script = v8::Script::compile(scope, check, None).unwrap();
+            let result = script.run(scope).unwrap();
+            assert_eq!(result.to_rust_string_lossy(scope), "/");
+
+            // Create a V8 payload to override _processConfig
+            let payload_code = r#"({
+                processConfig: { cwd: "/app", env: { FOO: "bar" }, timing_mitigation: "off", frozen_time_ms: null },
+                osConfig: { homedir: "/home/user", tmpdir: "/tmp", platform: "linux", arch: "arm64" }
+            })"#;
+            let payload_source = v8::String::new(scope, payload_code).unwrap();
+            let payload_script = v8::Script::compile(scope, payload_source, None).unwrap();
+            let payload_val = payload_script.run(scope).unwrap();
+            let payload_bytes = serialize_v8_value(scope, payload_val)
+                .expect("serialize payload");
+
+            // Inject per-session globals (overrides snapshot defaults)
+            crate::execution::inject_globals_from_payload(scope, &payload_bytes);
+
+            // Verify _processConfig was overridden
+            let check = v8::String::new(scope, "_processConfig.cwd").unwrap();
+            let script = v8::Script::compile(scope, check, None).unwrap();
+            let result = script.run(scope).unwrap();
+            assert_eq!(
+                result.to_rust_string_lossy(scope),
+                "/app",
+                "_processConfig.cwd should be overridden from '/' to '/app'"
+            );
+
+            // Verify _osConfig was overridden
+            let check = v8::String::new(scope, "_osConfig.arch").unwrap();
+            let script = v8::Script::compile(scope, check, None).unwrap();
+            let result = script.run(scope).unwrap();
+            assert_eq!(
+                result.to_rust_string_lossy(scope),
+                "arm64",
+                "_osConfig.arch should be overridden to 'arm64'"
+            );
+        }
+
+        // --- Part 20: Multiple restores from same snapshot are independent ---
+        // Verifies that user code in one restored context does not leak to another.
+        {
+            let bridge_code = r#"
+                (function() {
+                    globalThis.__bridge_ok = true;
+                })();
+            "#;
+            let blob = create_snapshot(bridge_code).expect("snapshot creation");
+            let blob_bytes: Vec<u8> = blob.to_vec();
+
+            // Restore A: set a session-specific global
+            {
+                let mut isolate = create_isolate_from_snapshot(blob_bytes.clone(), None);
+                let scope = &mut v8::HandleScope::new(&mut isolate);
+                let context = v8::Context::new(scope, Default::default());
+                let scope = &mut v8::ContextScope::new(scope, context);
+
+                // Bridge state from snapshot should be present
+                let check = v8::String::new(scope, "String(__bridge_ok)").unwrap();
+                let script = v8::Script::compile(scope, check, None).unwrap();
+                let result = script.run(scope).unwrap();
+                assert_eq!(result.to_rust_string_lossy(scope), "true");
+
+                // Set session-specific state
+                let code = v8::String::new(scope, "globalThis.__user_data = 'session-a';").unwrap();
+                let script = v8::Script::compile(scope, code, None).unwrap();
+                script.run(scope);
+            }
+
+            // Restore B: session A's state should not be visible
+            {
+                let mut isolate = create_isolate_from_snapshot(blob_bytes.clone(), None);
+                let scope = &mut v8::HandleScope::new(&mut isolate);
+                let context = v8::Context::new(scope, Default::default());
+                let scope = &mut v8::ContextScope::new(scope, context);
+
+                // Bridge state should still be present
+                let check = v8::String::new(scope, "String(__bridge_ok)").unwrap();
+                let script = v8::Script::compile(scope, check, None).unwrap();
+                let result = script.run(scope).unwrap();
+                assert_eq!(result.to_rust_string_lossy(scope), "true");
+
+                // Session A's data should NOT be visible
+                let check = v8::String::new(scope, "typeof __user_data").unwrap();
+                let script = v8::Script::compile(scope, check, None).unwrap();
+                let result = script.run(scope).unwrap();
+                assert_eq!(
+                    result.to_rust_string_lossy(scope),
+                    "undefined",
+                    "session B should not see session A's user data"
+                );
+            }
+        }
     }
 }
