fix(kiloclaw): preserve legacy routing after user ID migration

Author: jeanduplessis

services/kiloclaw/src/db/index.ts

@@ -69,7 +69,14 @@ export async function getActiveInstance(db: WorkerDb, userId: string) {
       organization_id: kiloclaw_instances.organization_id,
     })
     .from(kiloclaw_instances)
-    .where(and(eq(kiloclaw_instances.user_id, userId), isNull(kiloclaw_instances.destroyed_at)))
+    .where(
+      and(
+        eq(kiloclaw_instances.user_id, userId),
+        isNull(kiloclaw_instances.organization_id),
+        isNull(kiloclaw_instances.destroyed_at)
+      )
+    )
+    .orderBy(kiloclaw_instances.created_at)
     .limit(1)
     .then(rows => rows[0] ?? null);
 

services/kiloclaw/src/durable-objects/kiloclaw-instance/index.ts

@@ -69,6 +69,7 @@ import {
   markRestartSuccessful,
 } from './reconcile';
 import { restoreFromPostgres, markDestroyedInPostgresHelper } from './postgres';
+import { legacyDoKeysForIdentity } from '../../lib/instance-routing';
 import {
   beginUnexpectedStopRecovery,
   cleanupPendingRecoveryVolumeIfNeeded,
@@ -1440,22 +1441,23 @@ export class KiloClawInstance extends DurableObject<KiloClawEnv> {
               instanceId: registryInstanceId,
             });
           } else {
-            // Legacy: find active entry by doKey=userId
+            const legacyDoKeys = legacyDoKeysForIdentity(preDestroyUserId, preDestroySandboxId);
             const entries = await registryStub.listInstances(registryKey);
-            const legacyEntry = entries.find(e => e.doKey === preDestroyUserId);
+            const legacyEntry = entries.find(e => legacyDoKeys.includes(e.doKey));
             if (legacyEntry) {
               await registryStub.destroyInstance(registryKey, legacyEntry.instanceId);
               console.log('[DO] Registry entry destroyed on finalization (legacy):', {
                 registryKey,
                 instanceId: legacyEntry.instanceId,
-                doKey: preDestroyUserId,
+                doKeysTried: legacyDoKeys,
+                matchedDoKey: legacyEntry.doKey,
               });
             } else {
               console.log(
                 '[DO] Registry cleanup: no active entry found (already cleaned or never existed):',
                 {
                   registryKey,
-                  doKey: preDestroyUserId,
+                  doKeysTried: legacyDoKeys,
                   activeEntryCount: entries.length,
                 }
               );

services/kiloclaw/src/durable-objects/kiloclaw-instance/postgres.test.ts

@@ -0,0 +1,26 @@
+import { describe, expect, it } from 'vitest';
+import { fallbackAppNameForRestore } from './postgres';
+import { sandboxIdFromUserId } from '../../auth/sandbox-id';
+import { appNameFromUserId, appNameFromInstanceId } from '../../fly/apps';
+
+describe('fallbackAppNameForRestore', () => {
+  it('keeps migrated legacy sandboxes on the acct-* naming path', async () => {
+    const legacyUserId = 'oauth/google:117453785559478190551';
+    const migratedUserId = '199e2b19-aa40-488d-9442-9a18a620ba68';
+
+    await expect(
+      fallbackAppNameForRestore(migratedUserId, sandboxIdFromUserId(legacyUserId))
+    ).resolves.toBe(await appNameFromUserId(legacyUserId));
+  });
+
+  it('keeps ki_ sandboxes on the inst-* naming path', async () => {
+    const instanceId = '11111111-1111-4111-8111-111111111111';
+
+    await expect(
+      fallbackAppNameForRestore(
+        '199e2b19-aa40-488d-9442-9a18a620ba68',
+        'ki_11111111111141118111111111111111'
+      )
+    ).resolves.toBe(await appNameFromInstanceId(instanceId));
+  });
+});

services/kiloclaw/src/durable-objects/kiloclaw-instance/postgres.ts

@@ -12,12 +12,24 @@ import { getAppKey, getFlyConfig } from './types';
 import { storageUpdate } from './state';
 import { attemptMetadataRecovery } from './reconcile';
 import { doError, doWarn, toLoggable, createReconcileContext } from './log';
+import { isInstanceKeyedSandboxId } from '@kilocode/worker-utils/instance-id';
 
 type RestoreOpts = {
   /** If the DO has a stored sandboxId, use it for precise lookup. */
   sandboxId?: string | null;
 };
 
+export async function fallbackAppNameForRestore(
+  userId: string,
+  sandboxId: string,
+  prefix?: string
+): Promise<string> {
+  const appKey = getAppKey({ userId, sandboxId });
+  return isInstanceKeyedSandboxId(sandboxId)
+    ? appNameFromInstanceId(appKey, prefix)
+    : appNameFromUserId(appKey, prefix);
+}
+
 /**
  * Restore DO state from Postgres backup if SQLite was wiped.
  *
@@ -63,10 +75,7 @@ export async function restoreFromPostgres(
     const appKey = getAppKey({ userId, sandboxId: instance.sandboxId });
     const appStub = env.KILOCLAW_APP.get(env.KILOCLAW_APP.idFromName(appKey));
     const prefix = env.WORKER_ENV === 'development' ? 'dev' : undefined;
-    const isInstanceKeyed = appKey !== userId;
-    const fallbackAppName = isInstanceKeyed
-      ? await appNameFromInstanceId(appKey, prefix)
-      : await appNameFromUserId(userId, prefix);
+    const fallbackAppName = await fallbackAppNameForRestore(userId, instance.sandboxId, prefix);
     const recoveredAppName = (await appStub.getAppName()) ?? fallbackAppName;
 
     await ctx.storage.put(

services/kiloclaw/src/durable-objects/kiloclaw-instance/types.test.ts

@@ -0,0 +1,25 @@
+import { describe, expect, it } from 'vitest';
+import { getAppKey } from './types';
+import { sandboxIdFromUserId } from '../../auth/sandbox-id';
+
+describe('getAppKey', () => {
+  it('derives the legacy app owner key from sandboxId instead of migrated userId', () => {
+    const legacyUserId = 'oauth/google:117453785559478190551';
+
+    expect(
+      getAppKey({
+        userId: '199e2b19-aa40-488d-9442-9a18a620ba68',
+        sandboxId: sandboxIdFromUserId(legacyUserId),
+      })
+    ).toBe(legacyUserId);
+  });
+
+  it('keeps instance-keyed sandboxes on the instanceId owner key', () => {
+    expect(
+      getAppKey({
+        userId: '199e2b19-aa40-488d-9442-9a18a620ba68',
+        sandboxId: 'ki_11111111111141118111111111111111',
+      })
+    ).toBe('11111111-1111-4111-8111-111111111111');
+  });
+});

services/kiloclaw/src/durable-objects/kiloclaw-instance/types.ts

@@ -1,6 +1,7 @@
 import type { KiloClawEnv } from '../../types';
 import type { GoogleCredentials, PersistedState, MachineSize } from '../../schemas/instance-config';
 import type { FlyClientConfig } from '../../fly/client';
+import { userIdFromSandboxId } from '../../auth/sandbox-id';
 import {
   isInstanceKeyedSandboxId,
   instanceIdFromSandboxId,
@@ -136,6 +137,14 @@ export function getAppKey(state: { userId: string | null; sandboxId: string | nu
   if (state.sandboxId && isInstanceKeyedSandboxId(state.sandboxId)) {
     return instanceIdFromSandboxId(state.sandboxId);
   }
+  if (state.sandboxId) {
+    try {
+      return userIdFromSandboxId(state.sandboxId);
+    } catch {
+      // Older tests and malformed legacy state can still carry placeholder
+      // sandboxIds that are not reversible base64url encodings.
+    }
+  }
   if (state.userId) return state.userId;
   throw new Error('Cannot derive app key: no sandboxId or userId');
 }

services/kiloclaw/src/durable-objects/kiloclaw-registry.ts

@@ -5,11 +5,8 @@ import { eq, isNull, and } from 'drizzle-orm';
 import migrations from '../../drizzle/migrations';
 import { registryInstances } from '../db/sqlite-schema';
 import { getWorkerDb, getActiveInstance } from '../db';
-import {
-  isInstanceKeyedSandboxId,
-  instanceIdFromSandboxId,
-} from '@kilocode/worker-utils/instance-id';
 import type { KiloClawEnv } from '../types';
+import { doKeyFromActiveInstance } from '../lib/instance-routing';
 
 export type RegistryEntry = {
   instanceId: string;
@@ -209,13 +206,7 @@ export class KiloClawRegistry extends DurableObject<KiloClawEnv> {
       const instance = await getActiveInstance(db, userId);
 
       if (instance) {
-        // Backfill registry entry from Postgres row.
-        // Derive doKey from the row's sandboxId format:
-        // - ki_ prefix → instance-keyed DO at idFromName(instanceId)
-        // - base64url  → legacy DO at idFromName(userId)
-        const doKey = isInstanceKeyedSandboxId(instance.sandboxId)
-          ? instanceIdFromSandboxId(instance.sandboxId)
-          : userId;
+        const doKey = doKeyFromActiveInstance(instance);
         this.db
           .insert(registryInstances)
           .values({

services/kiloclaw/src/index.ts

@@ -25,6 +25,7 @@ import { sandboxIdFromUserId } from './auth/sandbox-id';
 import { InstanceIdParam } from './schemas/instance-config';
 import { isValidInstanceId } from '@kilocode/worker-utils/instance-id';
 import { registerVersionIfNeeded } from './lib/image-version';
+import { resolveDoKeyForUser } from './lib/instance-routing';
 import { startingUpPage } from './pages/starting-up';
 import { buildForwardHeaders } from './utils/proxy-headers';
 import { KILOCLAW_ACTIVE_INSTANCE_COOKIE } from './config';
@@ -349,16 +350,21 @@ async function resolveRegistryEntry(c: Context<AppEnv>) {
     const stub = c.env.KILOCLAW_INSTANCE.get(c.env.KILOCLAW_INSTANCE.idFromName(entry.doKey));
     return { stub, entry };
   } catch (err) {
-    // Registry DO failed. Fall back to the legacy userId-keyed DO.
-    // Only preserves access for legacy instances (doKey = userId).
-    // For instance-keyed DOs (doKey = instanceId), this hits the wrong
-    // (empty) DO — the user sees "not provisioned" until the registry
-    // recovers. Acceptable: a broken registry is transient, and silently
-    // routing to the wrong DO would be worse.
-    console.error('[PROXY] Registry lookup failed, falling back to legacy DO:', err);
-    const stub = c.env.KILOCLAW_INSTANCE.get(c.env.KILOCLAW_INSTANCE.idFromName(userId));
+    console.error(
+      '[PROXY] Registry lookup failed, falling back to Postgres-backed DO lookup:',
+      err
+    );
+    const fallbackDoKey =
+      (await resolveDoKeyForUser(c.env.HYPERDRIVE?.connectionString, userId).catch(fallbackErr => {
+        console.error(
+          '[PROXY] Postgres-backed DO lookup failed, falling back to userId:',
+          fallbackErr
+        );
+        return null;
+      })) ?? userId;
+    const stub = c.env.KILOCLAW_INSTANCE.get(c.env.KILOCLAW_INSTANCE.idFromName(fallbackDoKey));
     const fallbackEntry: RegistryEntry = {
-      doKey: userId,
+      doKey: fallbackDoKey,
       instanceId: '',
       assignedUserId: userId,
       createdAt: '',

services/kiloclaw/src/lib/instance-routing.ts

@@ -0,0 +1,43 @@
+import { getActiveInstance, getWorkerDb } from '../db';
+import { userIdFromSandboxId } from '../auth/sandbox-id';
+import {
+  instanceIdFromSandboxId,
+  isInstanceKeyedSandboxId,
+} from '@kilocode/worker-utils/instance-id';
+
+type ActiveInstanceIdentity = {
+  id: string;
+  sandboxId: string;
+};
+
+export function legacyDoKeysForIdentity(userId: string, sandboxId: string): string[] {
+  const keys = new Set<string>([userId]);
+
+  if (!isInstanceKeyedSandboxId(sandboxId)) {
+    try {
+      keys.add(userIdFromSandboxId(sandboxId));
+    } catch {
+      // Placeholder sandboxIds can exist in old tests or malformed state.
+    }
+  }
+
+  return [...keys];
+}
+
+export function doKeyFromActiveInstance(instance: ActiveInstanceIdentity): string {
+  return isInstanceKeyedSandboxId(instance.sandboxId)
+    ? instanceIdFromSandboxId(instance.sandboxId)
+    : userIdFromSandboxId(instance.sandboxId);
+}
+
+export async function resolveDoKeyForUser(
+  connectionString: string | undefined,
+  userId: string
+): Promise<string | null> {
+  if (!connectionString) return null;
+
+  const instance = await getActiveInstance(getWorkerDb(connectionString), userId);
+  if (!instance) return null;
+
+  return doKeyFromActiveInstance(instance);
+}