diff --git a/scripts/check-module-boundaries.mjs b/scripts/check-module-boundaries.mjs index 483a7086f0..eed66df002 100644 --- a/scripts/check-module-boundaries.mjs +++ b/scripts/check-module-boundaries.mjs @@ -126,6 +126,8 @@ export function scanModuleBoundaries({ addViolation(violations, sourceLabel, "application", specifier, "application cannot import database packages"); } else if (targetLocation?.layer === "adapters" || targetSegments.includes("adapters")) { addViolation(violations, sourceLabel, "application", specifier, "application cannot import concrete adapters"); + } else if (targetSegments.includes("services") || targetSegments.includes("routes")) { + addViolation(violations, sourceLabel, "application", specifier, "application cannot import server services or routes"); } else if (targetSegments.join("/") === "errors.js" || targetSegments.join("/") === "errors.ts") { addViolation(violations, sourceLabel, "application", specifier, "application cannot import HTTP error helpers"); } diff --git a/scripts/check-module-boundaries.test.mjs b/scripts/check-module-boundaries.test.mjs index a73f87e855..671025cec3 100644 --- a/scripts/check-module-boundaries.test.mjs +++ b/scripts/check-module-boundaries.test.mjs @@ -41,6 +41,7 @@ test("scanModuleBoundaries rejects outward dependencies and module-internal impo 'import { adapter } from "../adapters/postgres.js";', 'import { parse } from "../../../adapters/application-utils.js";', 'import { forbidden } from "../../../errors.js";', + 'import { helper } from "../../../services/example.js";', 'import db = require("@paperclipai/db");', ].join("\n"), ); @@ -58,6 +59,10 @@ test("scanModuleBoundaries rejects outward dependencies and module-internal impo reason: "application cannot import concrete adapters", }, { specifier: "../../../errors.js", reason: "application cannot import HTTP error helpers" }, + { + specifier: "../../../services/example.js", + reason: "application cannot import server services or routes", + }, { specifier: "@paperclipai/db", reason: "application cannot import database packages" }, { specifier: "drizzle-orm", reason: "domain cannot import database packages" }, { diff --git a/server/src/modules/wake-queue/adapters/postgres.test.ts b/server/src/modules/wake-queue/adapters/postgres.test.ts index 0fbc086385..c7a787bca2 100644 --- a/server/src/modules/wake-queue/adapters/postgres.test.ts +++ b/server/src/modules/wake-queue/adapters/postgres.test.ts @@ -1,6 +1,7 @@ import { randomUUID } from "node:crypto"; import { eq } from "drizzle-orm"; import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; +import type { Db } from "@paperclipai/db"; import { agentWakeupRequests, agents, @@ -14,8 +15,13 @@ import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase, } from "../../../__tests__/helpers/embedded-postgres.js"; -import { createPostgresWakeQueueAdapter } from "./postgres.js"; +import { + createAdmissionTransactionScope, + createPostgresWakeQueueAdapter, + createWakeAdmissionWriter, +} from "./postgres.js"; import type { WakeQueuePostgresAdapterDeps } from "./postgres.js"; +import type { TransactionScope } from "../application/ports.js"; // Proves the atomicity and company-scope properties the security review // requires: every mutation names `companyId` in its own SQL `WHERE` clause, @@ -436,4 +442,105 @@ describeEmbeddedPostgres("wake-queue postgres adapter", () => { expect(issueRow?.executionRunId).toBe(retryRunId); expect(issueRow?.checkoutRunId).toBeNull(); }); + + // The admission half opens no transaction of its own: `heartbeat.ts` still + // owns it. These tests drive the admission writer directly against a + // transaction they open themselves, the same way `heartbeat.ts` will. + describe("wake admission", () => { + // Review test (b): an admission adapter mutation with a foreign company + // affects no row. + it("refuses to merge into a deferred wake for a company that does not own it, and leaves the wake untouched", async () => { + const companyId = await seedCompany(); + const otherCompanyId = await seedCompany(); + const agentId = await seedAgent({ companyId }); + const issueId = await seedIssue({ companyId, assigneeAgentId: agentId }); + const wakeId = await seedDeferredWake({ companyId, agentId, issueId }); + const writer = createWakeAdmissionWriter(); + + await expect( + db.transaction(async (tx) => { + const scope = createAdmissionTransactionScope(otherCompanyId, tx as unknown as Db); + await writer.mergeIntoExistingDeferredWake(scope, { + companyId: otherCompanyId, + existingDeferredWakeId: wakeId, + mergedPayload: { issueId, foo: "bar" }, + nextCoalescedCount: 5, + }); + }), + ).rejects.toThrow(); + + const wakeRow = (await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wakeId)))[0]; + expect(wakeRow?.status).toBe("deferred_issue_execution"); + expect(wakeRow?.coalescedCount).not.toBe(5); + expect(wakeRow?.payload).not.toHaveProperty("foo"); + }); + + // Review test (b), the coalesce write: a foreign company affects no row + // on the active execution run either. + it("refuses to coalesce into a run for a company that does not own it, and leaves the run and the wake table untouched", async () => { + const companyId = await seedCompany(); + const otherCompanyId = await seedCompany(); + const agentId = await seedAgent({ companyId }); + const runId = await seedRun({ companyId, agentId, status: "running", contextSnapshot: { taskKey: "issue-1" } }); + const writer = createWakeAdmissionWriter(); + + await expect( + db.transaction(async (tx) => { + const scope = createAdmissionTransactionScope(otherCompanyId, tx as unknown as Db); + await writer.coalesceIntoActiveExecutionRun(scope, { + companyId: otherCompanyId, + activeExecutionRunId: runId, + mergedContextSnapshot: { taskKey: "issue-1", commentId: "c1" }, + agentId, + source: "on_demand", + triggerDetail: null, + payload: null, + requestedByActorType: "user", + requestedByActorId: null, + idempotencyKey: null, + }); + }), + ).rejects.toThrow(); + + const runRow = (await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)))[0]; + expect(runRow?.contextSnapshot).toEqual({ taskKey: "issue-1" }); + const wakeRows = await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.companyId, otherCompanyId)); + expect(wakeRows).toHaveLength(0); + }); + + // Review test (d): a transaction-port call cannot use the root database + // executor. The module rejects a missing scope handle and a mismatched + // scope handle, and neither case writes a row. + it("rejects a missing transaction scope and a mismatched one, without writing through the root executor", async () => { + const companyId = await seedCompany(); + const otherCompanyId = await seedCompany(); + const agentId = await seedAgent({ companyId }); + const issueId = await seedIssue({ companyId, assigneeAgentId: agentId }); + const writer = createWakeAdmissionWriter(); + const newWakeInput = { + companyId, + agentId, + source: "automation", + triggerDetail: null, + payload: { issueId }, + requestedByActorType: "system", + requestedByActorId: null, + idempotencyKey: null, + }; + + await expect( + writer.insertNewDeferredWake(undefined as unknown as TransactionScope, newWakeInput), + ).rejects.toThrow(/transaction scope/); + + await expect( + db.transaction(async (tx) => { + const mismatchedScope = createAdmissionTransactionScope(otherCompanyId, tx as unknown as Db); + await writer.insertNewDeferredWake(mismatchedScope, newWakeInput); + }), + ).rejects.toThrow(/different company/); + + const rows = await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.companyId, companyId)); + expect(rows).toHaveLength(0); + }); + }); }); diff --git a/server/src/modules/wake-queue/adapters/postgres.ts b/server/src/modules/wake-queue/adapters/postgres.ts index fd9e434535..50b5b4e048 100644 --- a/server/src/modules/wake-queue/adapters/postgres.ts +++ b/server/src/modules/wake-queue/adapters/postgres.ts @@ -33,6 +33,7 @@ import { parseObject, readNonEmptyString, } from "../domain/values.js"; +import { requireTransactionScopeTx, TransactionScope } from "../application/ports.js"; import type { DeferredWakeCandidate, InvokableAgentSnapshot, @@ -41,6 +42,8 @@ import type { LockedIssueExecution, ReleaseTransactionResult, RunSnapshot, + WakeAdmissionReader, + WakeAdmissionWriter, WakeQueueHost, WakeQueueTransaction, } from "../application/ports.js"; @@ -628,6 +631,142 @@ async function recordNativeTerminalRecoveryIfNeeded(tx: Db, run: HeartbeatRunRow return true; } +/** + * Builds the temporary transaction-scope handle the admission port needs. + * `heartbeat.ts` calls this through `createWakeQueue`'s own wrapper; it + * never builds a `TransactionScope` itself. + */ +export function createAdmissionTransactionScope(companyId: string, tx: Db): TransactionScope { + return TransactionScope.create(companyId, tx); +} + +function requireAdmissionTx(scope: TransactionScope | null | undefined, companyId: string): Db { + return requireTransactionScopeTx(scope, companyId) as Db; +} + +export function createWakeAdmissionReader(): WakeAdmissionReader { + return { + async isSameExecutionAgent( + scope, + { companyId, activeExecutionRunAgentId, issueExecutionAgentNameKey, agentNameKey }, + ) { + const tx = requireAdmissionTx(scope, companyId); + const executionAgent = await tx + .select({ name: agents.name }) + .from(agents) + .where(and(eq(agents.id, activeExecutionRunAgentId), eq(agents.companyId, companyId))) + .then((rows) => rows[0] ?? null); + const executionAgentNameKey = + normalizeAgentNameKey(issueExecutionAgentNameKey) ?? normalizeAgentNameKey(executionAgent?.name); + return Boolean(executionAgentNameKey) && executionAgentNameKey === agentNameKey; + }, + + async findExistingDeferredWake(scope, { companyId, agentId, issueId }) { + const tx = requireAdmissionTx(scope, companyId); + const row = await tx + .select() + .from(agentWakeupRequests) + .where( + and( + eq(agentWakeupRequests.companyId, companyId), + eq(agentWakeupRequests.agentId, agentId), + eq(agentWakeupRequests.status, DEFERRED_WAKE_STATUS), + sql`${agentWakeupRequests.payload} ->> 'issueId' = ${issueId}`, + ), + ) + .orderBy(asc(agentWakeupRequests.requestedAt)) + .limit(1) + .then((rows) => rows[0] ?? null); + if (!row) return null; + const payload = parseObject(row.payload); + return { + id: row.id, + payload, + deferredContext: parseObject(payload[DEFERRED_WAKE_CONTEXT_KEY]), + coalescedCount: row.coalescedCount, + }; + }, + }; +} + +export function createWakeAdmissionWriter(): WakeAdmissionWriter { + return { + async coalesceIntoActiveExecutionRun(scope, input) { + const tx = requireAdmissionTx(scope, input.companyId); + const now = new Date(); + const mergedRun = await tx + .update(heartbeatRuns) + .set({ contextSnapshot: input.mergedContextSnapshot, updatedAt: now }) + .where(and(eq(heartbeatRuns.id, input.activeExecutionRunId), eq(heartbeatRuns.companyId, input.companyId))) + .returning() + .then((rows) => rows[0] ?? null); + if (!mergedRun) { + // The compare-and-set write affected no row. Throw to roll the + // transaction back instead of recording a coalesced wake against a + // run this write never touched. + throw new Error("wake-queue: the coalesce target run was not found for this company"); + } + await tx.insert(agentWakeupRequests).values({ + companyId: input.companyId, + agentId: input.agentId, + source: input.source, + triggerDetail: input.triggerDetail, + reason: "issue_execution_same_name", + payload: input.payload, + status: "coalesced", + coalescedCount: 1, + requestedByActorType: input.requestedByActorType, + requestedByActorId: input.requestedByActorId, + idempotencyKey: input.idempotencyKey, + runId: mergedRun.id, + finishedAt: now, + }); + return mergedRun as unknown as Record; + }, + + async mergeIntoExistingDeferredWake(scope, input) { + const tx = requireAdmissionTx(scope, input.companyId); + const rows = await tx + .update(agentWakeupRequests) + .set({ + payload: input.mergedPayload, + coalescedCount: input.nextCoalescedCount, + updatedAt: new Date(), + }) + .where( + and( + eq(agentWakeupRequests.id, input.existingDeferredWakeId), + eq(agentWakeupRequests.companyId, input.companyId), + eq(agentWakeupRequests.status, DEFERRED_WAKE_STATUS), + ), + ) + .returning({ id: agentWakeupRequests.id }); + if (rows.length === 0) { + // The compare-and-set write affected no row: a concurrent writer + // already moved this wake off `deferred_issue_execution`. Roll the + // transaction back instead of leaving the merge half-applied. + throw new Error("wake-queue: the deferred wake to merge into was not found for this company"); + } + }, + + async insertNewDeferredWake(scope, input) { + const tx = requireAdmissionTx(scope, input.companyId); + await tx.insert(agentWakeupRequests).values({ + companyId: input.companyId, + agentId: input.agentId, + source: input.source, + triggerDetail: input.triggerDetail, + reason: "issue_execution_deferred", + payload: input.payload, + status: DEFERRED_WAKE_STATUS, + requestedByActorType: input.requestedByActorType, + requestedByActorId: input.requestedByActorId, + idempotencyKey: input.idempotencyKey, + }); + }, + }; +} + export function createPostgresWakeQueueAdapter(db: Db, deps: WakeQueuePostgresAdapterDeps): IssueLockWriter { return { async withIssueExecutionLock(input, fn): Promise { diff --git a/server/src/modules/wake-queue/application/ports.ts b/server/src/modules/wake-queue/application/ports.ts index 24c82d5624..ec45559628 100644 --- a/server/src/modules/wake-queue/application/ports.ts +++ b/server/src/modules/wake-queue/application/ports.ts @@ -244,3 +244,160 @@ export interface RecoveryEscalationPort { } export type { PostCommitEffect, ReleaseOutcome }; + +/** + * Temporary port: `heartbeat.ts` still opens and owns the transaction that + * admits a wake behind an active issue execution; the module does not own + * that transaction yet. A later change will decompose `enqueueWakeup` so + * the module owns the transaction itself. That change removes this handle + * and replaces it with a transaction the module opens on its own. + * + * Only this module builds a scope. `createWakeQueue` exposes a method that + * builds one for a caller outside the module. That caller receives an + * opaque handle back and never touches this class directly. + */ +export class TransactionScope { + private constructor( + private readonly companyId: string, + private readonly rawTx: unknown, + ) {} + + static create(companyId: string, rawTx: unknown): TransactionScope { + return new TransactionScope(companyId, rawTx); + } + + /** Returns the bound transaction only when `companyId` matches the scope's own company. */ + requireTx(companyId: string): unknown { + if (companyId !== this.companyId) { + throw new Error( + "wake-queue: the transaction scope belongs to a different company than the requested write", + ); + } + return this.rawTx; + } +} + +/** Reads a scope's bound transaction. Rejects a missing scope with the same clear error as a mismatched one; never falls back to any other executor. */ +export function requireTransactionScopeTx( + scope: TransactionScope | null | undefined, + companyId: string, +): unknown { + if (!scope) { + throw new Error("wake-queue: this call carries no transaction scope"); + } + return scope.requireTx(companyId); +} + +/** The active execution run a new wake arrives behind. */ +export type WakeAdmissionActiveExecutionRun = { + id: string; + agentId: string; + status: string; + contextSnapshot: unknown; +}; + +export type ExistingDeferredWake = { + id: string; + payload: Record; + /** `payload._paperclipWakeContext`, already parsed to a plain object. */ + deferredContext: Record; + coalescedCount: number | null; +}; + +/** + * The four wake-admission decision helpers that stay in `heartbeat.ts` + * today. The application layer receives them through this port so it never + * imports the service it is extracted from. + */ +export type WakeAdmissionHeartbeatHelpers = { + /** `filterZombieCoalesceTarget` in `heartbeat.ts`. */ + filterZombieCoalesceTarget( + target: WakeAdmissionActiveExecutionRun | null, + liveRunExecutions: { has(id: string): boolean }, + ): WakeAdmissionActiveExecutionRun | null; + /** `mergeCoalescedContextSnapshot` in `heartbeat.ts`. */ + mergeCoalescedContextSnapshot( + existingRaw: unknown, + incoming: Record, + options?: { preserveExistingInteractionContinuation?: boolean }, + ): Record; + /** `shouldDeferFollowupWakeForSameIssue` in `heartbeat.ts`. */ + shouldDeferFollowupWakeForSameIssue(input: { + activeRunStatus: string | null | undefined; + isSameExecutionAgent: boolean; + wakeCommentId: string | null | undefined; + forceFreshSession: boolean; + }): boolean; + /** `shouldQueueFollowupForRunningIssueWake` in `heartbeat.ts`. */ + shouldQueueFollowupForRunningIssueWake(input: { + contextSnapshot: Record | null | undefined; + wakeCommentId: string | null; + }): boolean; +}; + +export type AdmitWakeBehindIssueExecutionResult = + | { kind: "proceed" } + | { kind: "coalesced"; run: Record } + | { kind: "deferred" }; + +/** Read-only lookups the admission use case needs, each scoped to a company. */ +export interface WakeAdmissionReader { + /** True when the active execution run's agent and this wake's own agent share an execution-agent-name key. */ + isSameExecutionAgent( + scope: TransactionScope, + input: { + companyId: string; + activeExecutionRunAgentId: string; + issueExecutionAgentNameKey: string | null; + agentNameKey: string | null; + }, + ): Promise; + findExistingDeferredWake( + scope: TransactionScope, + input: { companyId: string; agentId: string; issueId: string }, + ): Promise; +} + +/** The transaction-scoped write operations that admit a wake behind an active issue execution. */ +export interface WakeAdmissionWriter { + /** Merges the wake's context into the active execution run and records the wake as coalesced. Returns the updated run row. */ + coalesceIntoActiveExecutionRun( + scope: TransactionScope, + input: { + companyId: string; + activeExecutionRunId: string; + mergedContextSnapshot: Record; + agentId: string; + source: string; + triggerDetail: string | null; + payload: Record | null; + requestedByActorType: string | null; + requestedByActorId: string | null; + idempotencyKey: string | null; + }, + ): Promise>; + /** Merges the wake's context into an already-queued deferred wake, guarded by its current status. */ + mergeIntoExistingDeferredWake( + scope: TransactionScope, + input: { + companyId: string; + existingDeferredWakeId: string; + mergedPayload: Record; + nextCoalescedCount: number; + }, + ): Promise; + /** Queues a new deferred wake behind the active execution run. */ + insertNewDeferredWake( + scope: TransactionScope, + input: { + companyId: string; + agentId: string; + source: string; + triggerDetail: string | null; + payload: Record; + requestedByActorType: string | null; + requestedByActorId: string | null; + idempotencyKey: string | null; + }, + ): Promise; +} diff --git a/server/src/modules/wake-queue/application/use-cases.test.ts b/server/src/modules/wake-queue/application/use-cases.test.ts index c6a7a7f677..ddba08ad09 100644 --- a/server/src/modules/wake-queue/application/use-cases.test.ts +++ b/server/src/modules/wake-queue/application/use-cases.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from "vitest"; -import { createReleaseIssueExecution } from "./use-cases.js"; +import { createAdmitWakeBehindIssueExecution, createReleaseIssueExecution } from "./use-cases.js"; +import type { AdmitWakeBehindIssueExecutionInput } from "./use-cases.js"; import { WakeQueueApplicationError } from "./types.js"; import type { DeferredWakeCandidate, @@ -10,6 +11,11 @@ import type { RecoveryEscalationPort, RunSnapshot, RunSummary, + TransactionScope, + WakeAdmissionActiveExecutionRun, + WakeAdmissionHeartbeatHelpers, + WakeAdmissionReader, + WakeAdmissionWriter, WakeQueueHost, WakeQueueTransaction, } from "./ports.js"; @@ -405,3 +411,161 @@ describe("releaseIssueExecution", () => { expect(recovery.escalateStrandedAssignedIssue).toHaveBeenCalledTimes(1); }); }); + +const ACTIVE_EXECUTION_RUN: WakeAdmissionActiveExecutionRun = { + id: "active-run-1", + agentId: "execution-agent", + status: "running", + contextSnapshot: { taskKey: "issue-1" }, +}; + +// The scope is opaque to the use case; the fakes below never inspect it. +const SCOPE = {} as TransactionScope; + +function admissionInput( + overrides: Partial = {}, +): AdmitWakeBehindIssueExecutionInput { + return { + companyId: "company-1", + issueId: "issue-1", + agentId: "wake-agent", + agentNameKey: "codexcoder", + issueExecutionAgentNameKey: null, + activeExecutionRun: ACTIVE_EXECUTION_RUN, + liveRunExecutions: { has: () => true }, + wakeCommentId: null, + forceFreshSession: false, + contextSnapshot: { wakeReason: "issue_commented" }, + source: "on_demand", + triggerDetail: null, + payload: { issueId: "issue-1" }, + requestedByActorType: "user", + requestedByActorId: "user-1", + idempotencyKey: null, + ...overrides, + }; +} + +function createFakeAdmissionReader(overrides: Partial = {}): WakeAdmissionReader { + return { + isSameExecutionAgent: vi.fn(async () => true), + findExistingDeferredWake: vi.fn(async () => null), + ...overrides, + }; +} + +function createFakeAdmissionWriter(overrides: Partial = {}): WakeAdmissionWriter { + return { + coalesceIntoActiveExecutionRun: vi.fn(async () => ({ id: "merged-run-1" })), + mergeIntoExistingDeferredWake: vi.fn(async () => {}), + insertNewDeferredWake: vi.fn(async () => {}), + ...overrides, + }; +} + +// Test doubles for the four heartbeat.ts decision helpers the module +// receives as a port. The defaults mirror the real helpers' behaviour for +// the plain wake in `admissionInput()`: no comment id, no forced fresh +// session, and a live coalesce target when `liveRunExecutions.has` says so. +function createFakeAdmissionHelpers( + overrides: Partial = {}, +): WakeAdmissionHeartbeatHelpers { + return { + filterZombieCoalesceTarget: vi.fn((target, liveRunExecutions) => + target && liveRunExecutions.has(target.id) ? target : null, + ), + mergeCoalescedContextSnapshot: vi.fn((existingRaw, incoming) => ({ + ...(existingRaw && typeof existingRaw === "object" ? (existingRaw as Record) : {}), + ...incoming, + })), + shouldDeferFollowupWakeForSameIssue: vi.fn(() => false), + shouldQueueFollowupForRunningIssueWake: vi.fn(() => false), + ...overrides, + }; +} + +describe("admitWakeBehindIssueExecution", () => { + it("returns the coalesce outcome and calls the writer one time when the same agent's run absorbs the wake", async () => { + const writer = createFakeAdmissionWriter(); + const reader = createFakeAdmissionReader(); + const helpers = createFakeAdmissionHelpers(); + const admit = createAdmitWakeBehindIssueExecution({ reader, writer, helpers }); + + const result = await admit(SCOPE, admissionInput()); + + expect(result).toEqual({ kind: "coalesced", run: { id: "merged-run-1" } }); + expect(writer.coalesceIntoActiveExecutionRun).toHaveBeenCalledTimes(1); + expect(writer.mergeIntoExistingDeferredWake).not.toHaveBeenCalled(); + expect(writer.insertNewDeferredWake).not.toHaveBeenCalled(); + }); + + it("never reads for an existing deferred wake on the coalesce path", async () => { + const writer = createFakeAdmissionWriter(); + const reader = createFakeAdmissionReader(); + const helpers = createFakeAdmissionHelpers(); + const admit = createAdmitWakeBehindIssueExecution({ reader, writer, helpers }); + + const result = await admit(SCOPE, admissionInput()); + + expect(result.kind).toBe("coalesced"); + expect(reader.findExistingDeferredWake).not.toHaveBeenCalled(); + }); + + it("merges into the existing deferred wake when the policy returns a merge target", async () => { + const mergeIntoExistingDeferredWake = vi.fn( + async (_scope: TransactionScope, _input: Parameters[1]) => {}, + ); + const writer = createFakeAdmissionWriter({ mergeIntoExistingDeferredWake }); + const reader = createFakeAdmissionReader({ + isSameExecutionAgent: vi.fn(async () => false), + findExistingDeferredWake: vi.fn(async () => ({ + id: "deferred-1", + payload: { issueId: "issue-1", foo: "bar" }, + deferredContext: { wakeReason: "issue_commented" }, + coalescedCount: 2, + })), + }); + const helpers = createFakeAdmissionHelpers(); + const admit = createAdmitWakeBehindIssueExecution({ reader, writer, helpers }); + + const result = await admit(SCOPE, admissionInput()); + + expect(result).toEqual({ kind: "deferred" }); + expect(mergeIntoExistingDeferredWake).toHaveBeenCalledTimes(1); + const call = mergeIntoExistingDeferredWake.mock.calls[0]![1]; + expect(call.existingDeferredWakeId).toBe("deferred-1"); + expect(call.nextCoalescedCount).toBe(3); + expect(call.mergedPayload.foo).toBe("bar"); + expect(writer.coalesceIntoActiveExecutionRun).not.toHaveBeenCalled(); + expect(writer.insertNewDeferredWake).not.toHaveBeenCalled(); + }); + + it("inserts a new deferred wake when a different agent holds the lock and none is queued yet", async () => { + const writer = createFakeAdmissionWriter(); + const reader = createFakeAdmissionReader({ isSameExecutionAgent: vi.fn(async () => false) }); + const helpers = createFakeAdmissionHelpers(); + const admit = createAdmitWakeBehindIssueExecution({ reader, writer, helpers }); + + const result = await admit(SCOPE, admissionInput()); + + expect(result).toEqual({ kind: "deferred" }); + expect(writer.insertNewDeferredWake).toHaveBeenCalledTimes(1); + expect(writer.coalesceIntoActiveExecutionRun).not.toHaveBeenCalled(); + expect(writer.mergeIntoExistingDeferredWake).not.toHaveBeenCalled(); + }); + + it("proceeds, and never reads for an existing deferred wake, when the zombie-run filter leaves no live coalesce target", async () => { + const writer = createFakeAdmissionWriter(); + const reader = createFakeAdmissionReader(); + const helpers = createFakeAdmissionHelpers(); + const admit = createAdmitWakeBehindIssueExecution({ reader, writer, helpers }); + + const result = await admit(SCOPE, admissionInput({ liveRunExecutions: { has: () => false } })); + + expect(result).toEqual({ kind: "proceed" }); + expect(reader.findExistingDeferredWake).not.toHaveBeenCalled(); + expect(writer.coalesceIntoActiveExecutionRun).not.toHaveBeenCalled(); + expect(writer.mergeIntoExistingDeferredWake).not.toHaveBeenCalled(); + expect(writer.insertNewDeferredWake).not.toHaveBeenCalled(); + }); +}); diff --git a/server/src/modules/wake-queue/application/use-cases.ts b/server/src/modules/wake-queue/application/use-cases.ts index 355724391b..11a8748b33 100644 --- a/server/src/modules/wake-queue/application/use-cases.ts +++ b/server/src/modules/wake-queue/application/use-cases.ts @@ -2,6 +2,7 @@ import { enrichPromotedWakeContext } from "../domain/context.js"; import { decideQueuedCommentAction, decideReleaseRecovery, + decideWakeAdmission, decideWakeOutcome, deriveImmediateRecoveryContextLabels, } from "../domain/policy.js"; @@ -13,6 +14,7 @@ import { } from "../domain/values.js"; import { withRecoveryContext } from "../../../services/recovery/status-only-context.js"; import type { + AdmitWakeBehindIssueExecutionResult, DeferredWakeCandidate, InvokableAgentSnapshot, IssueLockWriter, @@ -21,12 +23,19 @@ import type { RecoveryEscalationPort, ReleaseTransactionResult, RunSnapshot, + TransactionScope, + WakeAdmissionActiveExecutionRun, + WakeAdmissionHeartbeatHelpers, + WakeAdmissionReader, + WakeAdmissionWriter, WakeQueueHost, WakeQueueTransaction, } from "./ports.js"; import type { PostCommitEffect, ReleaseOutcome } from "./types.js"; import { WakeQueueApplicationError } from "./types.js"; +const DEFERRED_WAKE_CONTEXT_KEY = "_paperclipWakeContext"; + const ISSUE_DISPOSITION_REPAIR_RETRY_REASON = "issue_disposition_repair"; const EXECUTION_REVIEW_PARTICIPANT_RECOVERY_WAKE_REASONS = new Set([ "execution_review_requested", @@ -553,6 +562,144 @@ function statusForBlock(issue: IssueSnapshot): "todo" | "in_progress" | "in_revi return issue.status === "todo" || issue.status === "in_review" ? issue.status : "in_progress"; } +export type AdmitWakeBehindIssueExecutionInput = { + companyId: string; + issueId: string; + agentId: string; + agentNameKey: string | null; + issueExecutionAgentNameKey: string | null; + activeExecutionRun: WakeAdmissionActiveExecutionRun; + /** Tracks which runs are still live in this process, for the zombie-run filter. */ + liveRunExecutions: { has(id: string): boolean }; + wakeCommentId: string | null; + forceFreshSession: boolean; + contextSnapshot: Record; + source: string; + triggerDetail: string | null; + payload: Record | null; + requestedByActorType: string | null; + requestedByActorId: string | null; + idempotencyKey: string | null; +}; + +export type { AdmitWakeBehindIssueExecutionResult }; + +/** + * Decides and applies the admission outcome for a wake that arrives while + * an active execution run already holds the issue's execution lock: merge + * it into that run (coalesce), hold it behind the run (defer), or leave it + * for the caller to queue as an ordinary wake (proceed). + */ +export function createAdmitWakeBehindIssueExecution(deps: { + reader: WakeAdmissionReader; + writer: WakeAdmissionWriter; + helpers: WakeAdmissionHeartbeatHelpers; +}) { + return async function admitWakeBehindIssueExecution( + scope: TransactionScope, + input: AdmitWakeBehindIssueExecutionInput, + ): Promise { + const isSameExecutionAgent = await deps.reader.isSameExecutionAgent(scope, { + companyId: input.companyId, + activeExecutionRunAgentId: input.activeExecutionRun.agentId, + issueExecutionAgentNameKey: input.issueExecutionAgentNameKey, + agentNameKey: input.agentNameKey, + }); + + const shouldDeferFollowupWake = deps.helpers.shouldDeferFollowupWakeForSameIssue({ + activeRunStatus: input.activeExecutionRun.status, + isSameExecutionAgent, + wakeCommentId: input.wakeCommentId, + forceFreshSession: input.forceFreshSession, + }); + const shouldQueueFollowupForRunningWake = + deps.helpers.shouldQueueFollowupForRunningIssueWake({ + contextSnapshot: input.contextSnapshot, + wakeCommentId: input.wakeCommentId, + }) && + input.activeExecutionRun.status === "running" && + isSameExecutionAgent; + const availableActiveExecutionRun = isSameExecutionAgent + ? deps.helpers.filterZombieCoalesceTarget(input.activeExecutionRun, input.liveRunExecutions) + : input.activeExecutionRun; + + const decision = decideWakeAdmission({ + isSameExecutionAgent, + shouldDeferFollowupWake, + shouldQueueFollowupForRunningWake, + availableActiveExecutionRunPresent: availableActiveExecutionRun !== null, + }); + + if (decision.kind === "proceed") return { kind: "proceed" }; + + if (decision.kind === "coalesce") { + const target = availableActiveExecutionRun!; + const mergedContextSnapshot = deps.helpers.mergeCoalescedContextSnapshot(target.contextSnapshot, input.contextSnapshot, { + preserveExistingInteractionContinuation: + target.status === "queued" || target.status === "scheduled_retry", + }); + const run = await deps.writer.coalesceIntoActiveExecutionRun(scope, { + companyId: input.companyId, + activeExecutionRunId: target.id, + mergedContextSnapshot, + agentId: input.agentId, + source: input.source, + triggerDetail: input.triggerDetail, + payload: input.payload, + requestedByActorType: input.requestedByActorType, + requestedByActorId: input.requestedByActorId, + idempotencyKey: input.idempotencyKey, + }); + return { kind: "coalesced", run }; + } + + // decision.kind === "defer": only now does the module read for an + // existing deferred wake, so the coalesce path (the common path) never + // pays for this query. + const existingDeferred = await deps.reader.findExistingDeferredWake(scope, { + companyId: input.companyId, + agentId: input.agentId, + issueId: input.issueId, + }); + + if (existingDeferred) { + const mergedDeferredContext = deps.helpers.mergeCoalescedContextSnapshot(existingDeferred.deferredContext, input.contextSnapshot, { + preserveExistingInteractionContinuation: true, + }); + const mergedPayload = { + ...existingDeferred.payload, + ...(input.payload ?? {}), + issueId: input.issueId, + [DEFERRED_WAKE_CONTEXT_KEY]: mergedDeferredContext, + }; + await deps.writer.mergeIntoExistingDeferredWake(scope, { + companyId: input.companyId, + existingDeferredWakeId: existingDeferred.id, + mergedPayload, + nextCoalescedCount: (existingDeferred.coalescedCount ?? 0) + 1, + }); + return { kind: "deferred" }; + } + + const deferredPayload = { + ...(input.payload ?? {}), + issueId: input.issueId, + [DEFERRED_WAKE_CONTEXT_KEY]: input.contextSnapshot, + }; + await deps.writer.insertNewDeferredWake(scope, { + companyId: input.companyId, + agentId: input.agentId, + source: input.source, + triggerDetail: input.triggerDetail, + payload: deferredPayload, + requestedByActorType: input.requestedByActorType, + requestedByActorId: input.requestedByActorId, + idempotencyKey: input.idempotencyKey, + }); + return { kind: "deferred" }; + }; +} + export function createReleaseIssueExecution(deps: { issueLock: IssueLockWriter; recovery: RecoveryEscalationPort; diff --git a/server/src/modules/wake-queue/domain/policy.test.ts b/server/src/modules/wake-queue/domain/policy.test.ts index b6c4a92a92..12119c0138 100644 --- a/server/src/modules/wake-queue/domain/policy.test.ts +++ b/server/src/modules/wake-queue/domain/policy.test.ts @@ -3,6 +3,7 @@ import { decidePreDrain, decideQueuedCommentAction, decideReleaseRecovery, + decideWakeAdmission, decideWakeOutcome, deriveImmediateRecoveryContextLabels, type DeferredWakeOutcomeFacts, @@ -10,6 +11,7 @@ import { type ImmediateRecoveryContextLabels, type PreDrainFacts, type ReleaseRecoveryFacts, + type WakeAdmissionFacts, } from "./policy.js"; const basePreDrainFacts: PreDrainFacts = { @@ -495,3 +497,50 @@ describe("deriveImmediateRecoveryContextLabels", () => { }); } }); + +const baseWakeAdmissionFacts: WakeAdmissionFacts = { + isSameExecutionAgent: true, + shouldDeferFollowupWake: false, + shouldQueueFollowupForRunningWake: false, + availableActiveExecutionRunPresent: true, +}; + +describe("decideWakeAdmission", () => { + const cases: Array<{ + name: string; + facts: WakeAdmissionFacts; + expected: ReturnType; + }> = [ + { + name: "coalesce: same execution agent, no defer condition, and a live coalesce target", + facts: baseWakeAdmissionFacts, + expected: { kind: "coalesce" }, + }, + { + name: "defer: same execution agent, but the running agent needs a fresh session", + facts: { ...baseWakeAdmissionFacts, shouldDeferFollowupWake: true }, + expected: { kind: "defer" }, + }, + { + name: "defer: same execution agent, but the running turn must finish first", + facts: { ...baseWakeAdmissionFacts, shouldQueueFollowupForRunningWake: true }, + expected: { kind: "defer" }, + }, + { + name: "defer: a different agent already holds the execution lock", + facts: { ...baseWakeAdmissionFacts, isSameExecutionAgent: false }, + expected: { kind: "defer" }, + }, + { + name: "proceed: the zombie-run filter leaves no live coalesce target", + facts: { ...baseWakeAdmissionFacts, availableActiveExecutionRunPresent: false }, + expected: { kind: "proceed" }, + }, + ]; + + for (const testCase of cases) { + it(testCase.name, () => { + expect(decideWakeAdmission(testCase.facts)).toEqual(testCase.expected); + }); + } +}); diff --git a/server/src/modules/wake-queue/domain/policy.ts b/server/src/modules/wake-queue/domain/policy.ts index 88feb2a7d2..0776721a9b 100644 --- a/server/src/modules/wake-queue/domain/policy.ts +++ b/server/src/modules/wake-queue/domain/policy.ts @@ -1,5 +1,7 @@ -// Pure decision rules for the release half of the deferred issue-execution -// wake state machine: +// Pure decision rules for the deferred issue-execution wake state machine: +// - the admission decision (decideWakeAdmission), applied to a new wake +// that arrives while an active execution run already holds the issue's +// execution lock // - the pre-drain decision (decidePreDrain), applied once per lock // acquisition, before the caller touches the deferred-wake queue at all // - the per-wake decision, split into decideQueuedCommentAction and @@ -11,6 +13,47 @@ // This file only branches on that facts object; it never queries a // database, reads the clock, or reads the wake context payload directly. +export type WakeAdmissionFacts = { + /** True when the active execution run's agent and this wake's own agent share an execution-agent-name key. */ + isSameExecutionAgent: boolean; + /** True when a same-agent, still-running execution run must not absorb this wake and needs a new run boundary instead. */ + shouldDeferFollowupWake: boolean; + /** True when a same-agent, running execution run must finish its current turn before this wake runs. */ + shouldQueueFollowupForRunningWake: boolean; + /** True when the active execution run still stands as a live coalesce target after the zombie-run filter runs. */ + availableActiveExecutionRunPresent: boolean; +}; + +export type WakeAdmissionDecision = + | { kind: "proceed" } + | { kind: "coalesce" } + | { kind: "defer" }; + +/** + * Decides what a new wake does when an active execution run already holds + * the issue's execution lock: run into that run (coalesce), wait behind it + * (defer), or proceed as an ordinary wake because no run currently holds the + * lock. The caller reads whether a deferred wake already exists only after + * this function returns `defer`, so that read never runs on the coalesce + * path. + */ +export function decideWakeAdmission(facts: WakeAdmissionFacts): WakeAdmissionDecision { + if ( + facts.isSameExecutionAgent && + !facts.shouldDeferFollowupWake && + !facts.shouldQueueFollowupForRunningWake && + facts.availableActiveExecutionRunPresent + ) { + return { kind: "coalesce" }; + } + + if (facts.availableActiveExecutionRunPresent) { + return { kind: "defer" }; + } + + return { kind: "proceed" }; +} + export type PreDrainFacts = { /** True when the transaction found the issue row the lock is for. */ issueRowPresent: boolean; diff --git a/server/src/modules/wake-queue/index.ts b/server/src/modules/wake-queue/index.ts index 62e369a27b..57f02cca2c 100644 --- a/server/src/modules/wake-queue/index.ts +++ b/server/src/modules/wake-queue/index.ts @@ -1,10 +1,17 @@ import type { Db } from "@paperclipai/db"; -import { createPostgresWakeQueueAdapter } from "./adapters/postgres.js"; -import { createReleaseIssueExecution } from "./application/use-cases.js"; +import { + createAdmissionTransactionScope as buildAdmissionTransactionScope, + createPostgresWakeQueueAdapter, + createWakeAdmissionReader, + createWakeAdmissionWriter, +} from "./adapters/postgres.js"; +import { createAdmitWakeBehindIssueExecution, createReleaseIssueExecution } from "./application/use-cases.js"; import type { IssueSnapshot, RecoveryEscalationPort, RunSnapshot, + TransactionScope, + WakeAdmissionHeartbeatHelpers, WakeQueueHost, } from "./application/ports.js"; @@ -19,8 +26,9 @@ export type { RunSnapshot, RecoveryEscalationPort, ReleaseRecoveryBlockedNoticeKind, + TransactionScope, } from "./application/ports.js"; -export type { ReleaseIssueExecutionInput } from "./application/use-cases.js"; +export type { AdmitWakeBehindIssueExecutionInput, AdmitWakeBehindIssueExecutionResult, ReleaseIssueExecutionInput } from "./application/use-cases.js"; export type WakeQueueDeps = { /** Stays in `heartbeat.ts`; resolves the responsible user for a promoted or recovery run seed. */ @@ -29,6 +37,12 @@ export type WakeQueueDeps = { getRoutineEnv: WakeQueueHost["getRoutineEnv"]; /** Stays in `heartbeat.ts`; resolves the session-before display id for a wakeup. */ resolveSessionBeforeForWakeup: WakeQueueHost["resolveSessionBeforeForWakeup"]; + /** + * The four wake-admission decision helpers stay in `heartbeat.ts` today; + * the module receives them here so it never imports the service it is + * extracted from. + */ + wakeAdmissionHelpers: WakeAdmissionHeartbeatHelpers; /** `services/recovery`'s stranded-issue escalation, called only after the release transaction commits. */ recovery: RecoveryEscalationPort; }; @@ -39,6 +53,11 @@ export type WakeQueueDeps = { * only caller: it builds one instance per process next to * `createRunDispatch(db)` and delegates `releaseIssueExecutionAndPromote`'s * body to `releaseIssueExecution`. + * + * The admission half is temporary: `heartbeat.ts` still opens and owns the + * transaction that admits a wake behind an active issue execution, so it + * builds a `TransactionScope` through `createAdmissionTransactionScope` + * before it calls `admitWakeBehindIssueExecution`. */ export function createWakeQueue(db: Db, deps: WakeQueueDeps) { const issueLock = createPostgresWakeQueueAdapter(db, { @@ -49,6 +68,14 @@ export function createWakeQueue(db: Db, deps: WakeQueueDeps) { return { releaseIssueExecution: createReleaseIssueExecution({ issueLock, recovery: deps.recovery }), + admitWakeBehindIssueExecution: createAdmitWakeBehindIssueExecution({ + reader: createWakeAdmissionReader(), + writer: createWakeAdmissionWriter(), + helpers: deps.wakeAdmissionHelpers, + }), + createAdmissionTransactionScope(companyId: string, tx: Db): TransactionScope { + return buildAdmissionTransactionScope(companyId, tx); + }, }; } diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index f91ca1fb9f..850a095379 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -8550,6 +8550,15 @@ export function heartbeatService( if (!agent) return null; return resolveSessionBeforeForWakeup(agent, input.taskKey); }, + // These four helpers stay in this file today; the wake-queue module + // receives them here so it never imports this file, the service it is + // extracted from. + wakeAdmissionHelpers: { + filterZombieCoalesceTarget, + mergeCoalescedContextSnapshot, + shouldDeferFollowupWakeForSameIssue, + shouldQueueFollowupForRunningIssueWake, + }, recovery: { escalateStrandedAssignedIssue: async (input) => { const rows = await loadStrandedEscalationRows(input); @@ -23797,147 +23806,50 @@ export function heartbeatService( // its fresh-session contract into unrelated work or create a second // deferred wake that could later replay the same reconciliation. if (reconciledSourceRunId) return { kind: "deferred" as const }; - const executionAgent = await tx - .select({ name: agents.name }) - .from(agents) - .where(eq(agents.id, activeExecutionRun.agentId)) - .then((rows) => rows[0] ?? null); - const executionAgentNameKey = - normalizeAgentNameKey(issue.executionAgentNameKey) ?? - normalizeAgentNameKey(executionAgent?.name); - const isSameExecutionAgent = - Boolean(executionAgentNameKey) && - executionAgentNameKey === agentNameKey; - const shouldDeferFollowupWake = shouldDeferFollowupWakeForSameIssue({ - activeRunStatus: activeExecutionRun.status, - isSameExecutionAgent, - wakeCommentId, - forceFreshSession: - enrichedContextSnapshot.forceFreshSession === true, - }); - const shouldQueueFollowupForRunningWake = - shouldQueueFollowupForRunningIssueWake({ - contextSnapshot: enrichedContextSnapshot, - wakeCommentId, - }) && - activeExecutionRun.status === "running" && - isSameExecutionAgent; - const availableActiveExecutionRun = isSameExecutionAgent - ? filterZombieCoalesceTarget(activeExecutionRun, liveRunExecutions) - : activeExecutionRun; - if ( - isSameExecutionAgent && - !shouldDeferFollowupWake && - !shouldQueueFollowupForRunningWake && - availableActiveExecutionRun - ) { - const mergedContextSnapshot = mergeCoalescedContextSnapshot( - availableActiveExecutionRun.contextSnapshot, - enrichedContextSnapshot, - { - preserveExistingInteractionContinuation: - availableActiveExecutionRun.status === "queued" || - availableActiveExecutionRun.status === "scheduled_retry", + const admissionScope = wakeQueue.createAdmissionTransactionScope( + agent.companyId, + tx as unknown as Db, + ); + const admission = await wakeQueue.admitWakeBehindIssueExecution( + admissionScope, + { + companyId: agent.companyId, + issueId: issue.id, + agentId, + agentNameKey, + issueExecutionAgentNameKey: issue.executionAgentNameKey, + activeExecutionRun: { + id: activeExecutionRun.id, + agentId: activeExecutionRun.agentId, + status: activeExecutionRun.status, + contextSnapshot: activeExecutionRun.contextSnapshot, }, - ); - const mergedRun = await tx - .update(heartbeatRuns) - .set({ - contextSnapshot: mergedContextSnapshot, - updatedAt: new Date(), - }) - .where(eq(heartbeatRuns.id, availableActiveExecutionRun.id)) - .returning() - .then((rows) => rows[0] ?? availableActiveExecutionRun); - - await tx.insert(agentWakeupRequests).values({ - companyId: agent.companyId, - agentId, + liveRunExecutions, + wakeCommentId, + forceFreshSession: + enrichedContextSnapshot.forceFreshSession === true, + contextSnapshot: enrichedContextSnapshot, source, triggerDetail, - reason: "issue_execution_same_name", payload, - status: "coalesced", - coalescedCount: 1, requestedByActorType: opts.requestedByActorType ?? null, requestedByActorId: opts.requestedByActorId ?? null, idempotencyKey: opts.idempotencyKey ?? null, - runId: mergedRun.id, - finishedAt: new Date(), - }); + }, + ); - return { kind: "coalesced" as const, run: mergedRun }; - } - - if (availableActiveExecutionRun) { - const deferredPayload = { - ...(payload ?? {}), - issueId, - [DEFERRED_WAKE_CONTEXT_KEY]: enrichedContextSnapshot, + if (admission.kind === "coalesced") { + return { + kind: "coalesced" as const, + run: admission.run as typeof heartbeatRuns.$inferSelect, }; - - const existingDeferred = await tx - .select() - .from(agentWakeupRequests) - .where( - and( - eq(agentWakeupRequests.companyId, agent.companyId), - eq(agentWakeupRequests.agentId, agentId), - eq(agentWakeupRequests.status, "deferred_issue_execution"), - sql`${agentWakeupRequests.payload} ->> 'issueId' = ${issue.id}`, - ), - ) - .orderBy(asc(agentWakeupRequests.requestedAt)) - .limit(1) - .then((rows) => rows[0] ?? null); - - if (existingDeferred) { - const existingDeferredPayload = parseObject( - existingDeferred.payload, - ); - const existingDeferredContext = parseObject( - existingDeferredPayload[DEFERRED_WAKE_CONTEXT_KEY], - ); - const mergedDeferredContext = mergeCoalescedContextSnapshot( - existingDeferredContext, - enrichedContextSnapshot, - { preserveExistingInteractionContinuation: true }, - ); - const mergedDeferredPayload = { - ...existingDeferredPayload, - ...(payload ?? {}), - issueId, - [DEFERRED_WAKE_CONTEXT_KEY]: mergedDeferredContext, - }; - - await tx - .update(agentWakeupRequests) - .set({ - payload: mergedDeferredPayload, - coalescedCount: (existingDeferred.coalescedCount ?? 0) + 1, - updatedAt: new Date(), - }) - .where(eq(agentWakeupRequests.id, existingDeferred.id)); - - return { kind: "deferred" as const }; - } - - await tx.insert(agentWakeupRequests).values({ - companyId: agent.companyId, - agentId, - source, - triggerDetail, - reason: "issue_execution_deferred", - payload: deferredPayload, - status: "deferred_issue_execution", - requestedByActorType: opts.requestedByActorType ?? null, - requestedByActorId: opts.requestedByActorId ?? null, - idempotencyKey: opts.idempotencyKey ?? null, - }); - + } + if (admission.kind === "deferred") { return { kind: "deferred" as const }; } + // admission.kind === "proceed": no active run absorbed this wake, + // so fall through to the ordinary queue path below. } // PAP-13775: no live run holds the lock, so this wake would start a