diff --git a/packages/data-schemas/misc/documentdb/audit.documentdb.spec.ts b/packages/data-schemas/misc/documentdb/audit.documentdb.spec.ts new file mode 100644 index 0000000000..4bd0a9c6d1 --- /dev/null +++ b/packages/data-schemas/misc/documentdb/audit.documentdb.spec.ts @@ -0,0 +1,433 @@ +import mongoose from 'mongoose'; +import { randomUUID } from 'crypto'; +import type { ConnectOptions } from 'mongoose'; +import { createAgentTriggerDeliveryMethods } from '~/methods/triggerDelivery'; +import { createMessageMethods } from '~/methods/message'; +import { createModels } from '~/models'; + +jest.mock('~/config/winston', () => ({ + error: jest.fn(), + warn: jest.fn(), + info: jest.fn(), + debug: jest.fn(), +})); + +/** + * Amazon DocumentDB live adjudication of the surface added between 2026-07-29 + * (when #14495 cleared the last known pipeline-form updates) and 2026-08-30. + * + * A static audit against AWS's supported-APIs tables predicts six rejections. + * Those tables are omission-based, so only a real cluster settles it. Every + * probe drives the PRODUCTION method rather than a re-implementation, so the + * suite cannot drift from the shapes the server actually emits. + * + * Query shapes are parsed before they are matched, so the no-fixture probes + * still adjudicate: an engine that accepts the shape returns an empty result, + * one that rejects it throws. + * + * Run (from packages/data-schemas, against a DEDICATED database): + * DOCUMENTDB_URI="mongodb://user:pass@127.0.0.1:27017/librechat_audit\ + * ?tls=true&retryWrites=false&authSource=admin&authMechanism=SCRAM-SHA-1&directConnection=true" \ + * DOCUMENTDB_TLS_CA_FILE="global-bundle.pem" \ + * DOCUMENTDB_TLS_ALLOW_INVALID_HOSTNAMES=true \ + * npx jest --config misc/documentdb/jest.documentdb.config.mjs audit + * + * Every URI parameter above is load-bearing when reaching a cluster through an + * SSH tunnel, and each was established against a real cluster: + * - authSource=admin the user lives in admin; a database in the path + * otherwise becomes the auth source and login fails + * - authMechanism=SCRAM-SHA-1 DocumentDB rejects SCRAM-SHA-256 ("Unsupported + * mechanism [ -301 ]") + * - directConnection=true replica-set discovery returns internal cluster + * hostnames that are unreachable through a tunnel + * - tlsAllowInvalidHostnames the tunnel endpoint never matches the cert + * Set DOCUMENTDB_STRICT=true to turn a rejected production shape into a failure. + */ +const DOCUMENTDB_URI = process.env.DOCUMENTDB_URI ?? ''; +const STRICT = process.env.DOCUMENTDB_STRICT === 'true'; +const describeLive = DOCUMENTDB_URI ? describe : describe.skip; + +const runId = randomUUID().slice(0, 8); +const ACCEPTED = 'accepted'; +const verdicts: Record = {}; + +function getDb() { + const db = mongoose.connection.db; + if (!db) { + throw new Error('MongoDB database handle not available'); + } + return db; +} + +/** Records whether the engine parsed the shape. Never throws: a rejection is + * the finding, not a harness failure. */ +async function probe(label: string, run: () => Promise): Promise { + try { + await run(); + verdicts[label] = ACCEPTED; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + verdicts[label] = `REJECTED - ${message.replace(/\s+/g, ' ').slice(0, 140)}`; + } + return verdicts[label]; +} + +function expectShape(label: string): void { + expect(verdicts[label]).toBeDefined(); + if (STRICT) { + expect(verdicts[label]).toBe(ACCEPTED); + } +} + +describeLive('Amazon DocumentDB - 2026-08-30 audit surface', () => { + let triggerMethods: ReturnType; + let messageMethods: ReturnType; + let probeCollection: string; + + const userId = new mongoose.Types.ObjectId(); + const conversationId = `audit-convo-${runId}`; + const messageId = `audit-message-${runId}`; + const deliveryKey = `audit-delivery-${runId}`; + const sourceId = `audit-source-${runId}`; + + beforeAll(async () => { + const options: ConnectOptions = { autoIndex: false, autoCreate: false }; + if (process.env.DOCUMENTDB_TLS_CA_FILE) { + options.tlsCAFile = process.env.DOCUMENTDB_TLS_CA_FILE; + } + if (process.env.DOCUMENTDB_TLS_ALLOW_INVALID_HOSTNAMES === 'true') { + options.tlsAllowInvalidHostnames = true; + } + await mongoose.connect(DOCUMENTDB_URI, options); + + Object.assign(mongoose.models, createModels(mongoose)); + triggerMethods = createAgentTriggerDeliveryMethods(mongoose); + messageMethods = createMessageMethods(mongoose); + + probeCollection = `audit_probe_${runId}`; + await getDb() + .collection(probeCollection) + .insertOne({ + probe: 1, + label: 'audit', + note: 'construct probe row', + values: [1, 2, 3], + }); + }); + + afterAll(async () => { + if (mongoose.connection.readyState !== 1) { + return; + } + await getDb() + .collection(probeCollection) + .drop() + .catch(() => undefined); + await mongoose.models.Message.deleteMany({ conversationId }).catch(() => undefined); + await mongoose.models.AgentTriggerDelivery.deleteMany({ deliveryKey }).catch(() => undefined); + await mongoose.models.AgentTriggerLaneSequence.deleteMany({ + orderingKey: { $regex: runId }, + }).catch(() => undefined); + + const width = Math.max(...Object.keys(verdicts).map((key) => key.length)); + const rows = Object.entries(verdicts).map( + ([label, verdict]) => ` ${label.padEnd(width)} ${verdict}`, + ); + console.log(`\nDocumentDB audit verdicts (run ${runId}):\n${rows.join('\n')}\n`); + await mongoose.disconnect(); + }); + + it('reports the engine build so the verdicts are attributable', async () => { + const info = await getDb().admin().command({ buildInfo: 1 }); + const version = typeof info.version === 'string' ? info.version : 'unknown'; + verdicts['engine version'] = version; + expect(version).toBeTruthy(); + }); + + describe('raw construct probes (isolates which primitive the engine refuses)', () => { + it('probes the pipeline-update form', async () => { + const verdict = await probe('pipeline-form findOneAndUpdate', () => + getDb() + .collection(probeCollection) + .findOneAndUpdate({ probe: 1 }, [{ $set: { touched: true } }]), + ); + expect(verdict).toBeTruthy(); + }); + + it('probes $$REMOVE, documented unsupported on every engine', async () => { + const verdict = await probe('$$REMOVE in $project', () => + getDb() + .collection(probeCollection) + .aggregate([ + { $project: { kept: '$label', dropped: { $cond: [false, '$note', '$$REMOVE'] } } }, + ]) + .toArray(), + ); + expect(verdict).toBeTruthy(); + }); + + it('probes $facet, documented unsupported on every engine', async () => { + const verdict = await probe('$facet stage', () => + getDb() + .collection(probeCollection) + .aggregate([{ $facet: { rows: [{ $project: { label: 1 } }] } }]) + .toArray(), + ); + expect(verdict).toBeTruthy(); + }); + + it('probes the operators the subagent projections rely on', async () => { + await probe('$regexMatch', () => + getDb() + .collection(probeCollection) + .aggregate([ + { $addFields: { matched: { $regexMatch: { input: '$label', regex: 'audit' } } } }, + ]) + .toArray(), + ); + await probe('$switch', () => + getDb() + .collection(probeCollection) + .aggregate([ + { + $addFields: { + branch: { $switch: { branches: [{ case: true, then: 'a' }], default: 'b' } }, + }, + }, + ]) + .toArray(), + ); + await probe('$let', () => + getDb() + .collection(probeCollection) + .aggregate([{ $addFields: { bound: { $let: { vars: { one: 1 }, in: '$$one' } } } }]) + .toArray(), + ); + await probe('$convert', () => + getDb() + .collection(probeCollection) + .aggregate([ + { + $addFields: { + text: { $convert: { input: '$label', to: 'string', onError: '', onNull: '' } }, + }, + }, + ]) + .toArray(), + ); + await probe('$strLenBytes + $substrCP', () => + getDb() + .collection(probeCollection) + .aggregate([ + { + $addFields: { + size: { $strLenBytes: '$label' }, + head: { $substrCP: ['$label', 0, 2] }, + }, + }, + ]) + .toArray(), + ); + await probe('$mergeObjects + $map', () => + getDb() + .collection(probeCollection) + .aggregate([ + { + $addFields: { + merged: { + $map: { + input: '$values', + as: 'value', + in: { $mergeObjects: [{ v: '$$value' }, { seen: true }] }, + }, + }, + }, + }, + ]) + .toArray(), + ); + for (const label of [ + '$regexMatch', + '$switch', + '$let', + '$convert', + '$strLenBytes + $substrCP', + '$mergeObjects + $map', + ]) { + expect(verdicts[label]).toBeTruthy(); + } + }); + + it('probes the classic-operator replacements the fix would use', async () => { + await probe('$max update operator', () => + getDb() + .collection(probeCollection) + .updateOne({ probe: 1 }, { $max: { leaseUntil: new Date() } }), + ); + await probe('$set + $unset classic update', () => + getDb() + .collection(probeCollection) + .updateOne({ probe: 1 }, { $set: { claimed: true }, $unset: { dropped: 1 } }), + ); + expectShape('$max update operator'); + expectShape('$set + $unset classic update'); + }); + }); + + describe('production shapes (drives the real methods)', () => { + it('site 1 - claimNextAgentTriggerDelivery', async () => { + await probe('enqueueAgentTriggerDelivery', () => + triggerMethods.enqueueAgentTriggerDelivery({ + deliveryKey, + fingerprint: `audit-fingerprint-${runId}`, + orderingKey: `audit-ordering-${runId}`, + envelope: { version: 1, audit: runId }, + user: userId, + availableAt: new Date(), + }), + ); + const now = new Date(); + await probe('claimNextAgentTriggerDelivery', () => + triggerMethods.claimNextAgentTriggerDelivery({ + workerId: `audit-worker-${runId}`, + claimToken: randomUUID(), + now, + leaseUntil: new Date(now.getTime() + 60_000), + }), + ); + expectShape('claimNextAgentTriggerDelivery'); + }); + + it('site 2 - renewAgentTriggerDeliveryProducerLease', async () => { + await probe('renewAgentTriggerDeliveryProducerLease', () => + triggerMethods.renewAgentTriggerDeliveryProducerLease({ + deliveryKey, + sourceId, + leaseUntil: new Date(Date.now() + 60_000), + }), + ); + expectShape('renewAgentTriggerDeliveryProducerLease'); + }); + + it('site 3 - claimBackgroundToolResults', async () => { + /** This method returns `not_found`/`not_ready` before it ever builds the + * update, so an unseeded probe would report a false `accepted`. Seed a + * terminal, wakeup-eligible, unclaimed task so the claim path is reached. */ + await mongoose.models.Message.create({ + messageId, + conversationId, + user: String(userId), + isCreatedByUser: false, + content: [ + { + type: 'tool_call', + tool_call: { + backgroundTask: { + taskId: `audit-task-${runId}`, + status: 'completed', + completionWakeup: true, + }, + }, + }, + ], + }); + await probe('claimBackgroundToolResults', () => + messageMethods.claimBackgroundToolResults({ + userId: String(userId), + conversationId, + messageId, + taskId: `audit-task-${runId}`, + kind: 'wakeup', + claimId: randomUUID(), + }), + ); + expectShape('claimBackgroundToolResults'); + }); + + it('site 4 - releaseBackgroundToolResultClaims', async () => { + await probe('releaseBackgroundToolResultClaims', () => + messageMethods.releaseBackgroundToolResultClaims({ + userId: String(userId), + conversationId, + messageId, + taskIds: [`audit-task-${runId}`], + kind: 'wakeup', + claimId: randomUUID(), + }), + ); + expectShape('releaseBackgroundToolResultClaims'); + }); + + it('site 5 - getMessagesForSubagentThreadView ($$REMOVE path)', async () => { + await probe('getMessagesForSubagentThreadView (list)', () => + messageMethods.getMessagesForSubagentThreadView({ + user: String(userId), + conversationId, + limit: 10, + textCodePointLimit: 512, + }), + ); + expectShape('getMessagesForSubagentThreadView (list)'); + }); + + it('site 6 - getMessagesForSubagentThreadView ($facet path)', async () => { + await probe('getMessagesForSubagentThreadView (selected)', () => + messageMethods.getMessagesForSubagentThreadView({ + user: String(userId), + conversationId, + selectedTaskId: `audit-task-${runId}`, + limit: 10, + textCodePointLimit: 512, + }), + ); + expectShape('getMessagesForSubagentThreadView (selected)'); + }); + + it('listSubagentTasksForThreads ($regexMatch path)', async () => { + await probe('listSubagentTasksForThreads', () => + messageMethods.listSubagentTasksForThreads({ + user: String(userId), + conversationIds: [conversationId], + limitPerThread: 4, + }), + ); + expectShape('listSubagentTasksForThreads'); + }); + + it('site 7 - updateToolCallResult', async () => { + const settleMessageId = `audit-settle-${runId}`; + await mongoose.models.Message.create({ + messageId: settleMessageId, + conversationId, + user: String(userId), + isCreatedByUser: false, + content: [ + { + type: 'tool_call', + tool_call: { id: `audit-call-${runId}`, name: 'execute_code', output: 'pending' }, + }, + ], + attachments: [{ file_id: `audit-file-${runId}`, toolCallId: `audit-call-${runId}` }], + }); + await probe('updateToolCallResult', () => + messageMethods.updateToolCallResult({ + userId: String(userId), + messageId: settleMessageId, + conversationId, + toolCallId: `audit-call-${runId}`, + output: 'settled output', + markBackgrounded: true, + backgroundTask: { + taskId: `audit-task-${runId}`, + toolName: 'execute_code', + status: 'completed', + settledAt: new Date(), + resultClaim: { kind: 'wakeup', claimId: `audit-claim-${runId}`, claimedAt: new Date() }, + }, + attachments: [{ file_id: `audit-file-${runId}`, toolCallId: `audit-call-${runId}` }], + }), + ); + expectShape('updateToolCallResult'); + }); + }); +}); diff --git a/packages/data-schemas/misc/documentdb/compat.documentdb.spec.ts b/packages/data-schemas/misc/documentdb/compat.documentdb.spec.ts index 33ea021d07..2f993d57bf 100644 --- a/packages/data-schemas/misc/documentdb/compat.documentdb.spec.ts +++ b/packages/data-schemas/misc/documentdb/compat.documentdb.spec.ts @@ -36,10 +36,14 @@ import { createModels } from '~/models'; * suite only runs when DOCUMENTDB_URI is set and skips otherwise. * * Run (from packages/data-schemas, against a DEDICATED database): - * DOCUMENTDB_URI="mongodb://user:pass@127.0.0.1:27017/librechat_compat?tls=true&retryWrites=false" \ + * DOCUMENTDB_URI="mongodb://user:pass@127.0.0.1:27017/librechat_compat\ + * ?tls=true&retryWrites=false&authSource=admin&authMechanism=SCRAM-SHA-1&directConnection=true" \ * DOCUMENTDB_TLS_CA_FILE="global-bundle.pem" \ * npx jest --config misc/documentdb/jest.documentdb.config.mjs * + * `authSource`/`authMechanism`/`directConnection` are load-bearing against a + * real cluster (see audit.documentdb.spec.ts for why each is required). + * * Through an SSH tunnel, additionally set * DOCUMENTDB_TLS_ALLOW_INVALID_HOSTNAMES=true * because the tunnel endpoint will not match the cluster certificate. diff --git a/packages/data-schemas/misc/documentdb/documentdb-compat.md b/packages/data-schemas/misc/documentdb/documentdb-compat.md index d3569d9687..eac86819de 100644 --- a/packages/data-schemas/misc/documentdb/documentdb-compat.md +++ b/packages/data-schemas/misc/documentdb/documentdb-compat.md @@ -1,9 +1,18 @@ # Amazon DocumentDB Compatibility Assessment (issue #14488) +> **Update 2026-08-30 — adjudicated live.** Everything below was originally +> decided from AWS's documentation, which omits unsupported operators rather +> than listing them. It has now been run against a real DocumentDB **5.0.0** +> cluster (the version this project supports, and the one the reference +> deployment runs). `audit.documentdb.spec.ts` records the verdicts; the +> corrections are in "Live verdicts" at the end of this document. Two claims +> below were wrong: pipeline-form updates had returned in new code, and the +> transaction probe reported a false negative. + Adjudicated against official AWS documentation on 2026-07-28. Engine columns throughout: DocumentDB **3.6 / 4.0 / 5.0 / 8.0 instance-based** and **elastic clusters**. AWS's supported-APIs page states that unsupported operators are -*omitted* from its tables, so several verdicts below are implicit-by-omission +_omitted_ from its tables, so several verdicts below are implicit-by-omission and flagged as such. ## Executive summary @@ -33,7 +42,7 @@ this and DocumentDB rejects it. AWS evidence: the [supported APIs page](https://docs.aws.amazon.com/documentdb/latest/developerguide/mongo-apis.html) lists only classic update operators (no pipeline form anywhere); the `$set`/ -`$unset` *stage* operators are marked unsupported for 3.6/4.0/5.0; `$$NOW` is +`$unset` _stage_ operators are marked unsupported for 3.6/4.0/5.0; `$$NOW` is absent from the System variables table entirely (`$$CURRENT` and `$$REMOVE` are explicitly "No"). Implicit-by-omission, but consistent with the reported `Failed to parse update: field must be of BSON type object` class of error @@ -150,7 +159,7 @@ in the docs. decides whether the partial-index caveat applies to them (5.0+: it doesn't). Worth asking directly on the issue. - **DocumentDB 8.0 pipeline-update acceptance** — 8.0 added `$set`/`$unset` - aggregation *stages*, but AWS never documents pipeline-form updates; the + aggregation _stages_, but AWS never documents pipeline-form updates; the harness probe answers this live. - **`collMod` is only "Partial"** on every version — avoid `Model.syncIndexes()` against DocumentDB (it may issue `collMod` beyond the @@ -180,16 +189,63 @@ in the docs. ## Support matrix and recommendation -| Capability (LibreChat dependency) | 3.6 | 4.0 | 5.0 | 8.0 | Elastic | -| --- | --- | --- | --- | --- | --- | -| Pipeline updates (**no longer used**) | ✗ | ✗ | ✗ | ? | ✗ | -| Plain update operators (all writes now) | ✓ | ✓ | ✓ | ✓ | ✓ | -| Unique indexes | ✓ | ✓ | ✓ | ✓ | ✗ | -| Partial unique indexes (OAuth ids) | ✗ | ✗ | ✓ | ✓ | ✗ | -| Transactions (runtime-probed) | ✗ | ✓ | ✓ | ✓ | ✗ | -| TTL indexes | ✓ | ✓ | ✓ | ✓ | ✓ | +| Capability (LibreChat dependency) | 3.6 | 4.0 | 5.0 | 8.0 | Elastic | +| --------------------------------------- | --- | --- | --- | --- | ------- | +| Pipeline updates (**no longer used**) | ✗ | ✗ | ✗ | ? | ✗ | +| Plain update operators (all writes now) | ✓ | ✓ | ✓ | ✓ | ✓ | +| Unique indexes | ✓ | ✓ | ✓ | ✓ | ✗ | +| Partial unique indexes (OAuth ids) | ✗ | ✗ | ✓ | ✓ | ✗ | +| Transactions (runtime-probed) | ✗ | ✓ | ✓ | ✓ | ✗ | +| TTL indexes | ✓ | ✓ | ✓ | ✓ | ✓ | **Recommendation**: support **DocumentDB 5.0+ instance-based** with `retryWrites=false` documented as required. 4.0 functions with partial-unique-index loss (now logged loudly at startup) — "works, with a documented caveat". Elastic clusters: unsupported, full stop. + +## Live verdicts (DocumentDB 5.0.0, 2026-08-30) + +Probed by `audit.documentdb.spec.ts`, which drives the production methods +themselves rather than re-implementations. + +| Construct | Verdict | Server error | +| ---------------------------------------------------- | -------- | ----------------------------------------------------------- | +| Aggregation-pipeline update | rejected | `Failed to parse update: field must be of BSON type object` | +| `$$REMOVE` | rejected | `Feature not supported: $$REMOVE` | +| `$facet` | rejected | `Aggregation stage not supported: '$facet'` | +| `$max`, `$set`/`$unset` | accepted | — | +| Filtered positional `$[]` | accepted | — | +| `$regexMatch`, `$switch`, `$let`, `$convert` | accepted | — | +| `$strLenBytes`, `$substrCP`, `$mergeObjects`, `$map` | accepted | — | +| Partial unique indexes, TTL indexes | accepted | — | + +### Correction 1 — pipeline updates returned after this document was written + +Six sites reintroduced unsupported constructs between 2026-07-29 and +2026-08-30, all in code added with the durable trigger and background-task +work. Nothing caught them: every unit suite runs `mongodb-memory-server`, which +is real MongoDB and accepts all of it. `src/methods/documentdb.spec.ts` is now +a static guard against the whole class. + +### Correction 2 — the transaction probe reported a false negative + +`supportsTransactions` read a canary collection that does not exist, and +DocumentDB rejects a transaction touching a non-existent collection +(`Feature not supported: non-existent collection in transaction`). The probe +therefore returned `false` on an engine that fully supports transactions, and +every caller silently took the non-transactional path. Verified directly: with +the collection materialized, both read-only and multi-write transactions +commit. The probe now creates the canary first. + +### Connection requirements for the live suites + +Established against the real cluster; all four are load-bearing through a +tunnel and none were documented before: + +- `authSource=admin` — the user lives in `admin`; a database in the URI path + otherwise becomes the auth source and authentication fails +- `authMechanism=SCRAM-SHA-1` — DocumentDB rejects SCRAM-SHA-256 + (`Unsupported mechanism [ -301 ]`) +- `directConnection=true` — replica-set discovery returns internal cluster + hostnames that are unreachable through a tunnel +- `tlsAllowInvalidHostnames` — the tunnel endpoint never matches the certificate diff --git a/packages/data-schemas/src/methods/documentdb.spec.ts b/packages/data-schemas/src/methods/documentdb.spec.ts new file mode 100644 index 0000000000..0a2f717579 --- /dev/null +++ b/packages/data-schemas/src/methods/documentdb.spec.ts @@ -0,0 +1,247 @@ +import fs from 'fs'; +import path from 'path'; +import ts from 'typescript'; + +/** + * Static guard for the Amazon DocumentDB write and read surface. + * + * DocumentDB is a supported deployment target (see the root README and + * `misc/documentdb/documentdb-compat.md`), but nothing in the normal test + * pyramid can catch an incompatibility: every suite runs against + * `mongodb-memory-server`, which is real MongoDB and accepts all of these + * constructs. Each verdict below was established against a live DocumentDB + * 5.0.0 cluster — the engine version this project documents as supported — and + * the exact server error is recorded beside it. + * + * The scan walks the TypeScript AST rather than source text, so it also + * catches a pipeline bound to a variable first, cast with `as`, or nested in a + * `bulkWrite` operation's `update` property. It covers every backend workspace + * that talks to MongoDB — this package, `packages/api`, and `api` — because the + * regression class is repo-wide and new backend code lands in `packages/api`. + * If a construct here becomes genuinely necessary, the fix is a compatible + * rewrite, not an exception list: `misc/documentdb/audit.documentdb.spec.ts` + * re-adjudicates any of this against a real cluster. + */ +const REPO_ROOT = path.join(__dirname, '../../../..'); +const SCAN_ROOTS = [ + path.join(REPO_ROOT, 'packages/data-schemas/src'), + path.join(REPO_ROOT, 'packages/api/src'), + path.join(REPO_ROOT, 'api'), +]; +const SKIPPED_DIRECTORIES = new Set(['node_modules', 'dist', 'build', 'coverage', '.turbo']); + +/** Aggregation-pipeline updates: `Failed to parse update: field must be of BSON type object`. */ +const UPDATE_METHODS = new Set([ + 'findOneAndUpdate', + 'findByIdAndUpdate', + 'findOneAndReplace', + 'updateOne', + 'updateMany', + 'replaceOne', +]); + +/** `$$REMOVE`: `Feature not supported: $$REMOVE`. `$facet`: `Aggregation stage not supported`. */ +const FORBIDDEN_TOKENS = ['$$REMOVE', '$$CURRENT', '$facet', '$graphLookup', '$unionWith']; + +function collectSourceFiles(directory: string, found: string[] = []): string[] { + for (const entry of fs.readdirSync(directory, { withFileTypes: true })) { + if (entry.isDirectory() && SKIPPED_DIRECTORIES.has(entry.name)) { + continue; + } + const target = path.join(directory, entry.name); + if (entry.isDirectory()) { + collectSourceFiles(target, found); + continue; + } + if (/\.(ts|js)$/.test(entry.name) && !/\.(spec|test)\.[tj]s$|\.d\.ts$/.test(entry.name)) { + found.push(target); + } + } + return found; +} + +function parse(fileName: string, source: string): ts.SourceFile { + return ts.createSourceFile(fileName, source, ts.ScriptTarget.Latest, true); +} + +/** Peels casts and parentheses so `[...] as PipelineStage[]` is still an array. */ +function unwrapExpression(expression: ts.Expression): ts.Expression { + let current = expression; + while ( + ts.isAsExpression(current) || + ts.isSatisfiesExpression(current) || + ts.isParenthesizedExpression(current) || + ts.isNonNullExpression(current) + ) { + current = current.expression; + } + return current; +} + +/** Names of variables initialized with array literals, so an indirect + * `const update = [...]; Model.updateOne(filter, update)` is still caught. + * Scope-naive by design: a false positive here names a variable that holds an + * array and is passed as an update, which deserves a look regardless. */ +function collectArrayVariableNames(sourceFile: ts.SourceFile): Set { + const names = new Set(); + const visit = (node: ts.Node): void => { + if ( + ts.isVariableDeclaration(node) && + ts.isIdentifier(node.name) && + node.initializer != null && + ts.isArrayLiteralExpression(unwrapExpression(node.initializer)) + ) { + names.add(node.name.text); + } + ts.forEachChild(node, visit); + }; + visit(sourceFile); + return names; +} + +function isArrayValued(expression: ts.Expression, arrayNames: Set): boolean { + const unwrapped = unwrapExpression(expression); + if (ts.isArrayLiteralExpression(unwrapped)) { + return true; + } + if (ts.isIdentifier(unwrapped)) { + return arrayNames.has(unwrapped.text); + } + if (ts.isConditionalExpression(unwrapped)) { + return ( + isArrayValued(unwrapped.whenTrue, arrayNames) || + isArrayValued(unwrapped.whenFalse, arrayNames) + ); + } + return false; +} + +function offenseAt(sourceFile: ts.SourceFile, node: ts.Node, label: string): string { + const { line } = sourceFile.getLineAndCharacterOfPosition(node.getStart(sourceFile)); + return `${sourceFile.fileName}:${line + 1} ${label}`; +} + +/** Reports every update argument that is an array — the pipeline-update form — + * whether passed directly to an update method or carried inside a `bulkWrite` + * operation's `update` property. */ +function findPipelineUpdates(sourceFile: ts.SourceFile): string[] { + const arrayNames = collectArrayVariableNames(sourceFile); + const offenses: string[] = []; + const visit = (node: ts.Node): void => { + if ( + ts.isCallExpression(node) && + ts.isPropertyAccessExpression(node.expression) && + UPDATE_METHODS.has(node.expression.name.text) && + node.arguments.length >= 2 && + isArrayValued(node.arguments[1], arrayNames) + ) { + offenses.push(offenseAt(sourceFile, node, node.expression.name.text)); + } + if ( + ts.isPropertyAssignment(node) && + (ts.isIdentifier(node.name) || ts.isStringLiteral(node.name)) && + node.name.text === 'update' && + isArrayValued(node.initializer, arrayNames) + ) { + offenses.push(offenseAt(sourceFile, node, 'update property')); + } + if ( + ts.isShorthandPropertyAssignment(node) && + node.name.text === 'update' && + arrayNames.has(node.name.text) + ) { + offenses.push(offenseAt(sourceFile, node, 'update property')); + } + ts.forEachChild(node, visit); + }; + visit(sourceFile); + return offenses; +} + +/** Reports forbidden operator tokens in string literals and property names, + * ignoring prose — the rewrites explain themselves by naming the construct. */ +function findForbiddenTokens(sourceFile: ts.SourceFile): string[] { + const offenses: string[] = []; + const visit = (node: ts.Node): void => { + if (ts.isStringLiteralLike(node) || ts.isIdentifier(node)) { + for (const token of FORBIDDEN_TOKENS) { + if (node.text === token || node.text.startsWith(`${token}.`)) { + offenses.push(offenseAt(sourceFile, node, token)); + } + } + } + ts.forEachChild(node, visit); + }; + visit(sourceFile); + return offenses; +} + +describe('Amazon DocumentDB compatibility', () => { + /** Each file is read and parsed once; both detectors walk the same tree. */ + const parsedSources = SCAN_ROOTS.flatMap((root) => collectSourceFiles(root)).map((file) => + parse(path.relative(REPO_ROOT, file), fs.readFileSync(file, 'utf8')), + ); + + it('scans the backend workspaces', () => { + expect(parsedSources.length).toBeGreaterThan(0); + }); + + it('uses no aggregation-pipeline updates', () => { + expect(parsedSources.flatMap(findPipelineUpdates)).toEqual([]); + }); + + it('uses no aggregation constructs the engine rejects', () => { + expect(parsedSources.flatMap(findForbiddenTokens)).toEqual([]); + }); + + /** A guard that cannot fail protects nothing, so every shape the detectors + * exist to catch is proven against fixtures here rather than by trusting a + * one-time manual injection. */ + describe('detector coverage', () => { + it.each([ + ['direct literal', `Model.updateOne(filter, [{ $set: { a: 1 } }]);`], + ['cast literal', `Model.updateMany(filter, [{ $set: { a: 1 } }] as PipelineStage[]);`], + [ + 'indirect variable', + `const update = [{ $set: { a: 1 } }];\nModel.findOneAndUpdate(filter, update);`, + ], + [ + 'conditional branch', + `Model.updateOne(filter, flag ? [{ $set: { a: 1 } }] : { $set: { a: 1 } });`, + ], + [ + 'bulkWrite operation payload', + `await Model.bulkWrite([{ updateOne: { filter, update: [{ $set: { a: 1 } }] } }]);`, + ], + [ + 'bulkWrite indirect payload', + `const update = [{ $set: { a: 1 } }];\nawait Model.bulkWrite([{ updateMany: { filter, update } }]);`, + ], + ])('flags a pipeline update: %s', (_shape, source) => { + expect(findPipelineUpdates(parse('fixture.ts', source))).not.toEqual([]); + }); + + it.each([ + ['classic update', `Model.updateOne(filter, { $set: { a: 1 } });`], + [ + 'classic bulk payload', + `await Model.bulkWrite([{ updateOne: { filter, update: { $set: { a: 1 } } } }]);`, + ], + ['unrelated array variable', `const stages = [{ $match: {} }];\nModel.aggregate(stages);`], + ])('accepts a supported shape: %s', (_shape, source) => { + expect(findPipelineUpdates(parse('fixture.ts', source))).toEqual([]); + }); + + it('flags forbidden operators in code but not in prose', () => { + expect( + findForbiddenTokens(parse('fixture.ts', `const projection = { x: '$$REMOVE' };`)), + ).not.toEqual([]); + expect( + findForbiddenTokens(parse('fixture.ts', `pipeline.push({ $facet: { rows: [] } });`)), + ).not.toEqual([]); + expect( + findForbiddenTokens(parse('fixture.ts', `/** $$REMOVE and $facet are unsupported. */`)), + ).toEqual([]); + }); + }); +}); diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index 28de4cc2f6..549eb73400 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -955,6 +955,145 @@ describe('Message Operations', () => { }); }); + it('stamps a supplied claim only on a part that carries none', async () => { + await saveMessage(mockCtx, { ...mockMessageData, content: toolCallContent() }); + + await updateToolCallResult({ + userId: 'user123', + messageId: 'msg123', + conversationId: mockMessageData.conversationId as string, + toolCallId: 'call_bg', + backgroundTask: { + taskId: 'task-1', + toolName: 'execute_code', + status: 'completed', + settledAt: new Date('2026-08-30T12:00:02Z'), + resultClaim: { + kind: 'wakeup', + claimId: 'wakeup-1', + claimedAt: new Date('2026-08-30T12:00:02Z'), + }, + }, + }); + + const saved = await Message.findOne({ messageId: 'msg123', user: 'user123' }).lean(); + const task = ( + saved?.content?.[1] as { + tool_call?: { backgroundTask?: Record }; + } + ).tool_call?.backgroundTask; + expect(task).toMatchObject({ + taskId: 'task-1', + resultClaim: { kind: 'wakeup', claimId: 'wakeup-1' }, + }); + }); + + it('reports exhausted attachment-merge contention as unmatched instead of throwing', async () => { + await saveMessage(mockCtx, { ...mockMessageData, content: toolCallContent() }); + const Message = mongoose.models.Message; + const casSpy = jest.spyOn(Message.collection, 'findOneAndUpdate').mockResolvedValue(null); + try { + /** A throw here escapes the settle retry loop and lands on ambiguous + * failure handling that can retire the completion; contention must + * surface as a retryable unmatched result. */ + await expect( + updateToolCallResult({ + userId: 'user123', + messageId: 'msg123', + conversationId: mockMessageData.conversationId as string, + toolCallId: 'call_bg', + output: 'settled output', + attachments: [{ file_id: 'file-contended', toolCallId: 'call_bg' }], + }), + ).resolves.toEqual({ matched: false, unfinished: false }); + expect(casSpy.mock.calls.length).toBeGreaterThanOrEqual(8); + } finally { + casSpy.mockRestore(); + } + }); + + it('disarms a previously armed completion wakeup when the new receipt omits it', async () => { + const content = toolCallContent(); + (content[1].tool_call as Record).backgroundTask = { + version: 1, + taskId: 'task-1', + toolName: 'execute_code', + status: 'completed', + settledAt: new Date('2026-08-30T12:00:00Z'), + completionWakeup: true, + }; + await saveMessage(mockCtx, { ...mockMessageData, content }); + + await updateToolCallResult({ + userId: 'user123', + messageId: 'msg123', + conversationId: mockMessageData.conversationId as string, + toolCallId: 'call_bg', + backgroundTask: { + taskId: 'task-1', + toolName: 'execute_code', + status: 'error', + settledAt: new Date('2026-08-30T12:00:02Z'), + }, + }); + + const saved = await Message.findOne({ messageId: 'msg123', user: 'user123' }).lean(); + const task = ( + saved?.content?.[1] as { + tool_call?: { backgroundTask?: Record }; + } + ).tool_call?.backgroundTask; + expect(task).toMatchObject({ taskId: 'task-1', status: 'error' }); + expect(task).not.toHaveProperty('completionWakeup'); + }); + + it('claims a terminal result whose persisted claim stamp is a stored null', async () => { + /** A full-row save can persist `resultClaim: null`, and the + * subfield-preserving settle write keeps it where the old whole-object + * rewrite dropped it. Every layer must read that as unclaimed — a + * `$exists: false` fence would strand the part as terminal but + * permanently unclaimable. */ + await saveMessage(mockCtx, { + ...mockMessageData, + content: [ + { + type: 'tool_call', + tool_call: { + id: 'call-null-claim', + name: 'slow_tool', + output: 'done', + backgroundTask: { + version: 1, + taskId: 'task-null-claim', + toolName: 'slow_tool', + status: 'completed', + settledAt: new Date(), + completionWakeup: true, + resultClaim: null, + }, + }, + }, + ], + }); + + /** The stored stamp must actually BE null, or this test proves nothing. */ + const stored = await Message.findOne({ messageId: 'msg123', user: 'user123' }).lean(); + const storedTask = ( + stored?.content?.[0] as { tool_call?: { backgroundTask?: Record } } + ).tool_call?.backgroundTask; + expect(storedTask).toHaveProperty('resultClaim', null); + + const claim = await claimBackgroundToolResults({ + userId: 'user123', + conversationId: mockMessageData.conversationId as string, + messageId: 'msg123', + taskId: 'task-null-claim', + kind: 'wakeup', + claimId: 'delivery-null', + }); + expect(claim.status).toBe('acquired'); + }); + it('elects one result consumer and batches terminal siblings for a wakeup', async () => { const terminal = (id: string, taskId: string, output: string) => ({ type: 'tool_call', @@ -1975,6 +2114,9 @@ describe('Message Operations', () => { message.subagentTranscriptProjectionTruncated === true, ), ).toHaveLength(1); + /** Three reads: the page, recent sources, and one single-snapshot read + * that yields both the selected task's messages and its source — the + * replacement for a `$facet`, which Amazon DocumentDB does not support. */ expect(aggregateSpy).toHaveBeenCalledTimes(3); const messagesPipeline = aggregateSpy.mock.calls[0][0] as unknown as Array< Record @@ -1985,10 +2127,12 @@ describe('Message Operations', () => { const selectedPipeline = aggregateSpy.mock.calls[2][0] as unknown as Array< Record >; + for (const [pipeline] of aggregateSpy.mock.calls) { + expect(pipeline).not.toEqual( + expect.arrayContaining([expect.objectContaining({ $facet: expect.anything() })]), + ); + } expect(messagesPipeline[2]).toEqual({ $limit: SUBAGENT_TRANSCRIPT_PAGE_LIMIT }); - expect(messagesPipeline).not.toEqual( - expect.arrayContaining([expect.objectContaining({ $facet: expect.anything() })]), - ); expect(recentSourcesPipeline[2]).toEqual({ $limit: SUBAGENT_TRANSCRIPT_PAGE_LIMIT * 2, }); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 5d2729de60..14f1da3078 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -445,6 +445,61 @@ export type SubagentThreadViewMessageRecord = Pick< }; }; +/** Amazon DocumentDB does not support `$$REMOVE`, so the bounded thread-view + * projections emit `null` where they mean "omit this key". This is the shape as + * it leaves the aggregation, before those sentinels are pruned back to absent. */ +/** Widens the given keys to admit the projection's `null` sentinel. */ +type WithNullSentinels = Omit & { + [P in K]?: NonNullable | null; +}; + +type ThreadViewRecord = SubagentThreadViewMessageRecord; +export type ProjectedSubagentThreadViewMessage = WithNullSentinels< + Omit, + 'subagentTranscriptProjectionTruncated' +> & { + subagentTriggerProjection?: WithNullSentinels< + NonNullable, + 'expectedActionToolName' + > | null; + subagentTask?: + | (Omit, 'controlReceipts'> & { + controlReceipts?: Array< + WithNullSentinels< + NonNullable['controlReceipts']>[number], + 'controlId' | 'reason' | 'message' + > + >; + }) + | null; +}; + +/** Restores the absent-vs-present contract by dropping the `null` sentinels the + * projection emitted. Bounded by the projection's own row, receipt, and byte + * limits, and folded into the pass that already materializes each record. */ +function pruneProjectedThreadViewMessage( + message: ProjectedSubagentThreadViewMessage, +): SubagentThreadViewMessageRecord { + if (message.subagentTranscriptProjectionTruncated === null) { + delete message.subagentTranscriptProjectionTruncated; + } + if (message.subagentTriggerProjection === null) { + delete message.subagentTriggerProjection; + } else if (message.subagentTriggerProjection?.expectedActionToolName === null) { + delete message.subagentTriggerProjection.expectedActionToolName; + } + if (message.subagentTask === null) { + delete message.subagentTask; + } else { + for (const receipt of message.subagentTask?.controlReceipts ?? []) { + if (receipt.controlId === null) delete receipt.controlId; + if (receipt.reason === null) delete receipt.reason; + if (receipt.message === null) delete receipt.message; + } + } + return message as SubagentThreadViewMessageRecord; +} + export type ParentSubagentTaskRecord = { conversationId: string; /** The shared bounded source window filled, so this child's history may be incomplete. */ @@ -621,6 +676,21 @@ export interface MessageMethods { deleteMessages(filter: FilterQuery): Promise; } +/** The agent-ownership rule shared by background-tool settling and claiming: + * `agentId` on the part wins, then `tool_call.agentId`, and a part without + * agent identity belongs to any caller (single-agent runs). `field: null` + * matches both a missing field and a stored null. The settle and claim paths + * MUST apply the identical rule, or a settled part becomes unclaimable. */ +function agentOwnershipFilter(prefix: string, agentId: string): Record { + return { + $or: [ + { [`${prefix}agentId`]: agentId }, + { [`${prefix}agentId`]: null, [`${prefix}tool_call.agentId`]: agentId }, + { [`${prefix}agentId`]: null, [`${prefix}tool_call.agentId`]: null }, + ], + }; +} + export function createMessageMethods(mongoose: typeof import('mongoose')): MessageMethods { /** * Saves a message in the database. @@ -959,193 +1029,184 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa }; }; }): Promise<{ matched: boolean; unfinished: boolean }> { - const targetPartMatch = { - $and: [ - { $eq: ['$$part.type', 'tool_call'] }, - { $eq: ['$$part.tool_call.id', toolCallId] }, - ...(stepId != null ? [{ $eq: ['$$part.tool_call.stepId', stepId] }] : []), - ...(agentId != null - ? [ - { - $in: [ - { - $ifNull: [{ $ifNull: ['$$part.agentId', '$$part.tool_call.agentId'] }, null], - }, - [null, agentId], - ], - }, - ] - : []), - ], + /** One source of truth for which content part this settle may touch: + * `prefix: ''` yields the `$elemMatch` document filter, `prefix: 'part.'` + * the arrayFilters element filter — the same predicate in one dialect, + * where the old code kept an aggregation `$expr` twin of it. `field: null` + * matches both a missing field and a stored null, mirroring the previous + * `$ifNull` chains. */ + const partScope = (prefix: string): Record => ({ + [`${prefix}type`]: 'tool_call', + [`${prefix}tool_call.id`]: toolCallId, + ...(stepId != null ? { [`${prefix}tool_call.stepId`]: stepId } : {}), + ...(agentId != null ? agentOwnershipFilter(prefix, agentId) : {}), + }); + const messageFilter = { + messageId, + user: userId, + conversationId, + content: { $elemMatch: partScope('') }, }; - const stages: Record[] = []; + const partIdentityFilter = partScope('part.'); + /** Amazon DocumentDB rejects aggregation-pipeline updates, so the part + * patch addresses the matching tool-call parts with the filtered positional + * operator and plain `$set`/`$unset`. Setting `backgroundTask` subfields + * individually preserves an existing `resultClaim` without the per-part + * branching the old pipeline needed, and clearing `completionWakeup` keeps + * the old whole-object replacement's disarm semantics. */ + const partPatch: Record = {}; + let disarmWakeup = false; if (output !== undefined || backgroundTask != null) { - stages.push({ - $set: { - content: { - $map: { - input: { $ifNull: ['$content', []] }, - as: 'part', - in: { - $cond: [ - targetPartMatch, - { - $mergeObjects: [ - '$$part', - { - tool_call: { - $mergeObjects: [ - '$$part.tool_call', - { - ...(output !== undefined ? { output: { $literal: output } } : {}), - ...(markBackgrounded === true ? { backgrounded: true } : {}), - ...(backgroundTask != null - ? { - backgroundTask: { - $mergeObjects: [ - { - $literal: { - version: 1, - taskId: backgroundTask.taskId, - toolName: backgroundTask.toolName, - status: backgroundTask.status, - settledAt: backgroundTask.settledAt, - ...(backgroundTask.completionWakeup === true - ? { completionWakeup: true } - : {}), - }, - }, - { - $cond: [ - { - $ne: [ - { - $ifNull: [ - '$$part.tool_call.backgroundTask.resultClaim', - null, - ], - }, - null, - ], - }, - { - resultClaim: - '$$part.tool_call.backgroundTask.resultClaim', - }, - backgroundTask.resultClaim != null - ? { - $literal: { - resultClaim: backgroundTask.resultClaim, - }, - } - : {}, - ], - }, - ], - }, - } - : {}), - }, - ], - }, - }, - ], - }, - '$$part', - ], - }, - }, - }, - }, - }); + if (output !== undefined) { + partPatch['content.$[part].tool_call.output'] = output; + } + if (markBackgrounded === true) { + partPatch['content.$[part].tool_call.backgrounded'] = true; + } + if (backgroundTask != null) { + partPatch['content.$[part].tool_call.backgroundTask.version'] = 1; + partPatch['content.$[part].tool_call.backgroundTask.taskId'] = backgroundTask.taskId; + partPatch['content.$[part].tool_call.backgroundTask.toolName'] = backgroundTask.toolName; + partPatch['content.$[part].tool_call.backgroundTask.status'] = backgroundTask.status; + partPatch['content.$[part].tool_call.backgroundTask.settledAt'] = backgroundTask.settledAt; + if (backgroundTask.completionWakeup === true) { + partPatch['content.$[part].tool_call.backgroundTask.completionWakeup'] = true; + } else { + disarmWakeup = true; + } + } } - if (attachments !== undefined && attachments.length > 0) { - /** Dedupe key mirrors the resume merge: `file_id ?? filepath`, so - * download-fallback attachments (no `file_id`, only a filepath) stay - * idempotent across re-applications instead of duplicating per poll. */ - const attachmentKeys = attachments - .map((attachment) => { - const { file_id, filepath } = attachment as { file_id?: unknown; filepath?: unknown }; - return typeof file_id === 'string' ? file_id : filepath; - }) - .filter((key): key is string => typeof key === 'string'); - stages.push({ - $set: { - attachments: { - $concatArrays: [ - { - $filter: { - input: { $ifNull: ['$attachments', []] }, - as: 'existing', - /** Replace only THIS tool call's prior entries: sibling calls - * can legitimately share a `file_id` (the filename claim is - * per-conversation), and the client anchors attachments to - * cards by `toolCallId`. */ - cond: { - $not: [ - { - $and: [ - { - $in: [ - { $ifNull: ['$$existing.file_id', '$$existing.filepath'] }, - { $literal: attachmentKeys }, - ], - }, - { $eq: ['$$existing.toolCallId', toolCallId] }, - /** Provider tool-call ids repeat across agents in - * handoff messages; a sibling agent's attachment - * under the same id/key must survive (missing - * agent identity = legacy wildcard). */ - ...(agentId != null - ? [ - { - $in: [{ $ifNull: ['$$existing.agentId', null] }, [null, agentId]], - }, - ] - : []), - ...(stepId != null - ? [ - { - $in: [{ $ifNull: ['$$existing.stepId', null] }, [null, stepId]], - }, - ] - : []), - ], - }, - ], - }, - }, - }, - { $literal: attachments }, - ], - }, - }, - }); - } - if (stages.length === 0) { + const mergingAttachments = attachments !== undefined && attachments.length > 0; + if (Object.keys(partPatch).length === 0 && !mergingAttachments) { return { matched: false, unfinished: false }; } + const settleUpdate = { + ...(Object.keys(partPatch).length > 0 ? { $set: partPatch } : {}), + ...(disarmWakeup + ? { $unset: { 'content.$[part].tool_call.backgroundTask.completionWakeup': 1 } } + : {}), + }; + const settleOptions = { + new: true, + projection: { unfinished: 1 }, + ...(Object.keys(partPatch).length > 0 || disarmWakeup + ? { arrayFilters: [partIdentityFilter] } + : {}), + }; try { const Message = mongoose.models.Message as Model; - const result = await Message.findOneAndUpdate( - { - messageId, - user: userId, - conversationId, - $expr: { - $anyElementTrue: { - $map: { - input: { $ifNull: ['$content', []] }, - as: 'part', - in: targetPartMatch, - }, + /** A supplied claim is stamped BEFORE the settle write, on parts that + * carry none (`resultClaim: null` also admits a stored null, as the old + * `$ifNull` did). Ordering is what keeps the split safe: until the settle + * write lands, the part is not terminal, so a concurrent wakeup or manual + * poll sees `not_ready` and stands down; the old single pipeline wrote + * claim and receipt atomically, and a claim-after-settle order would + * instead expose a claimable terminal part that lets a second consumer + * deliver the same result. A crash between the writes is healed by the + * settle retry, whose claim write no-ops against its own stamp. */ + if (backgroundTask?.resultClaim != null) { + await Message.updateOne( + messageFilter, + { + $set: { + 'content.$[part].tool_call.backgroundTask.resultClaim': backgroundTask.resultClaim, }, }, - }, - stages, - { new: true, projection: { unfinished: 1 } }, - ).lean<{ unfinished?: boolean } | null>(); - return { matched: result != null, unfinished: result?.unfinished === true }; + { + arrayFilters: [ + { ...partIdentityFilter, 'part.tool_call.backgroundTask.resultClaim': null }, + ], + }, + ); + } + if (!mergingAttachments) { + const result = await Message.findOneAndUpdate( + messageFilter, + settleUpdate, + settleOptions, + ).lean<{ unfinished?: boolean } | null>(); + return { matched: result != null, unfinished: result?.unfinished === true }; + } + /** Dedupe key mirrors the resume merge: `file_id ?? filepath`, so + * download-fallback attachments (no `file_id`, only a filepath) stay + * idempotent across re-applications instead of duplicating per poll. + * Replace only THIS tool call's prior entries: sibling calls can + * legitimately share a `file_id` (the filename claim is per-conversation), + * and the client anchors attachments to cards by `toolCallId`. Provider + * tool-call ids repeat across agents in handoff messages, so a sibling + * agent's attachment under the same id/key must survive (missing agent + * identity = legacy wildcard). */ + const attachmentKeys = new Set( + attachments + .map((attachment) => { + const { file_id, filepath } = attachment as { file_id?: unknown; filepath?: unknown }; + return typeof file_id === 'string' ? file_id : filepath; + }) + .filter((key): key is string => typeof key === 'string'), + ); + const replacesEntry = (existing: unknown): boolean => { + if (existing == null || typeof existing !== 'object') return false; + const entry = existing as { + file_id?: unknown; + filepath?: unknown; + toolCallId?: unknown; + agentId?: unknown; + stepId?: unknown; + }; + const key = entry.file_id ?? entry.filepath; + if (typeof key !== 'string' || !attachmentKeys.has(key)) return false; + if (entry.toolCallId !== toolCallId) return false; + const entryAgent = entry.agentId ?? null; + if (agentId != null && entryAgent !== null && entryAgent !== agentId) return false; + const entryStep = entry.stepId ?? null; + if (stepId != null && entryStep !== null && entryStep !== stepId) return false; + return true; + }; + /** The old pipeline replaced-and-appended `attachments` in one atomic + * write; classic `$pull` and `$push` conflict on one field and splitting + * them opens a window where a crash strips attachments and concurrent + * re-applications duplicate them. Instead: read the array, merge it here + * with the exact semantics the old `$filter`/`$concatArrays` had, and + * write everything back in ONE update fenced on the array being unchanged + * — the same guarded full-array compare-and-swap this file already uses + * for `subagentTask.controlReceipts`. A lost fence means a concurrent + * writer advanced the array; re-reading converges to exactly one copy. */ + for (let attempt = 0; attempt < ATTACHMENT_MERGE_CAS_ATTEMPTS; attempt += 1) { + const row = await Message.findOne(messageFilter) + .select({ _id: 1, attachments: 1 }) + .lean<{ _id: Types.ObjectId; attachments?: unknown[] } | null>(); + if (row == null) { + return { matched: false, unfinished: false }; + } + const prior = Array.isArray(row.attachments) ? row.attachments : []; + const merged = [...prior.filter((entry) => !replacesEntry(entry)), ...attachments]; + const result = await Message.findOneAndUpdate( + { + ...messageFilter, + _id: row._id, + attachments: row.attachments == null ? null : row.attachments, + }, + { + ...settleUpdate, + $set: { ...partPatch, attachments: merged }, + }, + settleOptions, + ).lean<{ unfinished?: boolean } | null>(); + if (result != null) { + return { matched: true, unfinished: result.unfinished === true }; + } + } + /** Losing every fence round means concurrent writers kept advancing the + * array — the document exists and is healthy, so persistence must stay + * retryable ABOVE this bounded loop. The settle retry treats an + * unmatched result as retry-then-heal-later; a throw would land on + * ambiguous-failure handling that can retire the completion outright, + * losing a completed tool result to mere attachment contention. */ + logger.warn( + `[updateToolCallResult] Attachment merge for tool call ${toolCallId} lost ` + + `${ATTACHMENT_MERGE_CAS_ATTEMPTS} fence rounds; leaving the retry to the caller`, + ); + return { matched: false, unfinished: false }; } catch (err) { logger.error('Error updating tool call result:', err); throw err; @@ -1153,6 +1214,9 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa } const MAX_BACKGROUND_TOOL_RESULT_BATCH = 8; + /** Bounds the attachments-merge fence retries. Each lost round means a + * concurrent writer changed the array, so re-reading converges. */ + const ATTACHMENT_MERGE_CAS_ATTEMPTS = 8; function readBackgroundToolResultClaim( row: Pick, @@ -1342,6 +1406,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa candidates.splice(boundedLimit); } const claimedAt = new Date(); + const claimStamp = { kind, claimId, claimedAt }; const updated = await Message.findOneAndUpdate( { user: userId, @@ -1359,7 +1424,13 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa 'tool_call.backgroundTask.resultClaim.kind': kind, 'tool_call.backgroundTask.resultClaim.claimId': claimId, } - : { 'tool_call.backgroundTask.resultClaim': { $exists: false } }), + : /** Missing OR stored null: the in-memory claimable scan, the + * claim arrayFilters, and the settle stamp all treat a null + * claim as unclaimed, and the subfield-preserving settle write + * keeps a persisted null a whole-object rewrite used to drop. + * `$exists: false` here would strand such a part as terminal + * but permanently unclaimable. */ + { 'tool_call.backgroundTask.resultClaim': null }), }, }, ...(agentId != null @@ -1393,100 +1464,39 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa } : {}), }, - [ - { - $set: { - content: { - $map: { - input: { $ifNull: ['$content', []] }, - as: 'part', - in: { - $cond: [ - { - $and: [ - { $eq: ['$$part.type', 'tool_call'] }, - { $in: ['$$part.tool_call.backgroundTask.taskId', candidates] }, - { - $in: ['$$part.tool_call.backgroundTask.status', ['completed', 'error']], - }, - ...(kind === 'wakeup' - ? [{ $eq: ['$$part.tool_call.backgroundTask.completionWakeup', true] }] - : []), - ...(agentId != null - ? [ - { - $in: [ - { - $ifNull: [ - { $ifNull: ['$$part.agentId', '$$part.tool_call.agentId'] }, - null, - ], - }, - [null, agentId], - ], - }, - ] - : []), - { - $or: [ - { - $eq: [ - { $ifNull: ['$$part.tool_call.backgroundTask.resultClaim', null] }, - null, - ], - }, - { - $and: [ - { - $eq: ['$$part.tool_call.backgroundTask.resultClaim.kind', kind], - }, - { - $eq: [ - '$$part.tool_call.backgroundTask.resultClaim.claimId', - claimId, - ], - }, - ], - }, - ], - }, - ], - }, - { - $mergeObjects: [ - '$$part', - { - tool_call: { - $mergeObjects: [ - '$$part.tool_call', - { - backgroundTask: { - $mergeObjects: [ - '$$part.tool_call.backgroundTask', - { - resultClaim: { - kind, - claimId, - claimedAt, - }, - }, - ], - }, - }, - ], - }, - }, - ], - }, - '$$part', - ], - }, + /** Stamps the claim onto every part this pass admitted. The filtered + * positional operator selects those parts by predicate, so the write + * touches only them instead of re-emitting the whole content array, and + * needs no read-modify-write. Amazon DocumentDB rejects the + * aggregation-pipeline form this replaces. */ + { $set: { 'content.$[part].tool_call.backgroundTask.resultClaim': claimStamp } }, + { + new: true, + projection: { content: 1 }, + arrayFilters: [ + { + 'part.type': 'tool_call', + 'part.tool_call.backgroundTask.taskId': { $in: candidates }, + 'part.tool_call.backgroundTask.status': { $in: ['completed', 'error'] }, + ...(kind === 'wakeup' + ? { 'part.tool_call.backgroundTask.completionWakeup': true } + : {}), + $and: [ + { + /** Unclaimed, or already held by this exact claimant (replay). */ + $or: [ + { 'part.tool_call.backgroundTask.resultClaim': null }, + { + 'part.tool_call.backgroundTask.resultClaim.kind': kind, + 'part.tool_call.backgroundTask.resultClaim.claimId': claimId, + }, + ], }, - }, + ...(agentId == null ? [] : [agentOwnershipFilter('part.', agentId)]), + ], }, - }, - ], - { new: true, projection: { content: 1 } }, + ], + }, ).lean(); if (updated == null) { return { status: 'not_ready' }; @@ -1520,63 +1530,38 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa return true; } const Message = mongoose.models.Message as Model; + /** Drops the claim stamp from every content part this claimant owns. The + * filtered positional operator addresses those parts by predicate in one + * atomic write, so no read-modify-write and no rewrite of untouched parts; + * Amazon DocumentDB rejects the aggregation-pipeline form this replaces. + * The filter requires a `content` array because the filtered positional + * operator errors on a row without one, where the old pipeline's + * `$ifNull` no-op'd — a legacy row with no array simply has no claims. */ const updated = await Message.findOneAndUpdate( - { user: userId, conversationId, messageId }, - [ - { - $set: { - content: { - $map: { - input: { $ifNull: ['$content', []] }, - as: 'part', - in: { - $cond: [ - { - $and: [ - ...(taskIds == null - ? [] - : [{ $in: ['$$part.tool_call.backgroundTask.taskId', taskIds] }]), - { $eq: ['$$part.tool_call.backgroundTask.resultClaim.kind', kind] }, - { $eq: ['$$part.tool_call.backgroundTask.resultClaim.claimId', claimId] }, - ], - }, - { - $mergeObjects: [ - '$$part', - { - tool_call: { - $mergeObjects: [ - '$$part.tool_call', - { - backgroundTask: { - $arrayToObject: { - $filter: { - input: { - $objectToArray: '$$part.tool_call.backgroundTask', - }, - as: 'field', - cond: { $ne: ['$$field.k', 'resultClaim'] }, - }, - }, - }, - }, - ], - }, - }, - ], - }, - '$$part', - ], - }, - }, - }, + { user: userId, conversationId, messageId, content: { $type: 'array' } }, + { $unset: { 'content.$[part].tool_call.backgroundTask.resultClaim': 1 } }, + { + new: true, + projection: { content: 1 }, + arrayFilters: [ + { + 'part.tool_call.backgroundTask.resultClaim.kind': kind, + 'part.tool_call.backgroundTask.resultClaim.claimId': claimId, + ...(taskIds == null + ? {} + : { 'part.tool_call.backgroundTask.taskId': { $in: taskIds } }), }, - }, - ], - { new: true, projection: { content: 1 } }, + ], + }, ).lean(); if (updated == null) { - return false; + const arraylessRow = await Message.exists({ + user: userId, + conversationId, + messageId, + content: { $not: { $type: 'array' } }, + }); + return arraylessRow != null; } const remaining = parseBackgroundToolResults(updated, { kind, claimId }); return taskIds == null @@ -2170,7 +2155,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa $cond: [ { $eq: [{ $type: '$$receipt.controlId' }, 'string'] }, boundedString('$$receipt.controlId', SUBAGENT_VIEW_CONTROL_STRING_CODE_POINT_LIMIT), - '$$REMOVE', + null, ], }, action: '$$receipt.action', @@ -2182,14 +2167,14 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa $cond: [ { $eq: [{ $type: '$$receipt.reason' }, 'string'] }, boundedString('$$receipt.reason', SUBAGENT_VIEW_CONTROL_STRING_CODE_POINT_LIMIT), - '$$REMOVE', + null, ], }, message: { $cond: [ { $eq: [{ $type: '$$receipt.message' }, 'string'] }, boundedString('$$receipt.message', SUBAGENT_VIEW_CONTROL_STRING_CODE_POINT_LIMIT), - '$$REMOVE', + null, ], }, messageTruncated: { @@ -2302,7 +2287,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa ], }, }, - '$$REMOVE', + null, ], }; const boundedActivityContent = { @@ -2434,9 +2419,14 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa cond: { $ne: ['$$activity', null] }, }, }; - type ActivitySourceProjection = Pick< - SubagentThreadViewMessageRecord, - | 'messageId' + type ActivitySourceProjection = WithNullSentinels< + Pick< + SubagentThreadViewMessageRecord, + | 'messageId' + | 'subagentTranscript' + | 'subagentActivityProjectionJson' + | 'subagentActivityProjectionTruncated' + >, | 'subagentTranscript' | 'subagentActivityProjectionJson' | 'subagentActivityProjectionTruncated' @@ -2456,11 +2446,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa error: 1, unfinished: 1, subagentTranscriptProjectionTruncated: { - $cond: [ - { $ne: [{ $type: '$subagentTranscript.messagesJson' }, 'missing'] }, - true, - '$$REMOVE', - ], + $cond: [{ $ne: [{ $type: '$subagentTranscript.messagesJson' }, 'missing'] }, true, null], }, subagentActivity: boundedActivityContent, subagentActivityProjectionTruncated: { @@ -2495,11 +2481,11 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa '$subagentTriggerProjection.expectedActionToolName', SUBAGENT_MESSAGE_ACTIVITY_ID_CODE_POINT_LIMIT, ), - '$$REMOVE', + null, ], }, }, - '$$REMOVE', + null, ], }, }; @@ -2509,23 +2495,17 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa _subagentActivityProjectionSourceBytes: activityProjectionJsonBytes, _subagentActivityProjectionSourceIsString: activityProjectionIsString, }; - const activitySourceProjection = { - _id: 0, - messageId: 1, + const activitySourcePayload = { subagentActivityProjectionJson: { - $cond: [ - activityProjectionAvailable, - '$subagentActivityProjection.activityJson', - '$$REMOVE', - ], + $cond: [activityProjectionAvailable, '$subagentActivityProjection.activityJson', null], }, subagentActivityProjectionTruncated: { - $cond: [activityProjectionAvailable, '$subagentActivityProjection.truncated', '$$REMOVE'], + $cond: [activityProjectionAvailable, '$subagentActivityProjection.truncated', null], }, subagentTranscript: { $cond: [ activityProjectionAvailable, - '$$REMOVE', + null, { taskId: '$subagentTranscript.taskId', mode: '$subagentTranscript.mode', @@ -2534,6 +2514,11 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa ], }, }; + const activitySourceProjection = { + _id: 0, + messageId: 1, + ...activitySourcePayload, + }; const baseMatch = { user: input.user, conversationId: input.conversationId, @@ -2561,7 +2546,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa /** Keep rows as independent MongoDB results. A `$facet` would combine the * complete page into one BSON document and could exceed MongoDB's 16 MiB * document limit before the API applies its smaller public byte budget. */ - const messagesPromise = Message.aggregate([ + const messagesPromise = Message.aggregate([ { $match: pageMatch }, { $sort: { createdAt: -1, _id: -1 } }, { $limit: input.limit }, @@ -2616,18 +2601,24 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa { $limit: SUBAGENT_TRANSCRIPT_PAGE_LIMIT - (input.selectedTaskId == null ? 0 : 1) }, { $project: activitySourceProjection }, ]); + /** One read replaces the unsupported `$facet` while keeping its guarantee: + * the bounded message fields and the source payload for each selected row + * come from the same document version, so a transcript persisted between + * two separate reads can never yield a message without its source. Rows + * stay independent results, so no combined document approaches the 16 MiB + * limit the surrounding reads already avoid. Only operators Amazon + * DocumentDB accepts are used. */ const selectedProjectionPromise = input.selectedTaskId == null - ? Promise.resolve([ - { - selectedMessages: [] as SubagentThreadViewMessageRecord[], - selectedSources: [] as ActivitySourceProjection[], - }, - ]) - : Message.aggregate<{ - selectedMessages: SubagentThreadViewMessageRecord[]; - selectedSources: ActivitySourceProjection[]; - }>([ + ? Promise.resolve({ + selectedMessages: [] as ProjectedSubagentThreadViewMessage[], + selectedSources: [] as ActivitySourceProjection[], + }) + : Message.aggregate< + ProjectedSubagentThreadViewMessage & { + _selectedSource: ActivitySourceProjection | null; + } + >([ { $match: { ...baseMatch, @@ -2637,58 +2628,85 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa }, }, { $limit: 2 }, + { $addFields: sourceMetadataProjection }, { - $facet: { - selectedMessages: [{ $project: boundedMessageProjection }], - selectedSources: [ - { $match: { messageId: `${input.selectedTaskId}:assistant` } }, - { $limit: 1 }, - { $addFields: sourceMetadataProjection }, - { - $match: { - $or: [ + $project: { + ...boundedMessageProjection, + _selectedSource: { + $cond: [ + { + $and: [ + { $eq: ['$messageId', `${input.selectedTaskId}:assistant`] }, { - 'subagentActivityProjection.taskId': input.selectedTaskId, - 'subagentActivityProjection.version': 1, - _subagentActivityProjectionSourceIsString: true, - _subagentActivityProjectionSourceBytes: { - $lte: SUBAGENT_ACTIVITY_PROJECTION_SOURCE_BYTE_LIMIT, - }, - }, - { - 'subagentTranscript.taskId': input.selectedTaskId, - _subagentTranscriptSourceIsString: true, - _subagentTranscriptSourceBytes: { - $lte: SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT, - }, + $or: [ + { + $and: [ + { + $eq: [ + '$subagentActivityProjection.taskId', + input.selectedTaskId, + ], + }, + activityProjectionAvailable, + ], + }, + { + $and: [ + { $eq: ['$subagentTranscript.taskId', input.selectedTaskId] }, + transcriptAvailable, + ], + }, + ], }, ], }, - }, - { $project: activitySourceProjection }, - ], + { messageId: '$messageId', ...activitySourcePayload }, + null, + ], + }, }, }, - ]); - const [messages, recentSources, [selectedProjection]] = await Promise.all([ + ]).then((rows) => { + const selectedSources: ActivitySourceProjection[] = []; + const selectedMessages = rows.map(({ _selectedSource, ...message }) => { + if (_selectedSource != null) { + selectedSources.push(_selectedSource); + } + return message; + }); + return { selectedMessages, selectedSources }; + }); + const [messages, recentSources, selectedProjection] = await Promise.all([ messagesPromise, recentSourcesPromise, selectedProjectionPromise, ]); - if (selectedProjection == null) return []; const sourcesByMessageId = new Map( [...selectedProjection.selectedSources, ...recentSources].map((record) => [ record.messageId, record, ]), ); - const retainedMessageIds = new Set(messages.map((message) => message.messageId)); - for (const message of selectedProjection.selectedMessages) { - if (retainedMessageIds.has(message.messageId)) continue; - messages.push(message); - retainedMessageIds.add(message.messageId); + /** The page rows come from an independent concurrent read; for the + * selected task, the single-snapshot pair is authoritative — a page + * duplicate can predate the source read and claim a transcript that + * `selectedSources` does not carry. Replace duplicates in place (keeping + * page order) instead of discarding the snapshot version. */ + const selectedByMessageId = new Map( + selectedProjection.selectedMessages.map((message) => [message.messageId, message]), + ); + for (let index = 0; index < messages.length; index += 1) { + const snapshot = selectedByMessageId.get(messages[index].messageId); + if (snapshot != null) { + messages[index] = snapshot; + selectedByMessageId.delete(snapshot.messageId); + } } - return messages.map((message) => { + for (const message of selectedByMessageId.values()) { + messages.push(message); + } + return messages.map((row) => { + const message = pruneProjectedThreadViewMessage(row); const source = sourcesByMessageId.get(message.messageId); if (source == null) return message; const projected = { ...message }; diff --git a/packages/data-schemas/src/methods/triggerDelivery.spec.ts b/packages/data-schemas/src/methods/triggerDelivery.spec.ts index d516f1c4f4..b4a7ded00c 100644 --- a/packages/data-schemas/src/methods/triggerDelivery.spec.ts +++ b/packages/data-schemas/src/methods/triggerDelivery.spec.ts @@ -13,6 +13,8 @@ import { createAgentTriggerDeliveryMethods, recordAgentEventActorReceiptMetric, setAgentEventActorReceiptMetricObserver, + AgentTriggerClaimContentionError, + CLAIM_CAS_MAX_ATTEMPTS, type AgentTriggerDeliveryMethods, } from './triggerDelivery'; import { @@ -241,6 +243,28 @@ describe('agent trigger delivery methods', () => { expect(claims.filter((claim) => claim != null)).toHaveLength(1); }); + it('surfaces exhausted claim contention as a retryable failure, not an empty queue', async () => { + const enqueued = await methods.enqueueAgentTriggerDelivery(enqueueInput()); + const casSpy = jest.spyOn(Delivery.collection, 'findOneAndUpdate').mockResolvedValue(null); + try { + await expect( + methods.claimNextAgentTriggerDelivery({ + workerId: 'worker-contended', + claimToken: 'claim-contended', + now: START, + leaseUntil: new Date(START.getTime() + 60_000), + }), + ).rejects.toBeInstanceOf(AgentTriggerClaimContentionError); + expect(casSpy).toHaveBeenCalledTimes(CLAIM_CAS_MAX_ATTEMPTS); + } finally { + casSpy.mockRestore(); + } + /** The losing worker owns nothing: the delivery stays claimable. */ + await expect(Delivery.findById(enqueued.delivery.id).lean()).resolves.toMatchObject({ + status: 'pending', + }); + }); + it('keeps capability work visible but nonclaimable to pre-capability consumers', async () => { const user = new mongoose.Types.ObjectId(); const queued = await methods.enqueueAgentTriggerDelivery( diff --git a/packages/data-schemas/src/methods/triggerDelivery.ts b/packages/data-schemas/src/methods/triggerDelivery.ts index 49b9aa8699..033a74ff60 100644 --- a/packages/data-schemas/src/methods/triggerDelivery.ts +++ b/packages/data-schemas/src/methods/triggerDelivery.ts @@ -1,5 +1,5 @@ import { createHash } from 'node:crypto'; -import type { Model, Types } from 'mongoose'; +import type { FilterQuery, Model, Types } from 'mongoose'; import type { AgentEventActorDetachedAction, AgentTriggerDeliveryClaim, @@ -34,6 +34,12 @@ const MAX_PURGE_RECOVERY_LIMIT = 200; const HISTORY_LIMIT = 64; const MAX_BATCH_SIZE = 8; const MAX_BATCH_BYTES = 512 * 1024; +/** Bounds the claim's compare-and-swap retries. Each lost round means another + * worker claimed the row this one read, so the queue is making progress. */ +export const CLAIM_CAS_MAX_ATTEMPTS = 16; +/** Candidates fetched per claim read; losing claimers advance through the + * batch instead of re-reading the same head-of-queue row. */ +const CLAIM_CANDIDATE_BATCH = 8; /** Capability work is inert to legacy claimers while preserving their lane * behavior: publishing is `staging`; queued work is `leased` without a lease * owner/deadline; execution adds a private lease; dead work is terminal. */ @@ -86,6 +92,13 @@ export function recordAgentEventActorReceiptMetric(metric: AgentEventActorReceip } } +export class AgentTriggerClaimContentionError extends Error { + constructor(attempts: number) { + super(`Agent trigger claim lost ${attempts} compare-and-swap rounds to concurrent workers`); + this.name = 'AgentTriggerClaimContentionError'; + } +} + export class AgentTriggerDeliveryConflictError extends Error { constructor(deliveryKey: string) { super(`Agent trigger delivery key ${deliveryKey} was reused with different content`); @@ -1266,106 +1279,115 @@ export function createAgentTriggerDeliveryMethods( availableAt: { $lte: input.now }, }, ]; - const claimed = await Delivery() - .findOneAndUpdate( + const claimableFilter: FilterQuery = { + $or: [ { - $or: [ - { - requiredWorkerCapability: { $exists: false }, - status: 'pending', - availableAt: { $lte: input.now }, - }, - { - requiredWorkerCapability: { $exists: false }, - status: 'leased', - leaseUntil: { $lte: expiredBefore }, - }, - ...capabilityConditions, - ], + requiredWorkerCapability: { $exists: false }, + status: 'pending', + availableAt: { $lte: input.now }, }, - [ - { - $set: { - status: { - $cond: [ - { $eq: [{ $type: '$capabilityStatus' }, 'missing'] }, - { - $cond: [ - { $in: ['$status', ['capability_pending', 'capability_leased']] }, - 'capability_leased', - 'leased', - ], - }, - 'leased', - ], - }, - claimAvailableAt: { - $cond: [ - { - $and: [ - { $eq: ['$status', 'capability_pending'] }, - { $eq: ['$capabilityStatus', 'dead'] }, - ], - }, - '$availableAt', - '$claimAvailableAt', - ], - }, - capabilityStatus: { - $cond: [{ $eq: [{ $type: '$capabilityStatus' }, 'missing'] }, '$$REMOVE', 'leased'], - }, - capabilityLeaseBy: { - $cond: [ - { $eq: [{ $type: '$capabilityStatus' }, 'missing'] }, - '$$REMOVE', - input.workerId, - ], - }, - capabilityLeaseUntil: { - $cond: [ - { $eq: [{ $type: '$capabilityStatus' }, 'missing'] }, - '$$REMOVE', - input.leaseUntil, - ], - }, - capabilityClaimToken: { - $cond: [ - { $eq: [{ $type: '$capabilityStatus' }, 'missing'] }, - '$$REMOVE', - input.claimToken, - ], - }, - leaseBy: { - $cond: [ - { $eq: [{ $type: '$capabilityStatus' }, 'missing'] }, - input.workerId, - LEGACY_CAPABILITY_SHIELD_OWNER, - ], - }, - leaseUntil: { - $cond: [ - { $eq: [{ $type: '$capabilityStatus' }, 'missing'] }, - input.leaseUntil, - LEGACY_CAPABILITY_SHIELD_AT, - ], - }, - claimToken: { - $cond: [ - { $eq: [{ $type: '$capabilityStatus' }, 'missing'] }, - input.claimToken, - LEGACY_CAPABILITY_SHIELD_OWNER, - ], - }, - settledAt: '$$REMOVE', - expiresAt: '$$REMOVE', - updatedAt: input.now, + { + requiredWorkerCapability: { $exists: false }, + status: 'leased', + leaseUntil: { $lte: expiredBefore }, + }, + ...capabilityConditions, + ], + }; + /** The claim's written values depend on the row's own lifecycle, which a + * classic update cannot branch on and Amazon DocumentDB will not accept as + * an aggregation pipeline. Read the oldest claimable candidates under the + * same sort, then claim one by pinning its exact `_id` and lifecycle state + * with a compare-and-swap. Ownership is still granted only by the fenced + * write, so concurrent workers cannot both win. Reading a BATCH is what + * keeps concurrent claimers from convoying on one head-of-queue row: a + * worker that loses a candidate moves to the next from the same read + * instead of re-reading, so a fleet of claimers spreads across the batch + * while each still prefers the oldest work. An empty queue costs the one + * read it costs today; a successful claim costs one additional round trip. */ + let casAttempts = 0; + while (casAttempts < CLAIM_CAS_MAX_ATTEMPTS) { + const candidates = await Delivery() + .find(claimableFilter) + .read('primary') + .sort({ claimAvailableAt: 1, availableAt: 1, createdAt: 1, _id: 1 }) + .limit(CLAIM_CANDIDATE_BATCH) + .select('_id status capabilityStatus availableAt') + .lean< + Array> + >(); + if (candidates.length === 0) { + return null; + } + for (const candidate of candidates) { + if (casAttempts >= CLAIM_CAS_MAX_ATTEMPTS) { + break; + } + casAttempts += 1; + const hasCapabilityLifecycle = candidate.capabilityStatus !== undefined; + const adoptsDeadCapability = + candidate.status === 'capability_pending' && candidate.capabilityStatus === 'dead'; + const claimedStatus: IAgentTriggerDelivery['status'] = + !hasCapabilityLifecycle && + (candidate.status === 'capability_pending' || candidate.status === 'capability_leased') + ? 'capability_leased' + : 'leased'; + const claimed = await Delivery() + .findOneAndUpdate( + { + ...claimableFilter, + _id: candidate._id, + status: candidate.status, + capabilityStatus: hasCapabilityLifecycle + ? candidate.capabilityStatus + : { $exists: false }, + /** The adoption write copies the read-time `availableAt` into the + * ordering key, so the fence must pin it: a concurrent requeue + * that bumps it (without changing the pinned lifecycle pair) + * would otherwise get a stale ordering stamp. */ + ...(adoptsDeadCapability ? { availableAt: candidate.availableAt } : {}), }, - }, - ], - { new: true, sort: { claimAvailableAt: 1, availableAt: 1, createdAt: 1, _id: 1 } }, - ) - .lean(); - return requireClaim(claimed); + { + $set: { + status: claimedStatus, + ...(adoptsDeadCapability ? { claimAvailableAt: candidate.availableAt } : {}), + ...(hasCapabilityLifecycle + ? { + capabilityStatus: 'leased', + capabilityLeaseBy: input.workerId, + capabilityLeaseUntil: input.leaseUntil, + capabilityClaimToken: input.claimToken, + } + : {}), + leaseBy: hasCapabilityLifecycle ? LEGACY_CAPABILITY_SHIELD_OWNER : input.workerId, + leaseUntil: hasCapabilityLifecycle ? LEGACY_CAPABILITY_SHIELD_AT : input.leaseUntil, + claimToken: hasCapabilityLifecycle + ? LEGACY_CAPABILITY_SHIELD_OWNER + : input.claimToken, + updatedAt: input.now, + }, + $unset: { + settledAt: 1, + expiresAt: 1, + ...(hasCapabilityLifecycle + ? {} + : { capabilityLeaseBy: 1, capabilityLeaseUntil: 1, capabilityClaimToken: 1 }), + }, + }, + { new: true, timestamps: false }, + ) + .lean(); + if (claimed != null) { + return requireClaim(claimed); + } + } + } + /** Exhausting the budget means every attempted row was taken by another + * worker — heavy contention, not an empty queue. `null` is the engine's + * confirmed-empty signal and advances its idle backoff, which would leave + * queued work waiting; the typed throw lands on the engine's retryable + * claim-failure path, which holds the polling cadence instead. */ + throw new AgentTriggerClaimContentionError(CLAIM_CAS_MAX_ATTEMPTS); } async function findEarlierAgentTriggerDelivery( @@ -2053,19 +2075,11 @@ export function createAgentTriggerDeliveryMethods( capabilityStatus: { $ne: 'dead' }, settledAt: { $exists: false }, }, - [ - { - $set: { - producerLeaseUntil: { - $cond: [ - { $gt: ['$producerLeaseUntil', input.leaseUntil] }, - '$producerLeaseUntil', - input.leaseUntil, - ], - }, - }, - }, - ], + /** `$max` is the classic-operator form of "never shorten a newer lease": + * it writes only when the incoming deadline is later, and sets the field + * when it is missing. Aggregation-pipeline updates are unsupported on + * Amazon DocumentDB (`misc/documentdb/documentdb-compat.md`). */ + { $max: { producerLeaseUntil: input.leaseUntil } }, { timestamps: false }, ); return renewed.matchedCount === 1; diff --git a/packages/data-schemas/src/utils/transactions.ts b/packages/data-schemas/src/utils/transactions.ts index 799a74e936..bc1ed55835 100644 --- a/packages/data-schemas/src/utils/transactions.ts +++ b/packages/data-schemas/src/utils/transactions.ts @@ -2,21 +2,17 @@ import logger from '~/config/winston'; export const CANCEL_RATE = 1.15; -/** - * Checks if the connected MongoDB deployment supports transactions - * This requires a MongoDB replica set configuration - * - * @returns True if transactions are supported, false otherwise - */ -export const supportsTransactions = async ( - mongoose: typeof import('mongoose'), -): Promise => { +const TRANSACTION_PROBE_COLLECTION = '__transaction_test__'; + +async function probeTransaction(mongoose: typeof import('mongoose')): Promise { try { const session = await mongoose.startSession(); try { session.startTransaction(); - await mongoose.connection.db?.collection('__transaction_test__').findOne({}, { session }); + await mongoose.connection.db + ?.collection(TRANSACTION_PROBE_COLLECTION) + .findOne({}, { session }); await session.commitTransaction(); logger.debug('MongoDB transactions are supported'); @@ -43,12 +39,54 @@ export const supportsTransactions = async ( ); return false; } +} + +/** + * Checks if the connected MongoDB deployment supports transactions + * This requires a MongoDB replica set configuration + * + * Amazon DocumentDB rejects a transaction that touches a collection which does + * not exist ("Feature not supported: non-existent collection in transaction"), + * so a failed first probe materializes the canary and probes once more — + * otherwise engines that fully support transactions would report unsupported + * and silently drop every caller to the non-transactional path. MongoDB allows + * transactional reads of missing collections, so deployments that pass the + * first probe never pay the create and never gain the canary collection. A + * failed create is logged rather than swallowed: on a hardened role it is the + * only visible explanation for a wrong "unsupported" verdict. + * + * @returns True if transactions are supported, false otherwise + */ +export const supportsTransactions = async ( + mongoose: typeof import('mongoose'), +): Promise => { + if (await probeTransaction(mongoose)) { + return true; + } + try { + await mongoose.connection.db?.createCollection(TRANSACTION_PROBE_COLLECTION); + } catch (error) { + const message = (error as Error)?.message ?? ''; + if (!/already exists|NamespaceExists/i.test(message)) { + logger.warn( + `[supportsTransactions] Could not create the "${TRANSACTION_PROBE_COLLECTION}" probe collection; ` + + 'if this deployment is Amazon DocumentDB, the transaction-support verdict below may be a false negative:', + error, + ); + } + } + return probeTransaction(mongoose); }; +let inflightProbe: Promise | null = null; + /** * Gets whether the current MongoDB deployment supports transactions * Caches the result for performance * + * Concurrent first callers share one in-flight probe instead of each paying + * the session/transaction round trips before the caller-side cache fills. + * * @returns True if transactions are supported, false otherwise */ export const getTransactionSupport = async ( @@ -58,5 +96,10 @@ export const getTransactionSupport = async ( if (transactionSupportCache !== null) { return transactionSupportCache; } - return await supportsTransactions(mongoose); + if (inflightProbe == null) { + inflightProbe = supportsTransactions(mongoose).finally(() => { + inflightProbe = null; + }); + } + return await inflightProbe; };