Restore KeyedSavedStateRegistryOwners just in time, before consumers can see them [UISA-95]

[zach-klippenstein]

Problem

WorkflowSavedStateRegistryAggregator restores its KeyedSavedStateRegistryOwner children from an observer on the parent SavedStateRegistryOwner's lifecycle. Observer dispatch order on a LifecycleRegistry follows registration order, and WorkflowLifecycleOwner lifecycles advance in a synchronous depth-first cascade on window attach — so a descendant's lifecycle observers can run before the aggregator's restore observer ever fires.

Concretely: a nested ComposeView's WrappedComposition consumes the view-tree SavedStateRegistry on ON_CREATE. If that dispatch wins the race against the aggregator's restore, the consumer reads an unrestored registry and throws:

IllegalStateException: You can 'consumeRestoredStateForKey' only after the corresponding
component has moved to the 'CREATED' state

On androidx.compose.ui 1.9.5 this is fatal in a second way: the exception is thrown mid-way through LayoutNode.attach()'s recursion during a lazy-item/pager subcompose, which leaves ancestor LayoutNodes half-attached (markAsAttached ran, runAttachLifecycle skipped). The tree then detonates at composition dispose with:

IllegalStateException: Must run runDetachLifecycle() once after runAttachLifecycle() and before markAsDetached()

Compose 1.7.x tolerated the poisoned state; 1.9.5 crashes the process on teardown. Reproduced deterministically (4/4) in Square's android-register with a workflow-hosted Compose screen containing a lazy list (DemoV2HomeTest#payTabAndButtonAppears_withBillPayEnabled); this crash currently blocks register's Compose 1.9.5 upgrade. (A minimal, workflow-free Compose repro of the downstream detonation exists and is being filed against androidx separately — but the state-restoration race fixed here is workflow's own bug regardless.)

Fix

Two cooperating parts, both local to workflow-ui/core-android's androidx integration:

1. Lifecycle clamp. KeyedSavedStateRegistryOwner no longer delegates its Lifecycle to the view's WorkflowLifecycleOwner. It now exposes its own LifecycleRegistry that mirrors the delegate's, but is held at INITIALIZED until the owner's registry has been restored. This enforces the androidx savedstate contract by construction: lifecycle-driven consumers (like Compose's DisposableSaveableStateRegistry) cannot observe an unrestored registry, and performRestore() remains legal no matter how far the delegate lifecycle has advanced (savedstate requires lifecycle < STARTED).

2. Just-in-time restore. When the delegate delivers ON_CREATE and the owner has not been restored yet, the owner synchronously asks the aggregator to restore it right then (restoreChildNow) — before the local lifecycle advances and any downstream consumer can see the registry. If the aggregator's own restore observer hasn't fired yet, it pulls its state from the parent registry on demand (consumeRestoredStateForKey only requires the parent registry to be restored, not the parent lifecycle to be CREATED). Children pruned from the aggregator restore empty rather than stealing a successor's saved state.

Because ancestors' lifecycles always advance before descendants', an owner's JIT restore can always find its parent's state — restoration order is preserved top-down and no state is lost.

This also structurally fixes #570: a child whose lifecycle is past INITIALIZED can now always be restored, because its local lifecycle cannot pass INITIALIZED unrestored.

API note

installChildRegistryOwnerOn now registers a lifecycle observer at install time, which LifecycleRegistry requires to happen on the main thread. This was already true of every production caller (container views during render); the KDoc now documents it, and ViewStateCacheTest routes its update() calls through runOnMainSync accordingly.

Testing

• New JVM regression tests in WorkflowSavedStateRegistryAggregatorTest covering: the dispatch-order race (child advancing before parent ON_CREATE is restored just in time with saved state), the lifecycle clamp, the detached-parent fallback (restores empty without throwing), and the pruned-child case (does not consume a replacement's state).
• :workflow-ui:core-android:connectedDebugAndroidTest: 39/39 pass (incl. BackStackContainerPersistenceLifecycleTest, ViewStateCacheTest, WorkflowViewStubLifecycleTest).
• :workflow-ui:compose:connectedDebugAndroidTest: 56/56 pass (incl. ComposeViewTreeIntegrationTest save/restore integration).
• Validated against the android-register repro via composite build: previously 4/4 crash, now 3/3 pass with zero occurrences of either crash signature in logcat.

🤖 Generated with Claude Code

[zach-klippenstein]

🤖 Validation complete — marking ready for review.

Evidence this fix is safe and effective, gathered against Square's android-register monorepo (where the crash was found):

1. Original crash repro: DemoV2HomeTest#payTabAndButtonAppears_withBillPayEnabled crashed 4/4 before this change (Must run runDetachLifecycle() once after runAttachLifecycle() and before markAsDetached() at teardown, Compose 1.9.5); with this branch it passes, along with the rest of its suite (6/6, composite build against register green-master).
2. No state-restoration regressions: workflow-ui:core-android device suites pass (ViewStateCacheTest, WorkflowSavedStateRegistryAggregatorTest instrumented: 39/39 and 56/56), plus the 4 new JVM regression tests in WorkflowSavedStateRegistryAggregatorTest covering the dispatch-order race, the lifecycle clamp, detached-parent fallback, and pruned-child restore.
3. Full register CI: a snapshot build of this branch (1.28.1-SNAPSHOT) ran register's entire suite — 1575/1576 parts green, zero test failures (the one red part is register's policy gate that forbids SNAPSHOT dependency versions, by design). Validation PR: squareup/android-register#159654.

Root-cause writeup is in the PR description; the underlying Compose-side issue (exception during LayoutNode.attach() recursion permanently corrupting the node tree) is being filed upstream separately.

[steve-the-edwards]

🤖 The class KDoc says the local lifecycle "is held at INITIALIZED until this owner's registry has been restored" — but that invariant is enforced by forcing restore here, not by withholding advancement. If onRestoreNeeded ever fails its contract (returns without restoring — say a future refactor breaks the children.put-before-installObserver ordering in installChildRegistryOwnerOn), the local lifecycle advances unrestored and we're back to the original downstream crash, far from the cause. A cheap check makes the clamp self-enforcing and fails loudly at the source:

suggestion

if (event == ON_CREATE && !savedStateRegistry.isRestored) {
   onRestoreNeeded(this@KeyedSavedStateRegistryOwner)
   check(savedStateRegistry.isRestored) {
      "onRestoreNeeded contract violation: registry for key '$key' was not restored"
   }
}

[zach-klippenstein]

🤖 Done in 74c72b3 — added the check(savedStateRegistry.isRestored) contract assertion right after the onRestoreNeeded call, plus a unit test (KeyedSavedStateRegistryOwnerTest) covering a callee that violates the contract.

[steve-the-edwards]

🤖 Nit: broken KDoc reference — the method is installObserver, not installObserverOn.

* any downstream consumer can see the registry. [installObserver] is called by the aggregator

[zach-klippenstein]

🤖 Fixed in 74c72b3 — [installObserverOn] → [installObserver].

[steve-the-edwards]

🤖 Nit: ViewStateCache.update is the public entry point most custom-container authors actually call, and it now inherits this main-thread requirement (as ViewStateCacheTest's runOnMainSync change shows). Worth a sentence in its KDoc too.

[zach-klippenstein]

🤖 Done in 74c72b3 — ViewStateCache.update's KDoc now states the main-thread / before-attach requirement and points at installChildRegistryOwnerOn as the source of it.

[steve-the-edwards]

🤖 Nit: this branch trades the old crash for silent state loss in the (theoretically unreachable in production) case where a child advances with no attached/restored parent. That's the right trade, but a debug-level log here would help diagnose any container that ever hits it.

[zach-klippenstein]

🤖 Done in 74c72b3 — the restore-empty fallback now logs at debug level (tag Workflow, matching the module's existing logging), including the child key and which of the two conditions was hit, so a container that reaches this branch leaves a breadcrumb for the lost state.