wip: fix whatever is going on with envoys

Author: NathanFlurry

.agent/notes/driver-test-fixes-summary.md

@@ -0,0 +1,56 @@
+# Driver Test Fixes Summary
+
+This was not one bug. It was a stack of separate issues that all showed up in the file-system and engine driver suites.
+
+## Main fixes
+
+1. Raw and Drizzle SQLite were not choosing the same backend path.
+
+Raw `rivetkit/db` could use the native addon while `rivetkit/db/drizzle` always went through the WASM VFS. I unified backend selection behind shared open logic in [open-database.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/db/open-database.ts), with a native `IDatabase` adapter in [native-adapter.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/db/native-adapter.ts). Both providers now reuse that in [mod.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/db/mod.ts) and [mod.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/db/drizzle/mod.ts). This removed a class of raw-path versus drizzle-path inconsistencies.
+
+2. The native SQLite KV channel had actor-resolution bugs.
+
+The KV channel was caching actor resolution too coarsely, which is wrong when one connection multiplexes many actors. That was fixed in [lib.rs](/home/nathan/r2/engine/packages/pegboard-kv-channel/src/lib.rs). I also fixed the open race where native SQLite could attempt `sqlite3_open_v2` before the actor was actually resolvable on the engine side, which was the source of the earlier `SQLITE_CANTOPEN` and internal-error failures. Related engine-side actor lifecycle and lookup changes were needed in [create.rs](/home/nathan/r2/engine/packages/pegboard/src/ops/actor/create.rs), [get_for_key.rs](/home/nathan/r2/engine/packages/pegboard/src/ops/actor/get_for_key.rs), [mod.rs](/home/nathan/r2/engine/packages/pegboard/src/workflows/actor2/mod.rs), and [runtime.rs](/home/nathan/r2/engine/packages/pegboard/src/workflows/actor2/runtime.rs).
+
+3. The native client path was not resilient to channel loss.
+
+When the native KV channel went stale, the old code surfaced connection-closed errors instead of reopening and retrying. I fixed that in [native-sqlite.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/db/native-sqlite.ts) and [native-adapter.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/db/native-adapter.ts). I also fixed crash cleanup ordering so hard-crash tests did not abort DB cleanup too early in [mod.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts).
+
+4. Query-backed actor handles and gateway routing had mismatches.
+
+Some engine failures were because `.get()` / `.getOrCreate()` / `.getForId()` were not consistently preserving the query-backed gateway path the tests expected. That was fixed in [actor-handle.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts), [actor-query.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/client/actor-query.ts), [resolve-gateway-target.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/driver-helpers/resolve-gateway-target.ts), [mod.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/remote-manager-driver/mod.ts), and [mod.rs](/home/nathan/r2/engine/packages/guard/src/routing/pegboard_gateway/mod.rs). That cleared the `.get().connect()` and `.getForId().connect()` engine proxy failures.
+
+5. Raw websocket lifecycle handling had real race bugs.
+
+There were two separate websocket problems:
+
+- `Conn.disconnect()` was re-entrant, so a local close and the resulting close event could double-run cleanup. I made it single-flight in [mod.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts).
+- Actors could auto-sleep before a wake-triggered raw websocket open had actually been delivered into RivetKit. I first fixed the local `prepareConn -> connectConn` gap in [connection-manager.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts), then fixed the earlier engine-only wake-to-open gap by adding a driver-level initial sleep override in [driver.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/actor/driver.ts), consumed in [mod.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts) and implemented in [actor-driver.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/drivers/engine/actor-driver.ts). That was the last real engine flake.
+
+6. AgentOS process callbacks were leaking shutdown-time errors into Vitest.
+
+The file-system suite ended up green on assertions but still failed because background process stdout and stderr callbacks were calling `broadcast()` after the actor had already shut down. I made those shutdown-safe in [process.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/process.ts). That removed the final unhandled rejections from the file-system run.
+
+## Supporting fixes
+
+There were also supporting changes needed to make the suites reliably green:
+
+- Better raw websocket open ordering and route handling in [router-websocket-endpoints.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts) and [actor-driver.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/drivers/engine/actor-driver.ts).
+- File-system driver and global-state cleanup and sleep behavior fixes in [global-state.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts), [actor.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/actor.ts), and [manager.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/manager.ts).
+- Remote websocket client and proxy fixes in [actor-websocket-client.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/remote-manager-driver/actor-websocket-client.ts) and [actor-conn.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts).
+- Test harness changes in [driver-engine.test.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/tests/driver-engine.test.ts), [utils.ts](/home/nathan/r2/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/utils.ts), and several driver-test files so the engine suite stopped fighting its own setup and teardown.
+
+## Net result
+
+The important root causes were:
+
+- backend split between raw and drizzle,
+- native KV and open races,
+- query-backed gateway mismatches,
+- raw-websocket sleep and disconnect races,
+- shutdown-time async callbacks escaping actor lifetime.
+
+Once those were fixed, both suites passed cleanly:
+
+- [driver-file-system-full-after-process-fix.log](/tmp/driver-file-system-full-after-process-fix.log)
+- [driver-engine-full-after-initial-sleep-fix.log](/tmp/driver-engine-full-after-initial-sleep-fix.log)

engine/packages/pegboard-envoy/src/conn.rs

@@ -147,7 +147,7 @@ pub async fn handle_init(
 				// Read existing data
 				let (create_ts_entry, old_last_ping_ts_entry, version_entry) = tokio::try_join!(
 					tx.read_opt(&create_ts_key, Serializable),
-					tx.read_opt(&create_ts_key, Serializable),
+					tx.read_opt(&last_ping_ts_key, Serializable),
 					tx.read_opt(&version_key, Serializable),
 				)?;
 
@@ -228,6 +228,7 @@ pub async fn handle_init(
 					);
 
 					tx.add_conflict_key(&old_lb_key, ConflictRangeType::Read)?;
+					tx.delete(&old_lb_key);
 				}
 
 				// Insert into LB

engine/packages/pegboard-kv-channel/src/lib.rs

@@ -434,9 +434,8 @@ async fn actor_request_task(
 	open_actors: Arc<Mutex<HashSet<String>>>,
 	mut rx: mpsc::Receiver<protocol::ToServerRequest>,
 ) {
-	// Cached actor resolution. Populated on first KV request, reused for all
-	// subsequent requests. Actor name is immutable so this never goes stale.
-	let mut cached_actor: Option<(Id, String)> = None;
+	// Cache keyed by actor id since a single connection multiplexes many actors.
+	let mut cached_actors: HashMap<String, (Id, String)> = HashMap::new();
 
 	while let Some(req) = rx.recv().await {
 		let is_close = matches!(req.data, protocol::RequestData::ActorCloseRequest);
@@ -464,11 +463,10 @@ async fn actor_request_task(
 						)
 					}
 				} else {
-					// Lazy-resolve and cache.
-					if cached_actor.is_none() {
+					if !cached_actors.contains_key(&req.actor_id) {
 						match resolve_actor(&ctx, &req.actor_id, namespace_id).await {
 							Ok(v) => {
-								cached_actor = Some(v);
+								cached_actors.insert(req.actor_id.clone(), v);
 							}
 							Err(resp) => {
 								// Don't cache failures. Next request will retry.
@@ -480,7 +478,8 @@ async fn actor_request_task(
 							}
 						}
 					}
-					let (parsed_id, actor_name) = cached_actor.as_ref().unwrap();
+					let (parsed_id, actor_name) =
+						cached_actors.get(&req.actor_id).unwrap();
 
 					let recipient = actor_kv::Recipient {
 						actor_id: *parsed_id,
@@ -795,10 +794,7 @@ async fn handle_kv_delete_range(
 
 /// Look up an actor by ID and return the parsed ID and actor name.
 ///
-/// Defense-in-depth: verifies the actor belongs to the authenticated namespace.
-/// The admin_token is a global credential, so this is not strictly necessary
-/// today, but prevents cross-namespace access if a less-privileged auth
-/// mechanism is introduced in the future.
+/// Verifies the actor belongs to the authenticated namespace.
 async fn resolve_actor(
 	ctx: &StandaloneCtx,
 	actor_id: &str,
@@ -812,11 +808,15 @@ async fn resolve_actor(
 	})?;
 
 	let actor = ctx
-		.op(pegboard::ops::actor::get_for_runner::Input {
-			actor_id: parsed_id,
+		.op(pegboard::ops::actor::get::Input {
+			actor_ids: vec![parsed_id],
+			fetch_error: false,
 		})
 		.await
-		.map_err(|err| internal_error(&err))?;
+		.map_err(|err| internal_error(&err))?
+		.actors
+		.into_iter()
+		.next();
 
 	match actor {
 		Some(actor) => {

engine/packages/pegboard/src/workflows/actor2/mod.rs

@@ -804,10 +804,15 @@ async fn process_signal(
 					runtime::reschedule_actor(ctx, &input, state, metrics_workflow_id).await?;
 				}
 				Transition::SleepIntent {
-					rewake_after_stop, ..
+					lost_timeout_ts,
+					rewake_after_stop,
+					..
 				} => {
 					if !*rewake_after_stop {
 						*rewake_after_stop = true;
+						let now = ctx.activity(runtime::GetTsInput {}).await?;
+						*lost_timeout_ts =
+							now + ctx.config().pegboard().actor_stop_threshold();
 
 						tracing::debug!(
 							actor_id=?input.actor_id,

engine/packages/pegboard/src/workflows/actor2/runtime.rs

@@ -565,6 +565,40 @@ pub async fn handle_stopped(
 
 		StoppedResult::Continue
 	} else if try_reallocate {
+		// A wake can arrive after the actor has already committed to sleeping
+		// but before the stop completes. In that case, mirror the v1 workflow
+		// and immediately reschedule after the stop instead of treating a
+		// graceful sleep exit as terminal destruction.
+		if let Transition::SleepIntent {
+			rewake_after_stop: true,
+			..
+		} = state.transition
+		{
+			let allocate_res = ctx.activity(AllocateInput {}).await?;
+
+			if let Some(allocation) = allocate_res.allocation {
+				state.generation += 1;
+
+				ctx.activity(SendOutboundInput {
+					generation: state.generation,
+					input: input.input.clone(),
+					allocation,
+				})
+				.await?;
+
+				state.transition = Transition::Allocating {
+					destroy_after_start: false,
+					lost_timeout_ts: allocate_res.now
+						+ ctx.config().pegboard().actor_allocation_threshold(),
+				};
+			} else {
+				state.transition = Transition::Reallocating {
+					since_ts: allocate_res.now,
+				};
+			}
+
+			StoppedResult::Continue
+		} else {
 		// An actor stopping with `StopCode::Ok` indicates a graceful exit
 		let graceful_exit = matches!(
 			variant,
@@ -574,63 +608,64 @@ pub async fn handle_stopped(
 			}
 		);
 
-		match (input.crash_policy, graceful_exit) {
-			(CrashPolicy::Restart, false) => {
-				let allocate_res = ctx.activity(AllocateInput {}).await?;
-
-				if let Some(allocation) = allocate_res.allocation {
-					state.generation += 1;
+			match (input.crash_policy, graceful_exit) {
+				(CrashPolicy::Restart, false) => {
+					let allocate_res = ctx.activity(AllocateInput {}).await?;
+
+					if let Some(allocation) = allocate_res.allocation {
+						state.generation += 1;
+
+						ctx.activity(SendOutboundInput {
+							generation: state.generation,
+							input: input.input.clone(),
+							allocation,
+						})
+						.await?;
+
+						// Transition to allocating
+						state.transition = Transition::Allocating {
+							destroy_after_start: false,
+							lost_timeout_ts: allocate_res.now
+								+ ctx.config().pegboard().actor_allocation_threshold(),
+						};
+					} else {
+						// Transition to retry backoff
+						state.transition = Transition::Reallocating {
+							since_ts: allocate_res.now,
+						};
+					}
 
-					ctx.activity(SendOutboundInput {
-						generation: state.generation,
-						input: input.input.clone(),
-						allocation,
-					})
-					.await?;
-
-					// Transition to allocating
-					state.transition = Transition::Allocating {
-						destroy_after_start: false,
-						lost_timeout_ts: allocate_res.now
-							+ ctx.config().pegboard().actor_allocation_threshold(),
-					};
-				} else {
-					// Transition to retry backoff
-					state.transition = Transition::Reallocating {
-						since_ts: allocate_res.now,
-					};
+					StoppedResult::Continue
 				}
+				(CrashPolicy::Sleep, false) => {
+					tracing::debug!(actor_id=?input.actor_id, "actor sleeping due to ungraceful exit");
 
-				StoppedResult::Continue
-			}
-			(CrashPolicy::Sleep, false) => {
-				tracing::debug!(actor_id=?input.actor_id, "actor sleeping due to ungraceful exit");
-
-				// Clear alarm
-				if let Some(alarm_ts) = state.alarm_ts {
-					let now = ctx.activity(GetTsInput {}).await?;
+					// Clear alarm
+					if let Some(alarm_ts) = state.alarm_ts {
+						let now = ctx.activity(GetTsInput {}).await?;
 
-					if now >= alarm_ts {
-						state.alarm_ts = None;
+						if now >= alarm_ts {
+							state.alarm_ts = None;
+						}
 					}
-				}
 
-				// Transition to sleeping
-				state.transition = Transition::Sleeping;
+					// Transition to sleeping
+					state.transition = Transition::Sleeping;
 
-				StoppedResult::Continue
-			}
-			_ => {
-				let now = ctx.activity(GetTsInput {}).await?;
+					StoppedResult::Continue
+				}
+				_ => {
+					let now = ctx.activity(GetTsInput {}).await?;
 
-				// Don't destroy on failed allocation, retry instead
-				if let StoppedVariant::FailedAllocation { .. } = &variant {
-					// Transition to retry backoff
-					state.transition = Transition::Reallocating { since_ts: now };
+					// Don't destroy on failed allocation, retry instead
+					if let StoppedVariant::FailedAllocation { .. } = &variant {
+						// Transition to retry backoff
+						state.transition = Transition::Reallocating { since_ts: now };
 
-					StoppedResult::Continue
-				} else {
-					StoppedResult::Destroy
+						StoppedResult::Continue
+					} else {
+						StoppedResult::Destroy
+					}
 				}
 			}
 		}