fix(server): wake chat milestones after run identity activation

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/__tests__/chat-run-milestone-identity-work.test.ts

@@ -0,0 +1,104 @@
+import { randomUUID } from "node:crypto";
+import { eq } from "drizzle-orm";
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+import { agentWakeupRequests, chatPublications, createDb, heartbeatRuns, issueComments, subscribeDatabaseWork } from "@paperclipai/db";
+import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js";
+import { resolveExternalTestDatabaseUrl } from "./helpers/external-test-database.js";
+import { createChatMilestoneFixtures } from "./helpers/chat-milestone-fixtures.js";
+import { createDeliveryWorkCoordinator } from "../services/delivery-work-coordinator.js";
+import { DELIVERY_QUEUES } from "../services/delivery-work-notifications.js";
+import { chatRunMilestoneWork } from "../services/chat-run-milestone-work.js";
+import { initializeRunIdentity, reserveSteeredIdentity, acceptSteeredIdentity } from "../services/run-identity.js";
+import { idleWorkSnapshot } from "../services/task-admission.js";
+
+const external = resolveExternalTestDatabaseUrl("PAPERCLIP_MILESTONE_IDENTITY_TEST_DATABASE_URL");
+const support = external ? { supported: true } : await getEmbeddedPostgresTestSupport();
+(support.supported ? describe : describe.skip)("chat milestones after identity activation", () => {
+  let db: ReturnType<typeof createDb>;
+  let temporary: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | undefined;
+  let fixtures: ReturnType<typeof createChatMilestoneFixtures>;
+  let f: Awaited<ReturnType<typeof fixtures.create>>;
+  let coordinator: ReturnType<typeof createDeliveryWorkCoordinator> | undefined;
+  const errors: unknown[] = [];
+  beforeAll(async () => {
+    if (!external) temporary = await startEmbeddedPostgresTestDatabase("chat-identity-milestones-");
+    db = createDb(external ?? temporary!.connectionString, { idleTimeoutSeconds: 0, maxLifetimeSeconds: 0 });
+    fixtures = createChatMilestoneFixtures(db);
+  }, 60_000);
+  beforeEach(async () => {
+    errors.length = 0; f = await fixtures.create();
+    vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"], shouldAdvanceTime: true });
+  });
+  afterEach(async () => {
+    await coordinator?.stop(); coordinator = undefined; vi.useRealTimers(); vi.restoreAllMocks();
+    await fixtures.cleanup(); expect(errors).toEqual([]);
+  });
+  afterAll(async () => { await db?.$client.end({ timeout: 1 }); await temporary?.cleanup(); });
+  async function flush() {
+    await vi.waitFor(async () => {
+      await vi.advanceTimersByTimeAsync(1000);
+      expect(idleWorkSnapshot().active).toBe(0); expect(coordinator!.nextWakeAt()).toBeNull();
+    }, { timeout: 10_000 });
+  }
+  async function quiet() {
+    const reads = vi.spyOn(db, "select"), writes = vi.spyOn(db, "transaction");
+    try { await vi.advanceTimersByTimeAsync(120_000); expect(reads).not.toHaveBeenCalled(); expect(writes).not.toHaveBeenCalled(); }
+    finally { reads.mockRestore(); writes.mockRestore(); }
+  }
+  const publications = () => db.select().from(chatPublications).where(eq(chatPublications.issueId, f.issueId));
+
+  it.each(["initialization", "steering"] as const)("wakes matching monitor lineage after %s commits, preserving rollback and replay", async kind => {
+    const nextCheckAt = new Date(Date.now() - 1000).toISOString();
+    const sourceId = randomUUID(), successorId = randomUUID();
+    await db.insert(heartbeatRuns).values({ id: sourceId, companyId: f.companyId, agentId: f.agentId,
+      nativeIssueId: f.issueId, runtimeMode: "native", status: "succeeded", responsibleUserId: f.userId,
+      contextSnapshot: { issueId: f.issueId, source: "chat:slack", endpointId: f.endpointId, wakeCommentId: f.commentId },
+      resultJson: { semanticToolReceipts: { schedule: { operationId: "set_task_monitor",
+        result: { taskId: f.issueId, monitor: { status: "scheduled", nextCheckAt } } } } },
+    });
+    await db.insert(heartbeatRuns).values({ id: successorId, companyId: f.companyId, agentId: f.agentId,
+      nativeIssueId: f.issueId, runtimeMode: "native", status: "running", responsibleUserId: null,
+      contextSnapshot: { issueId: f.issueId, source: "issue.monitor" },
+    });
+    await db.insert(agentWakeupRequests).values({ companyId: f.companyId, agentId: f.agentId, runId: successorId,
+      source: "automation", reason: "issue_monitor_due", status: "claimed",
+      idempotencyKey: `issue-monitor:${f.issueId}:${nextCheckAt}`, payload: { issueId: f.issueId, nextCheckAt },
+    });
+    const identityInput = { companyId: f.companyId, runId: successorId, issueId: f.issueId, responsibleUserId: f.userId, cause: "dispatch" };
+    if (kind === "steering") {
+      await initializeRunIdentity(db, { ...identityInput, responsibleUserId: null });
+      await db.update(issueComments).set({ authorUserId: f.userId }).where(eq(issueComments.id, f.commentId));
+    }
+    coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => true, canReconcile: () => true, onError: error => errors.push(error) });
+    await chatRunMilestoneWork(db).register(coordinator, () => true).ready;
+    await flush(); expect(await publications()).toEqual([]); await quiet();
+    const pending = kind === "steering" ? await reserveSteeredIdentity(db, {
+      companyId: f.companyId, runId: successorId, issueId: f.issueId, messageId: f.commentId,
+    }) : null;
+    // Reserving an unacknowledged message must not activate its identity.
+    if (kind === "steering") expect(pending?.status).toBe("pending");
+    await flush(); expect(await publications()).toEqual([]); await quiet();
+    const activate = (rollback: boolean) => db.transaction(async tx => {
+      if (kind === "steering") await acceptSteeredIdentity(tx, pending!);
+      else await initializeRunIdentity(tx as unknown as typeof db, identityInput);
+      await vi.advanceTimersByTimeAsync(1000);
+      expect(await publications()).toEqual([]);
+      if (rollback) throw new Error("outer rollback");
+    });
+    await expect(activate(true)).rejects.toThrow("outer rollback");
+    await flush(); expect(await publications()).toEqual([]); await quiet();
+    await activate(false); await flush();
+    expect(await publications()).toEqual([expect.objectContaining({ idempotencyKey: `run:${successorId}:working:${f.endpointId}` })]);
+    await quiet();
+    const notices: string[] = [];
+    const unsubscribe = subscribeDatabaseWork(db, DELIVERY_QUEUES.chatRunMilestones, event => notices.push(event));
+    try {
+      await db.transaction(async tx => {
+        if (kind === "steering") await acceptSteeredIdentity(tx, pending!);
+        else await initializeRunIdentity(tx as unknown as typeof db, identityInput);
+      });
+      expect(notices).toEqual([]);
+    } finally { unsubscribe(); }
+    await flush(); await quiet();
+  });
+});

server/src/__tests__/github-operation-credentials.test.ts

@@ -1,3 +1,4 @@
+import { resolveExternalTestDatabaseUrl } from "./helpers/external-test-database.js";
 import { createNativeGitHubAccess } from "../services/native-runtime/native-github-access.js";
 import express from "express";
 import request from "supertest";
@@ -32,7 +33,7 @@ import {
 import {
   initializeRunIdentity,
   reserveSteeredIdentity,
-  acceptSteeredIdentity,
+  reconcileSteeredIdentity,
 } from "../services/run-identity.js";
 import { resolveGitHubOperationCredentials } from "../services/github-operation-credentials.js";
 import {
@@ -49,23 +50,25 @@ const vault = vi.hoisted(() => ({
   resolveSecretValue: vi.fn(async () => "test-dedicated-token"),
 }));
 vi.mock("../services/secrets.js", () => ({ secretService: () => vault }));
-const support = await getEmbeddedPostgresTestSupport();
+const external = resolveExternalTestDatabaseUrl("PAPERCLIP_GITHUB_CREDENTIAL_TEST_DATABASE_URL");
+const support = external ? { supported: true } : await getEmbeddedPostgresTestSupport();
 (support.supported ? describe : describe.skip)(
   "operation-time GitHub credential resolution",
   () => {
-    let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>,
+    let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | undefined,
       db: ReturnType<typeof createDb>;
     beforeAll(async () => {
       vi.stubEnv(
         "PAPERCLIP_AGENT_JWT_SECRET",
         "test-github-broker-signing-secret",
       );
-      database = await startEmbeddedPostgresTestDatabase(
+      if (!external) database = await startEmbeddedPostgresTestDatabase(
         "paperclip-github-operation-",
       );
-      db = createDb(database.connectionString);
+      db = createDb(external ?? database!.connectionString);
     }, 30_000);
     afterAll(async () => {
+      await db?.$client.end({ timeout: 1 });
       await database?.cleanup();
       vi.unstubAllEnvs();
     }, 60_000);
@@ -208,7 +211,7 @@ const support = await getEmbeddedPostgresTestSupport();
         ...input,
         messageId: id,
       });
-      await acceptSteeredIdentity(db, context!);
+      await reconcileSteeredIdentity(db, context!);
     }
     it("resolves A → B → A without retaining tokens, and records only redacted diagnostics", async () => {
       const input = await seed();

server/src/__tests__/run-identity.test.ts

@@ -1,19 +1,21 @@
+import { resolveExternalTestDatabaseUrl } from "./helpers/external-test-database.js";
 import { randomUUID } from "node:crypto";
 import { eq, sql } from "drizzle-orm";
 import { afterAll, beforeAll, describe, expect, it } from "vitest";
 import { agentWakeupRequests, agents, companies, createDb, heartbeatRuns, heartbeatRunEvents, issueComments, issueThreadInteractions, issues, secretAccessEvents } from "@paperclipai/db";
 import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js";
-import { acceptSteeredIdentity, captureRunIdentity, initializeRunIdentity, listRunIdentityContexts, rejectSteeredIdentity, reserveSteeredIdentity } from "../services/run-identity.js";
+import { acceptSteeredIdentity, reconcileSteeredIdentity, captureRunIdentity, initializeRunIdentity, listRunIdentityContexts, rejectSteeredIdentity, reserveSteeredIdentity } from "../services/run-identity.js";
 
-const support = await getEmbeddedPostgresTestSupport();
+const external = resolveExternalTestDatabaseUrl("PAPERCLIP_RUN_IDENTITY_TEST_DATABASE_URL");
+const support = external ? { supported: true } : await getEmbeddedPostgresTestSupport();
 (support.supported ? describe : describe.skip)("durable execution identity", () => {
-  let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
+  let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | undefined;
   let db: ReturnType<typeof createDb>;
   beforeAll(async () => {
-    database = await startEmbeddedPostgresTestDatabase("paperclip-run-identity-");
-    db = createDb(database.connectionString);
+    if (!external) database = await startEmbeddedPostgresTestDatabase("paperclip-run-identity-");
+    db = createDb(external ?? database!.connectionString);
   }, 30_000);
-  afterAll(async () => { await database?.cleanup(); }, 60_000);
+  afterAll(async () => { await db?.$client.end({ timeout: 1 }); await database?.cleanup(); }, 60_000);
 
   async function seed() {
     const companyId = randomUUID(), agentId = randomUUID(), issueId = randomUUID(), runId = randomUUID();
@@ -121,8 +123,8 @@ const support = await getEmbeddedPostgresTestSupport();
     expect((await captureRunIdentity(db, input)).context?.responsibleUserId).toBe("B");
     expect(startedUnderA.context?.responsibleUserId).toBe("A");
     const a = await reserveSteeredIdentity(db, { ...input, messageId: input.messageIds[2] });
-    await acceptSteeredIdentity(db, a!);
-    await acceptSteeredIdentity(db, b!);
+    await reconcileSteeredIdentity(db, a!);
+    await reconcileSteeredIdentity(db, b!);
     expect((await captureRunIdentity(db, input)).context?.responsibleUserId).toBe("A");
   });
   it("keeps rejected steering unchanged and preserves the originating identity on retry", async () => {
@@ -180,7 +182,7 @@ const support = await getEmbeddedPostgresTestSupport();
       sourceIdentityContextId: origin.id, payload: { prompt: "Continue?" } as never,
     });
     const b = await reserveSteeredIdentity(db, { ...input, messageId: input.messageIds[1] });
-    await acceptSteeredIdentity(db, b!);
+    await reconcileSteeredIdentity(db, b!);
     for (const approval of [true, false]) {
       const runId = randomUUID();
       await db.insert(heartbeatRuns).values({ id: runId, companyId: input.companyId, agentId: input.agentId, status: "running" });
@@ -251,34 +253,40 @@ const support = await getEmbeddedPostgresTestSupport();
       // A real append holds KEY SHARE on both parent rows until it commits.
       // Bound the other connection's wait so a conflicting lock fails this
       // regression instead of leaving both transactions waiting for each other.
-      const identityDb = createDb(`${database.connectionString}?options=-c%20lock_timeout%3D1000`, {
+      const identityUrl = new URL(external ?? database!.connectionString);
+      identityUrl.searchParams.set("options", "-c lock_timeout=1000");
+      const identityDb = createDb(identityUrl.toString(), {
         maxConnections: 1,
       });
-      const [settings] = await identityDb.execute(sql`show lock_timeout`);
-      expect(settings?.lock_timeout).toBe("1s");
-      await db.transaction(async (audit) => {
-        await audit.insert(secretAccessEvents).values({
-          companyId: input.companyId,
-          heartbeatRunId: input.runId,
-          issueId: input.issueId,
-          provider: "local_encrypted",
-          actorType: "agent",
-          actorId: input.agentId,
-          consumerType: "agent",
-          consumerId: input.agentId,
-          outcome: "granted",
-        });
-        if (operation === "initialize") {
-          await expect(initializeRunIdentity(identityDb, {
-            ...input, messageIds: [], responsibleUserId: "A", cause: "instruction",
-          })).resolves.toMatchObject({ responsibleUserId: "A" });
-        } else {
-          await expect(captureRunIdentity(identityDb, input)).resolves.toMatchObject({
-            run: { responsibleUserId: "A" },
-            context: { responsibleUserId: "A" },
+      try {
+        const [settings] = await identityDb.execute(sql`show lock_timeout`);
+        expect(settings?.lock_timeout).toBe("1s");
+        await db.transaction(async (audit) => {
+          await audit.insert(secretAccessEvents).values({
+            companyId: input.companyId,
+            heartbeatRunId: input.runId,
+            issueId: input.issueId,
+            provider: "local_encrypted",
+            actorType: "agent",
+            actorId: input.agentId,
+            consumerType: "agent",
+            consumerId: input.agentId,
+            outcome: "granted",
           });
-        }
-      });
+          if (operation === "initialize") {
+            await expect(initializeRunIdentity(identityDb, {
+              ...input, messageIds: [], responsibleUserId: "A", cause: "instruction",
+            })).resolves.toMatchObject({ responsibleUserId: "A" });
+          } else {
+            await expect(captureRunIdentity(identityDb, input)).resolves.toMatchObject({
+              run: { responsibleUserId: "A" },
+              context: { responsibleUserId: "A" },
+            });
+          }
+        });
+      } finally {
+        await identityDb.$client.end({ timeout: 1 });
+      }
     },
   );
 

server/src/services/run-identity.ts

@@ -10,7 +10,8 @@ import {
   type Db,
 } from "@paperclipai/db";
 import { conflict, forbidden } from "../errors.js";
-import { boundRunIssueId } from "./run-task-binding.js";
+import { notifyChatRunMilestones } from "./chat-run-milestone-notifications.js";
+import { boundRunIssueId, boundRunIssueIdSql } from "./run-task-binding.js";
 import { isUuidLike } from "@paperclipai/shared";
 import { queuedCommentIdsFromRunContext, queuedCommentIdsFromWakePayload } from "./issue-queued-comment-queue.js";
 
@@ -137,6 +138,13 @@ async function append(
 }
 
 async function activate(executor: Executor, context: RunIdentityContext) {
+  const [binding] = await executor.select({ issueId: boundRunIssueIdSql() })
+    .from(heartbeatRuns).where(and(
+      eq(heartbeatRuns.id, context.runId), eq(heartbeatRuns.companyId, context.companyId),
+    ));
+  if (binding?.issueId) {
+    await notifyChatRunMilestones(executor, { companyId: context.companyId, issueId: binding.issueId });
+  }
   const [run] = await executor
     .update(heartbeatRuns)
     .set({
@@ -409,6 +417,7 @@ export async function rejectSteeredIdentity(
     );
 }
 
+/** Accept within the caller's task/run transaction; standalone recovery uses reconcileSteeredIdentity. */
 export async function acceptSteeredIdentity(
   executor: Executor,
   context: RunIdentityContext,