From 7aac0d73bc371294a9f0e695471e9fa043f4c430 Mon Sep 17 00:00:00 2001 From: Dotta Date: Tue, 21 Apr 2026 15:45:10 -0500 Subject: [PATCH] fix(memory-jobs): address final review hardening --- .../0063_motionless_gertrude_yorkes.sql | 2 + .../db/src/migrations/meta/0063_snapshot.json | 52 +++++++- .../db/src/schema/memory_extraction_jobs.ts | 8 ++ .../src/__tests__/memory-job-capture.test.ts | 2 +- server/src/services/memory-job-capture.ts | 4 +- server/src/services/memory-job-store.ts | 114 ++++++++++-------- ui/src/pages/MemoryJobs.tsx | 6 +- 7 files changed, 131 insertions(+), 57 deletions(-) diff --git a/packages/db/src/migrations/0063_motionless_gertrude_yorkes.sql b/packages/db/src/migrations/0063_motionless_gertrude_yorkes.sql index b1dc110039b..5f1c6307e1f 100644 --- a/packages/db/src/migrations/0063_motionless_gertrude_yorkes.sql +++ b/packages/db/src/migrations/0063_motionless_gertrude_yorkes.sql @@ -104,7 +104,9 @@ CREATE INDEX IF NOT EXISTS "memory_extraction_jobs_company_operation_submitted_i CREATE INDEX IF NOT EXISTS "memory_extraction_jobs_company_issue_submitted_idx" ON "memory_extraction_jobs" USING btree ("company_id","source_issue_id","submitted_at");--> statement-breakpoint CREATE INDEX IF NOT EXISTS "memory_extraction_jobs_company_run_submitted_idx" ON "memory_extraction_jobs" USING btree ("company_id","source_heartbeat_run_id","submitted_at");--> statement-breakpoint CREATE INDEX IF NOT EXISTS "memory_extraction_jobs_company_status_lease_expires_idx" ON "memory_extraction_jobs" USING btree ("company_id","status","lease_expires_at");--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "memory_extraction_jobs_dispatcher_queued_idx" ON "memory_extraction_jobs" USING btree ("dispatcher_kind","submitted_at") WHERE "status" = 'queued';--> statement-breakpoint CREATE INDEX IF NOT EXISTS "memory_extraction_jobs_retry_of_job_idx" ON "memory_extraction_jobs" USING btree ("retry_of_job_id");--> statement-breakpoint +CREATE UNIQUE INDEX IF NOT EXISTS "memory_extraction_jobs_retry_attempt_uq" ON "memory_extraction_jobs" USING btree ("company_id","retry_of_job_id","attempt_number") WHERE "retry_of_job_id" is not null;--> statement-breakpoint CREATE INDEX IF NOT EXISTS "routine_runs_dispatch_fingerprint_idx" ON "routine_runs" USING btree ("routine_id","dispatch_fingerprint");--> statement-breakpoint CREATE UNIQUE INDEX IF NOT EXISTS "issues_open_routine_execution_uq" ON "issues" USING btree ("company_id","origin_kind","origin_id","origin_fingerprint") WHERE "issues"."origin_kind" = 'routine_execution' and "issues"."origin_id" is not null diff --git a/packages/db/src/migrations/meta/0063_snapshot.json b/packages/db/src/migrations/meta/0063_snapshot.json index f5c02d2bedc..d14f38e490a 100644 --- a/packages/db/src/migrations/meta/0063_snapshot.json +++ b/packages/db/src/migrations/meta/0063_snapshot.json @@ -10677,6 +10677,28 @@ "method": "btree", "with": {} }, + "memory_extraction_jobs_dispatcher_queued_idx": { + "name": "memory_extraction_jobs_dispatcher_queued_idx", + "columns": [ + { + "expression": "dispatcher_kind", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "submitted_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"memory_extraction_jobs\".\"status\" = 'queued'", + "concurrently": false, + "method": "btree", + "with": {} + }, "memory_extraction_jobs_retry_of_job_idx": { "name": "memory_extraction_jobs_retry_of_job_idx", "columns": [ @@ -10691,6 +10713,34 @@ "concurrently": false, "method": "btree", "with": {} + }, + "memory_extraction_jobs_retry_attempt_uq": { + "name": "memory_extraction_jobs_retry_attempt_uq", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "retry_of_job_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "attempt_number", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "\"memory_extraction_jobs\".\"retry_of_job_id\" is not null", + "concurrently": false, + "method": "btree", + "with": {} } }, "foreignKeys": { @@ -14557,4 +14607,4 @@ "schemas": {}, "tables": {} } -} \ No newline at end of file +} diff --git a/packages/db/src/schema/memory_extraction_jobs.ts b/packages/db/src/schema/memory_extraction_jobs.ts index 354cf7b3b81..96f22bcecd1 100644 --- a/packages/db/src/schema/memory_extraction_jobs.ts +++ b/packages/db/src/schema/memory_extraction_jobs.ts @@ -6,8 +6,10 @@ import { pgTable, text, timestamp, + uniqueIndex, uuid, } from "drizzle-orm/pg-core"; +import { sql } from "drizzle-orm"; import type { MemoryExtractionJobAttributionMode, MemoryExtractionJobDispatcherKind, @@ -95,6 +97,12 @@ export const memoryExtractionJobs = pgTable( table.status, table.leaseExpiresAt, ), + dispatcherQueuedIdx: index("memory_extraction_jobs_dispatcher_queued_idx") + .on(table.dispatcherKind, table.submittedAt) + .where(sql`${table.status} = 'queued'`), retryOfJobIdx: index("memory_extraction_jobs_retry_of_job_idx").on(table.retryOfJobId), + retryAttemptUq: uniqueIndex("memory_extraction_jobs_retry_attempt_uq") + .on(table.companyId, table.retryOfJobId, table.attemptNumber) + .where(sql`${table.retryOfJobId} is not null`), }), ); diff --git a/server/src/__tests__/memory-job-capture.test.ts b/server/src/__tests__/memory-job-capture.test.ts index 44588e2e723..628f397293a 100644 --- a/server/src/__tests__/memory-job-capture.test.ts +++ b/server/src/__tests__/memory-job-capture.test.ts @@ -224,7 +224,7 @@ describeEmbeddedPostgres("memory job capture", () => { costCents: 31, }); expect(completed?.resultJson).toMatchObject({ - captureState: "succeeded", + captureState: "recorded_pending_extraction", sourceRunId: runId, persistedUsageSummary: true, }); diff --git a/server/src/services/memory-job-capture.ts b/server/src/services/memory-job-capture.ts index b0cad9e172d..f06d4447b2e 100644 --- a/server/src/services/memory-job-capture.ts +++ b/server/src/services/memory-job-capture.ts @@ -317,11 +317,13 @@ export function createPostRunMemoryCaptureHandler(): MemoryJobHandler { }; } + // The extraction engine is a follow-up layer; this handler only records the + // post-run context that a later extractor will consume. return { outcome: "succeeded", resultSummary: job.resultSummary ?? `Post-run capture recorded for run ${job.sourceHeartbeatRunId}`, resultJson: mergeCaptureResultJson(job.resultJson ?? null, { - captureState: "succeeded", + captureState: "recorded_pending_extraction", sourceRunId: job.sourceHeartbeatRunId, persistedUsageSummary: job.usageJson != null, persistedSourceIssueId: job.sourceIssueId, diff --git a/server/src/services/memory-job-store.ts b/server/src/services/memory-job-store.ts index f3ad2c1dc2e..a3d972ed4d2 100644 --- a/server/src/services/memory-job-store.ts +++ b/server/src/services/memory-job-store.ts @@ -9,6 +9,7 @@ import { isNull, lte, or, + sql, } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { memoryExtractionJobs } from "@paperclipai/db"; @@ -681,64 +682,71 @@ export function memoryJobStore(db: Db) { jobId: string, opts: { now?: Date } = {}, ): Promise { - const sourceJob = await ensureMemoryJobExists(db, companyId, jobId); + return db.transaction(async (tx) => { + const txDb = tx as unknown as Db; + const sourceJob = await ensureMemoryJobExists(txDb, companyId, jobId); - if (!rerunEligibleForJob(sourceJob)) { - throw conflict(`Memory job ${sourceJob.id} cannot be rerun from status ${sourceJob.status}`); - } + if (!rerunEligibleForJob(sourceJob)) { + throw conflict(`Memory job ${sourceJob.id} cannot be rerun from status ${sourceJob.status}`); + } - const retryRootId = retryRootIdForJob(sourceJob); - const lineageRows = await db - .select({ attemptNumber: memoryExtractionJobs.attemptNumber }) - .from(memoryExtractionJobs) - .where( - and( - eq(memoryExtractionJobs.companyId, companyId), - or( - eq(memoryExtractionJobs.id, retryRootId), - eq(memoryExtractionJobs.retryOfJobId, retryRootId), - ), - ), + const retryRootId = retryRootIdForJob(sourceJob); + await tx.execute( + sql`select id from ${memoryExtractionJobs} where ${memoryExtractionJobs.companyId} = ${companyId} and ${memoryExtractionJobs.id} = ${retryRootId} for update`, ); - const nextAttemptNumber = lineageRows.reduce((maxAttempt, row) => { - return Math.max(maxAttempt, row.attemptNumber); - }, 0) + 1; + const lineageRows = await txDb + .select({ attemptNumber: memoryExtractionJobs.attemptNumber }) + .from(memoryExtractionJobs) + .where( + and( + eq(memoryExtractionJobs.companyId, companyId), + or( + eq(memoryExtractionJobs.id, retryRootId), + eq(memoryExtractionJobs.retryOfJobId, retryRootId), + ), + ), + ); - return db - .insert(memoryExtractionJobs) - .values({ - companyId: sourceJob.companyId, - bindingId: sourceJob.bindingId, - bindingKey: sourceJob.bindingKey, - operationType: sourceJob.operationType, - status: "queued", - sourceAgentId: sourceJob.sourceAgentId, - sourceIssueId: sourceJob.sourceIssueId, - sourceProjectId: sourceJob.sourceProjectId, - sourceGoalId: sourceJob.sourceGoalId, - sourceHeartbeatRunId: sourceJob.sourceHeartbeatRunId, - hookKind: sourceJob.hookKind, - submittedAt: opts.now, - attributionMode: sourceJob.attributionMode, - costCents: 0, - resultSummary: null, - errorCode: null, - error: null, - sourceKind: sourceJob.sourceKind, - sourceRefJson: sourceJob.sourceRefJson, - retryOfJobId: retryRootId, - attemptNumber: nextAttemptNumber, - dispatcherKind: sourceJob.dispatcherKind, - leaseExpiresAt: null, - providerJobId: null, - usageJson: null, - resultJson: null, - startedAt: null, - finishedAt: null, - }) - .returning() - .then((rows) => rows[0]); + const nextAttemptNumber = lineageRows.reduce((maxAttempt, row) => { + return Math.max(maxAttempt, row.attemptNumber); + }, 0) + 1; + + return txDb + .insert(memoryExtractionJobs) + .values({ + companyId: sourceJob.companyId, + bindingId: sourceJob.bindingId, + bindingKey: sourceJob.bindingKey, + operationType: sourceJob.operationType, + status: "queued", + sourceAgentId: sourceJob.sourceAgentId, + sourceIssueId: sourceJob.sourceIssueId, + sourceProjectId: sourceJob.sourceProjectId, + sourceGoalId: sourceJob.sourceGoalId, + sourceHeartbeatRunId: sourceJob.sourceHeartbeatRunId, + hookKind: sourceJob.hookKind, + submittedAt: opts.now, + attributionMode: sourceJob.attributionMode, + costCents: 0, + resultSummary: null, + errorCode: null, + error: null, + sourceKind: sourceJob.sourceKind, + sourceRefJson: sourceJob.sourceRefJson, + retryOfJobId: retryRootId, + attemptNumber: nextAttemptNumber, + dispatcherKind: sourceJob.dispatcherKind, + leaseExpiresAt: null, + providerJobId: null, + usageJson: null, + resultJson: null, + startedAt: null, + finishedAt: null, + }) + .returning() + .then((rows) => rows[0]); + }); } return { diff --git a/ui/src/pages/MemoryJobs.tsx b/ui/src/pages/MemoryJobs.tsx index 3560638cb19..c7aeef7eec8 100644 --- a/ui/src/pages/MemoryJobs.tsx +++ b/ui/src/pages/MemoryJobs.tsx @@ -508,7 +508,11 @@ export function MemoryJobs() {
- Showing latest {JOB_LIST_LIMIT} + + {appliedFilters.state === "attention" + ? `Showing latest ${JOB_LIST_LIMIT} per state` + : `Showing latest ${JOB_LIST_LIMIT}`} + {jobs.length} visible {appliedFilters.issueLabel ? Issue {appliedFilters.issueLabel} : null}