fix(rivetkit): use startup kv preload

[NathanFlurry]

Description

Please include a summary of the changes and the related issue. Please also include relevant motivation and context.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

How Has This Been Tested?

Please describe the tests that you ran to verify your changes.

Checklist:

• My code follows the style guidelines of this project
• I have performed a self-review of my code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[claude]

PR #4956 Review: fix(rivetkit): use startup kv preload

Overview

This PR fixes a bug where KV preload data was being ignored during actor startup. Three distinct issues are corrected:

1. Oversized exact keys were incorrectly marked as preloaded (scanned) even when they exceeded the byte budget and were not included in the preload bundle.
2. The startup path in ActorTask was not consulting preloaded_kv when loading LAST_PUSHED_ALARM_KEY, PERSIST_DATA_KEY, and the inspector token, always live-fetching them.
3. set_has_initialized was unconditionally dirtying the state metrics even when the value did not change.

---

Code Quality

engine/packages/pegboard/src/actor_kv/preload.rs The core fix is clean and correct. Using if let Some(b) = builder to distinguish missing keys (stream yields nothing, builder stays None) from over-budget keys (builder built, size check fails) is the right idiom. The debug log on budget exhaustion is helpful.

Minor indentation artifact in the diff: the for-loop closing brace appears on its own line without matching context. The code is correct but rustfmt should clean this up.

rivetkit-rust/packages/rivetkit-core/src/actor/task.rs The new load_startup_key helper is a clean pattern. The preloaded_kv.is_some() branch in load_persisted_startup correctly handles the case where KV was preloaded but no pre-decoded actor bundle was provided.

One subtle issue in load_startup_key:

if let Some(entry) = self
    .preloaded_kv
    .as_ref()
    .and_then(|preloaded| preloaded.key_entry(key))
{
    return Ok(entry);
}

key_entry returns Option<Option<Vec<u8>>>:

• Some(Some(bytes)): key present and in budget. and_then returns Some(bytes). Early return. Correct.
• Some(None): key confirmed absent (scanned, missing from FDB). and_then collapses to None. Falls through to live KV fetch. Suboptimal.
• None: key not in preload scope. Falls through to live KV fetch. Correct.

For confirmed-absent keys (PERSIST_DATA_KEY on a fresh actor with no preloaded bundle), the code does an unnecessary live fetch. The result is functionally correct (both paths return None), but it wastes a round-trip. The fix is to check scope membership separately from value presence before calling into live KV.

---

Potential Bugs

rivetkit-typescript/packages/rivetkit/src/actor/keys.ts Both LAST_PUSHED_ALARM and WORKFLOW_PREFIX are assigned Uint8Array.from([6]):

LAST_PUSHED_ALARM: Uint8Array.from([6]),
WORKFLOW_PREFIX: Uint8Array.from([6]),   // same byte value

This mirrors the Rust side (LAST_PUSHED_ALARM_KEY = &[6] and WORKFLOW_PREFIX = [6] with actual workflow data at [6, 1, ...]), so it is intentional. Exact key [6] holds the alarm timestamp; prefix [6, ...] holds workflow data. They do not collide at the storage level.

However, two exported constants with identical byte values in the same object is a landmine. A future refactor or iteration over KEYS values could silently produce duplicate byte sequences. This should have a comment explaining the intentional byte overlap.

---

Tests

The new tests are well-structured and cover the critical cases:

• preload_oversized_exact_key_is_not_marked_requested: verifies the engine-side fix directly.
• preload_missing_exact_key_is_marked_requested: verifies the confirmed-absent path.
• startup_uses_preloaded_last_pushed_alarm_without_live_kv: verifies alarm is read from preload when all KV is unavailable.
• startup_decodes_persisted_actor_from_preloaded_kv_without_live_kv: verifies persisted actor decoded from preloaded KV.
• startup_uses_preloaded_alarm_with_partial_startup_preload: verifies only the missing key triggers a live fetch.

The EnvoyShutdownError approach for mocking unavailable KV gives strong guarantees. The saw_persist_live_get assertion in the partial-preload test is particularly good: it confirms exactly one live fetch occurred and it was the right key.

One gap: there is no test for the Some(None) path described above (key in preload scope, confirmed absent, no pre-decoded actor bundle). Adding a test where PERSIST_DATA_KEY is in requested_get_keys but not in entries, with a mock that asserts the live fetch still occurs, would document the current behavior and catch any future change.

---

Minor Notes

• set_has_initialized early return is a good defensive optimization. No concerns.
• Adding [6] (LAST_PUSHED_ALARM_KEY) to the startup preload key list in registry/mod.rs is consistent with the TypeScript buildActorNames change.
• init_inspector_token_with_preload is correctly scoped as pub(crate) with init_inspector_token forwarding to it. Good layering.

[NathanFlurry]

• fix(pegboard): align drain grace defaults with stop threshold #4958 : 2 dependent PRs (#4960 , #4966 )
• fix(rivetkit): preserve bridge error messages #4959
• fix(rivetkit): fire abort signal on shutdown finalize #4957
• fix(rivetkit): use startup kv preload #4956 👈 (View in Graphite)
• fix(client): flatten skipReadyWait option #4955
• ci(publish): add manual preview publish flow #4954
• fix(depot): add ltx decode diagnostics #4953
• fix(universaldb): preserve plain set values #4952
• fix(universaldb): align versionstamp operand encoding #4951
• fix(rivetkit): export database provider types #4950
• wip: kitchen sink agentic loop finish #4949
• chore(rivetkit): note duplicate admission logic #4948
• test(rivetkit): cover bypass fetch header forwarding #4947
• feat(kitchen-sink): add actor probe controls #4946
• fix(gateway): split actor response wait errors #4945
• fix(guard): handle bypass preflight before query parse #4944
• feat(kitchen-sink): add mock agent event log copy #4943
• fix(kitchen-sink): default mock agent to local endpoint #4942
• test(rivetkit): enable driver lifecycle regressions #4941
• test(rivetkit): verify cached websocket db refs exceed grace #4940
• fix(rivetkit-core): keep sleeping actors reachable during grace #4939
• test(rivetkit): cover async websocket close db access #4938
• fix(rivetkit): keep actors awake until keepAwake work finishes #4937
• fix(rivetkit): persist native onSleep state on errors #4936
• fix(rivetkit-core): allow dispatch during sleep grace #4935
• fix(rivetkit): make duplicate sleep requests idempotent #4934
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.