Rehydrate per-session audit verify key on operator restart

pkg/controllers/agentsession/controller.go

@@ -2042,9 +2042,10 @@ func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Resu
 
 	// Register the per-session audit public key for verify-on-write. Read from
 	// status every reconcile — cheap and idempotent — so a restarted operator
-	// re-registers a key it never minted.
-	if pub, decErr := provenance.DecodePubKey(sess.Status.AuditPublicKey); decErr == nil {
-		r.Tokens.SetPublisherKey(provenance.SessionPublisher(sess.Namespace, sess.Name), sess.Status.AuditKeyID, pub)
+	// re-registers a key it never minted. reregisterMemoryToken does the same at
+	// the top of Reconcile, through the same helper, so a session that
+	// short-circuits before here still verifies its own writes.
+	if r.registerAuditVerifyKey(&sess) {
 		// …and witness the same binding DURABLY, in the session's own scope. The
 		// registration above is process memory and the status field above it dies
 		// with the CR, while the records this key signs are permanent and sit in a

pkg/controllers/agentsession/memorytoken.go

@@ -9,11 +9,13 @@ import (
 
 	spiceboxv1alpha1 "github.com/authzed/openagentprimitives/pkg/apis/v1alpha1"
 	"github.com/authzed/openagentprimitives/pkg/memory"
+	"github.com/authzed/openagentprimitives/pkg/memory/provenance"
 )
 
-// reregisterMemoryToken re-installs this session's memory-API bearer token in
-// the operator's in-process token registry, reading the value back out of the
-// per-session Secret that already holds it.
+// reregisterMemoryToken re-installs this session's memory-API bearer token AND
+// its audit verify key in the operator's in-process token registry, reading the
+// token back out of the per-session Secret that already holds it and the key off
+// the status the operator anchored it on.
 //
 // WHY it exists. tokens.Registry is process memory — plain maps, nothing
 // rehydrates them at startup — while the <session>-memory-token Secret and
@@ -97,4 +99,41 @@ func (r *Reconciler) reregisterMemoryToken(ctx context.Context, sess *spiceboxv1
 	r.Tokens.Set(key, token, "", extras...)
 	log.FromContext(ctx).Info("memory token: restored this session's registration from its Secret",
 		"session", sess.Namespace+"/"+sess.Name, "extraScopes", len(extras))
+
+	// Restore the per-session audit VERIFY key in the same breath as the token.
+	// Both live in this process-memory registry and both die on restart, but the
+	// token was restored here — above every short-circuit — while the key was
+	// registered only at step 4 of Reconcile, ~900 lines and ~30 early returns
+	// below. The gap meant a restarted operator accepted the session's bearer (no
+	// 401) yet held no key to verify what that bearer signed, so every
+	// append-only write failed `403 ... no usable key <keyID>`. A terminal
+	// session, whose reap returns before step 4 forever, could never recover.
+	//
+	// It is best-effort: a session whose status carries no key yet has not
+	// completed a full reconcile, and step 4 registers it (and witnesses it
+	// durably) once reached. Unlike the token, restoring the key never widens
+	// anything — a verify key only lets the facade check a signature it would
+	// otherwise reject.
+	r.registerAuditVerifyKey(sess)
+}
+
+// registerAuditVerifyKey installs this session's audit public key into the
+// in-process verify-on-write registry, read from status.auditPublicKey/
+// auditKeyID — the K8s-witnessed trust root. It is the single source both the
+// restart-restoration above and step 4 of Reconcile register from, so
+// verify-on-write cannot disagree with itself across the two call sites.
+//
+// Idempotent, and a no-op when status carries no decodable key yet (a session
+// that has not completed a full reconcile). Returns whether a key was
+// registered, so the caller holding the context can chain the durable witness
+// (reregisterMemoryToken deliberately does not — a terminal session was already
+// witnessed during its life, and restoring the in-memory verify key is the only
+// thing a restart actually lost).
+func (r *Reconciler) registerAuditVerifyKey(sess *spiceboxv1alpha1.AgentSession) bool {
+	pub, err := provenance.DecodePubKey(sess.Status.AuditPublicKey)
+	if err != nil {
+		return false
+	}
+	r.Tokens.SetPublisherKey(provenance.SessionPublisher(sess.Namespace, sess.Name), sess.Status.AuditKeyID, pub)
+	return true
 }

pkg/controllers/agentsession/memorytoken_test.go

@@ -27,6 +27,7 @@ import (
 
 	spiceboxv1alpha1 "github.com/authzed/openagentprimitives/pkg/apis/v1alpha1"
 	"github.com/authzed/openagentprimitives/pkg/memory"
+	"github.com/authzed/openagentprimitives/pkg/memory/provenance"
 	"github.com/authzed/openagentprimitives/pkg/memory/tokens"
 )
 
@@ -89,6 +90,54 @@ func TestReconcile_SucceededSession_RestoresMemoryTokenAfterOperatorRestart(t *t
 	assert.True(t, f.r.Tokens.AuthorizesMutation(tok, key), "restored token keeps the write reach it had before the restart")
 }
 
+// TestReconcile_SucceededSession_RestoresAuditVerifyKeyAfterOperatorRestart is
+// the sibling regression to the memory-token restore above, for the OTHER half
+// of the per-session registration a restart wipes: the Ed25519 audit VERIFY key
+// the facade resolves on every append-only write.
+//
+// Both live in the same process-memory registry. The token was restored on
+// every short-circuit path (reregisterMemoryToken, at the top of Reconcile) but
+// the verify key was registered only at step 4, ~900 lines and ~30 early
+// returns below. So a terminal session — whose reap returns long before step 4,
+// forever — got its token back but not its key, and every append-only write
+// (and every provenance-verifying read) then failed with
+// `403 ... no usable key <keyID> for session:...`, exactly the error a
+// production operator OOMKill produced.
+//
+// The registry starts empty ON PURPOSE — that, not the phase, is the restart.
+func TestReconcile_SucceededSession_RestoresAuditVerifyKeyAfterOperatorRestart(t *testing.T) {
+	const tok = "restored-bearer-value"
+	ctx := memory.WithSystemApproval(context.Background(), "test")
+	f := succeededFixture(t, tok)
+
+	// The session's durable audit identity, as it stands on the CR after the
+	// full reconcile that minted it — BEFORE the restart. generateAuditKeypair
+	// is the operator's own mint, so status carries a real (publicKey, keyID).
+	_, pubB64, keyID, err := generateAuditKeypair()
+	require.NoError(t, err, "mint the session's audit keypair")
+	f.sess.Status.AuditPublicKey = pubB64
+	f.sess.Status.AuditKeyID = keyID
+	require.NoError(t, f.c.Status().Update(ctx, f.sess),
+		"persist the audit key on status (the K8s-witnessed trust root the operator registers from)")
+
+	publisher := provenance.SessionPublisher(f.sess.Namespace, f.sess.Name)
+	_, had := f.r.Tokens.PublisherKey(publisher, keyID)
+	require.False(t, had, "precondition: a restarted operator holds no verify key for this session")
+
+	f.reconcile(t)
+
+	// The reap ran, so this reconcile really did take the short-circuit that
+	// step 4's key registration sits below — the same gate the token test pins.
+	assert.True(t, f.sandboxGone(t), "the terminal reap must have run: this is the short-circuiting path")
+
+	gotPub, ok := f.r.Tokens.PublisherKey(publisher, keyID)
+	require.True(t, ok,
+		"the session's audit verify key must be re-registered after the reconcile, or append-only writes 403 with 'no usable key'")
+	wantPub, decErr := provenance.DecodePubKey(pubB64)
+	require.NoError(t, decErr, "decode the status public key")
+	assert.Equal(t, wantPub, gotPub, "the restored verify key must be the one anchored on status")
+}
+
 // TestReconcile_RestoredTokenScopes pins that the restoration reaches exactly
 // what the registration it replaces reached — the session plus the bundle
 // SpiceboxSessions recorded on status, for READS only — and nothing else.