From edf025dda9dea680814b608c4df05f26a156f403 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:02:10 +0000 Subject: [PATCH 01/13] feat(cloud-agent-sdk): require a client message id on cloud agent sends The Cloud Agent treats a message id as the durable user-message identity and replays a repeated id instead of admitting a second turn. The cloud API binding made the id optional, so a caller could leave it out and lose replay on a retry. CloudAgentApi.send now requires the id, and the cloud transport refuses a send without one before the request leaves the client. --- .../src/cloud-agent-transport.test.ts | 5 +++++ .../cloud-agent-sdk/src/cloud-agent-transport.ts | 15 +++++++++++---- .../cloud-agent-sdk/src/session-transport.test.ts | 4 ++++ packages/cloud-agent-sdk/src/transport.ts | 7 ++++++- 4 files changed, 26 insertions(+), 5 deletions(-) diff --git a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts index ce05f428379..da9a7e91057 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts @@ -886,6 +886,7 @@ describe('CloudAgentTransport command delegation', () => { const { transport } = createTransportWithSinks(undefined, undefined, api); await transport.send!({ + messageId: 'msg_123456789abc123456789AB005', payload: { type: 'prompt', prompt: 'hello', @@ -896,6 +897,7 @@ describe('CloudAgentTransport command delegation', () => { expect(api.send).toHaveBeenCalledWith({ sessionId: 'ses-1', + messageId: 'msg_123456789abc123456789AB005', payload: { type: 'prompt', prompt: 'hello', @@ -913,6 +915,7 @@ describe('CloudAgentTransport command delegation', () => { await expect( transport.send!({ + messageId: 'msg_123456789abc123456789AB006', payload: { type: 'prompt', prompt: 'hello', @@ -935,6 +938,7 @@ describe('CloudAgentTransport command delegation', () => { }; await transport.send!({ + messageId: 'msg_123456789abc123456789AB007', payload: { type: 'prompt', prompt: 'read it', @@ -946,6 +950,7 @@ describe('CloudAgentTransport command delegation', () => { expect(api.send).toHaveBeenCalledWith({ sessionId: 'ses-1', + messageId: 'msg_123456789abc123456789AB007', payload: { type: 'prompt', prompt: 'read it', diff --git a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts index 18c7454c45f..f74252fe1a7 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts @@ -399,16 +399,23 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport closeConnection('destroy'); }, - send: input => - runMutation(() => + send: input => { + if (!input.messageId) { + // A Cloud Agent send without a client id cannot be replayed: the + // server would admit a second turn for a retry of the same message. + throw new Error('Cloud Agent sends require a client message id'); + } + const { messageId } = input; + return runMutation(() => config.api.send({ sessionId: config.sessionId, payload: normalizeCloudAgentPayload(input.payload), - ...(input.messageId ? { messageId: input.messageId } : {}), + messageId, ...(input.attachments ? { attachments: input.attachments } : {}), ...(input.images ? { images: input.images } : {}), }) - ), + ); + }, interrupt: () => runMutation(() => config.api.interrupt({ sessionId: config.sessionId })), dropQueuedMessage: async messageId => { const cancelQueuedMessage = config.api.cancelQueuedMessage; diff --git a/packages/cloud-agent-sdk/src/session-transport.test.ts b/packages/cloud-agent-sdk/src/session-transport.test.ts index 0b69eb6b216..f8725e092f3 100644 --- a/packages/cloud-agent-sdk/src/session-transport.test.ts +++ b/packages/cloud-agent-sdk/src/session-transport.test.ts @@ -459,6 +459,7 @@ describe('session transport delegation (cloud agent)', () => { await connectSession(session); await session.send({ + messageId: 'msg_123456789abc123456789AB001', payload: { type: 'prompt', prompt: 'hello', @@ -470,6 +471,7 @@ describe('session transport delegation (cloud agent)', () => { expect(api.send).toHaveBeenCalledTimes(1); expect(api.send).toHaveBeenCalledWith({ sessionId: cloudAgentSessionId, + messageId: 'msg_123456789abc123456789AB001', payload: { type: 'prompt', prompt: 'hello', mode: 'auto', model: 'test/model-1' }, }); @@ -486,6 +488,7 @@ describe('session transport delegation (cloud agent)', () => { await connectSession(session); await session.send({ + messageId: 'msg_123456789abc123456789AB002', payload: { type: 'prompt', prompt: 'hello', @@ -497,6 +500,7 @@ describe('session transport delegation (cloud agent)', () => { expect(api.send).toHaveBeenCalledWith({ sessionId: cloudAgentSessionId, + messageId: 'msg_123456789abc123456789AB002', payload: { type: 'prompt', prompt: 'hello', mode: 'auto', model: 'test/model-1' }, attachments, }); diff --git a/packages/cloud-agent-sdk/src/transport.ts b/packages/cloud-agent-sdk/src/transport.ts index 8bd58b13a70..f44538f4e99 100644 --- a/packages/cloud-agent-sdk/src/transport.ts +++ b/packages/cloud-agent-sdk/src/transport.ts @@ -190,7 +190,12 @@ type CloudAgentApi = { send: (payload: { sessionId: CloudAgentSessionId; payload: CloudAgentSendPayload; - messageId?: string; + /** + * Durable user-message identity minted by the caller. The Cloud Agent + * replays a repeated id instead of admitting a second turn, so it must be + * the same value on a retry of the same turn. + */ + messageId: string; attachments?: CloudAgentAttachments; images?: Images; }) => Promise; From 2cdc7c027c2c477f1ae8d677e9d8e9b11bcd6449 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:02:21 +0000 Subject: [PATCH 02/13] fix(web): link message attachments before dispatch and report released rows linkPendingUploads now reports the requested keys whose row is already 'reaped', and the new linkAttachmentsForDispatch helper links one message's attachment rows before a dispatch and refuses the dispatch when a row was already released, with an actionable error instead of a turn whose attachment object is gone. The abandoned-upload reaper re-checks the lease in the claim statement that moves a row to 'reaped', so a row refreshed by a retry presign after the batch was picked is skipped. The claim comments now state the invariant they hold. --- .../r2/cloud-agent-pending-uploads.test.ts | 161 +++++++++++++++++- .../src/lib/r2/cloud-agent-pending-uploads.ts | 117 ++++++++++--- 2 files changed, 247 insertions(+), 31 deletions(-) diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts index 04cba1830db..ccc60a6cbc4 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts @@ -1,6 +1,7 @@ import { afterEach, beforeAll, describe, expect, it, jest } from '@jest/globals'; import type { admitPendingUpload as AdmitPendingUpload, + linkAttachmentsForDispatch as LinkAttachmentsForDispatch, linkPendingUploads as LinkPendingUploads, reapAbandonedUploads as ReapAbandonedUploads, releasePendingUploads as ReleasePendingUploads, @@ -29,6 +30,7 @@ import { db } from '@kilocode/web-shared/lib/drizzle'; import { cloud_agent_pending_uploads } from '@kilocode/db/schema'; let admitPendingUpload: typeof AdmitPendingUpload; +let linkAttachmentsForDispatch: typeof LinkAttachmentsForDispatch; let linkPendingUploads: typeof LinkPendingUploads; let reapAbandonedUploads: typeof ReapAbandonedUploads; let releasePendingUploads: typeof ReleasePendingUploads; @@ -72,6 +74,7 @@ describe('cloud-agent pending-upload ledger', () => { beforeAll(async () => { const pendingUploadsModule = await import('./cloud-agent-pending-uploads'); admitPendingUpload = pendingUploadsModule.admitPendingUpload; + linkAttachmentsForDispatch = pendingUploadsModule.linkAttachmentsForDispatch; linkPendingUploads = pendingUploadsModule.linkPendingUploads; reapAbandonedUploads = pendingUploadsModule.reapAbandonedUploads; releasePendingUploads = pendingUploadsModule.releasePendingUploads; @@ -553,7 +556,7 @@ describe('cloud-agent pending-upload ledger', () => { expect(byId.get(otherUserRow)).toBe('pending'); }); - it('claims the row before deleting its R2 object so a concurrent link cannot lose it', async () => { + it('moves a released row out of pending before deleting its R2 object', async () => { const kiloUserId = `ca-user-${crypto.randomUUID()}`; const messageUuid = crypto.randomUUID(); const attachment = crypto.randomUUID(); @@ -614,7 +617,7 @@ describe('cloud-agent pending-upload ledger', () => { ).toBe('pending'); }); - it('claims an expired row before deleting its object so a late link cannot lose it', async () => { + it('moves an expired row out of pending before deleting its object', async () => { const kiloUserId = `ca-user-${crypto.randomUUID()}`; const messageUuid = crypto.randomUUID(); const attachment = crypto.randomUUID(); @@ -628,9 +631,9 @@ describe('cloud-agent pending-upload ledger', () => { expiresAt: past, }); - // A row picked by the reaper must already be claimed when DeleteObject - // runs. If it were still 'pending', linkPendingUploads could flip it to - // 'linked' and the sent attachment would lose its object. + // A row picked by the reaper must be out of 'pending' before DeleteObject + // runs, so a caller that links afterwards reads 'reaped' and refuses the + // dispatch instead of promising an object that is gone. const statusesAtDelete: (string | undefined)[] = []; mockSend.mockImplementation(async () => { const rows = await db @@ -645,4 +648,152 @@ describe('cloud-agent pending-upload ledger', () => { expect(statusesAtDelete).toEqual(['reaped']); }); + + it('keeps a linked attachment object when the reaper runs afterwards', async () => { + const kiloUserId = `ca-user-${crypto.randomUUID()}`; + const messageUuid = crypto.randomUUID(); + const attachment = crypto.randomUUID(); + const objectKey = buildObjectKey(kiloUserId, messageUuid, attachment); + const past = new Date(Date.now() - 60_000).toISOString(); + mockSend.mockResolvedValue({}); + + const row = await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: attachment, + expiresAt: past, + }); + + const linked = await linkPendingUploads(kiloUserId, messageUuid, [objectKey]); + expect(linked).toEqual({ reapedObjectKeys: [] }); + + const summary = await reapAbandonedUploads(); + + expect(summary.reaped).toBe(0); + expect(mockSend).not.toHaveBeenCalled(); + const [after] = await db + .select({ status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.id, row)); + expect(after?.status).toBe('linked'); + }); + + it('reports the attachment row a reaper claim took first and does not link it', async () => { + const kiloUserId = `ca-user-${crypto.randomUUID()}`; + const messageUuid = crypto.randomUUID(); + const attachment = crypto.randomUUID(); + const objectKey = buildObjectKey(kiloUserId, messageUuid, attachment); + const past = new Date(Date.now() - 60_000).toISOString(); + mockSend.mockResolvedValue({}); + + const row = await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: attachment, + expiresAt: past, + }); + + await reapAbandonedUploads(); + + const linked = await linkPendingUploads(kiloUserId, messageUuid, [objectKey]); + expect(linked).toEqual({ reapedObjectKeys: [objectKey] }); + const [after] = await db + .select({ status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.id, row)); + expect(after?.status).toBe('reaped'); + }); + + it('does not reap a lease a retry refreshes before the claim runs', async () => { + const kiloUserId = `ca-user-${crypto.randomUUID()}`; + const messageUuid = crypto.randomUUID(); + const attachment = crypto.randomUUID(); + const past = new Date(Date.now() - 60_000).toISOString(); + const future = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString(); + mockSend.mockResolvedValue({}); + + const row = await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: attachment, + expiresAt: past, + }); + + const hold = Promise.withResolvers(); + const refreshed = Promise.withResolvers(); + + // The retry presign extends the lease while the reaper is already running. + // Whichever statement reaches the row first, the lease re-check in the + // claim keeps the object; only an expired row may be deleted. + const refreshing = db.transaction(async tx => { + await tx + .update(cloud_agent_pending_uploads) + .set({ expires_at: future }) + .where(eq(cloud_agent_pending_uploads.id, row)); + refreshed.resolve(); + await hold.promise; + }); + + await refreshed.promise; + const reaping = reapAbandonedUploads(); + hold.resolve(); + await refreshing; + + expect(await reaping).toEqual({ reaped: 0 }); + expect(mockSend).not.toHaveBeenCalled(); + const [after] = await db + .select({ status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.id, row)); + expect(after?.status).toBe('pending'); + }); + + it('links a message attachment row before dispatch and tolerates a key with no row', async () => { + const kiloUserId = `ca-user-${crypto.randomUUID()}`; + const messageUuid = crypto.randomUUID(); + const attachment = crypto.randomUUID(); + const legacyFile = `${crypto.randomUUID()}.pdf`; + const objectKey = buildObjectKey(kiloUserId, messageUuid, attachment); + const future = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString(); + + const row = await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: attachment, + expiresAt: future, + }); + + await linkAttachmentsForDispatch(kiloUserId, { + path: messageUuid, + files: [`${attachment}.bin`, legacyFile], + }); + + const [after] = await db + .select({ status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.id, row)); + expect(after?.status).toBe('linked'); + }); + + it('refuses a dispatch whose attachment row was already reaped', async () => { + const kiloUserId = `ca-user-${crypto.randomUUID()}`; + const messageUuid = crypto.randomUUID(); + const attachment = crypto.randomUUID(); + const past = new Date(Date.now() - 60_000).toISOString(); + + await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: attachment, + expiresAt: past, + status: 'reaped', + }); + + await expect( + linkAttachmentsForDispatch(kiloUserId, { + path: messageUuid, + files: [`${attachment}.bin`], + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + }); }); diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts index f706f093a47..cad3169acbf 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts @@ -37,9 +37,9 @@ export type AdmitPendingUploadParams = { * row for objectKey in any other state keeps its key, and the upload gets freshObjectKey: * a release or reaper claim sets 'reaped' before its DeleteObject runs, so a * PUT to the old key could land before that delete, and a 'linked' key holds - * a sent attachment. The row is later flipped to 'linked' by - * linkPendingUploads when the message is actually sent, or to 'reaped' once - * its object is deleted. + * an attachment a dispatched message refers to. The row is later flipped to + * 'linked' by linkPendingUploads before the message is dispatched, or to + * 'reaped' once its object is deleted. */ export async function admitPendingUpload({ kiloUserId, @@ -133,31 +133,89 @@ export async function admitPendingUpload({ }); } +export type LinkPendingUploadsResult = { + /** + * Requested keys whose row is 'reaped', so the object behind them is gone or + * being deleted. A message referencing one of these must not be dispatched. + */ + reapedObjectKeys: string[]; +}; + /** - * Flip pending ledger rows to 'linked' for a user's message in one - * transaction. Rows whose object keys are not in the supplied set stay - * 'pending', so a partial key set leaves unmatched uploads to be reaped. + * Flip pending ledger rows to 'linked' for a user's message in one guarded + * update. Rows whose object keys are not in the supplied set stay 'pending', + * so a partial key set leaves unmatched uploads to be reaped. A row already + * 'linked' belongs to an earlier dispatch of the same message and is reported + * as neither linked now nor reaped. */ export async function linkPendingUploads( kiloUserId: string, messageUuid: string, objectKeys: string[] +): Promise { + if (objectKeys.length === 0) return { reapedObjectKeys: [] }; + + await db + .update(cloud_agent_pending_uploads) + .set({ status: 'linked' }) + .where( + and( + eq(cloud_agent_pending_uploads.kilo_user_id, kiloUserId), + eq(cloud_agent_pending_uploads.message_uuid, messageUuid), + eq(cloud_agent_pending_uploads.status, 'pending'), + inArray(cloud_agent_pending_uploads.object_key, objectKeys) + ) + ); + + // Read back what the guarded update could not link. A row that a release or + // a reaper claim already took out of 'pending' is 'reaped', and its object is + // going away, so the caller must not dispatch a message that references it. + const rows = await db + .select({ + object_key: cloud_agent_pending_uploads.object_key, + status: cloud_agent_pending_uploads.status, + }) + .from(cloud_agent_pending_uploads) + .where( + and( + eq(cloud_agent_pending_uploads.kilo_user_id, kiloUserId), + eq(cloud_agent_pending_uploads.message_uuid, messageUuid), + eq(cloud_agent_pending_uploads.status, 'reaped'), + inArray(cloud_agent_pending_uploads.object_key, objectKeys) + ) + ); + + return { reapedObjectKeys: rows.map(row => row.object_key) }; +} + +/** Message for a send or prepare whose attachment object is already gone. */ +export const CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE = + 'Attachment upload expired. Add the file again to send it.'; + +/** + * Link the ledger rows of one message's attachments before its dispatch, so + * the objects survive the reaper, and refuse the dispatch when a row was + * already reaped. Keys with no ledger row belong to an uploader that predates + * the ledger and are left alone. + */ +export async function linkAttachmentsForDispatch( + kiloUserId: string, + attachments: { path: string; files: string[] } | undefined ): Promise { - if (objectKeys.length === 0) return; + if (!attachments || attachments.files.length === 0) return; - await db.transaction(async tx => { - await tx - .update(cloud_agent_pending_uploads) - .set({ status: 'linked' }) - .where( - and( - eq(cloud_agent_pending_uploads.kilo_user_id, kiloUserId), - eq(cloud_agent_pending_uploads.message_uuid, messageUuid), - eq(cloud_agent_pending_uploads.status, 'pending'), - inArray(cloud_agent_pending_uploads.object_key, objectKeys) - ) - ); - }); + const { reapedObjectKeys } = await linkPendingUploads( + kiloUserId, + attachments.path, + attachments.files.map(file => `${kiloUserId}/cloud-agent/${attachments.path}/${file}`) + ); + + if (reapedObjectKeys.length > 0) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE, + }); + } } export type ReapAbandonedUploadsSummary = { @@ -193,10 +251,13 @@ async function deletePendingUploadObjects( /** * Claim these rows in one statement, delete their objects, then put back to * 'pending' any row whose delete failed so a later reaper run retries it. - * Claiming first is what keeps a concurrent linkPendingUploads from losing a - * sent attachment: after the claim the row is no longer 'pending', so the link - * matches nothing and the object it needs is never deleted. Returns the number - * of objects actually deleted. + * Claiming in the same guarded update that selects the rows is what makes the + * ledger the single authority for a delete: a row that a link or a release has + * already moved out of the predicate is never claimed, and the flip to + * 'reaped' commits before DeleteObject runs. Every dispatch links a message's + * attachment rows before the worker can accept the turn, so a row that is + * still claimable belongs to no accepted message. Returns the number of objects + * actually deleted. */ async function claimAndDeleteUploadObjects(claim: SQL | undefined): Promise { const claimed = await db @@ -263,8 +324,9 @@ export async function purgeUserPendingUploads(kiloUserId: string): Promise /** * Delete the private R2 objects behind abandoned pending uploads whose * 24-hour lease has lapsed, one bounded batch per run. The batch is picked in - * the same statement that claims it, and the claim still requires 'pending', - * so a row linked since the pick is skipped rather than stripped of its object. + * the same statement that claims it, and the claim re-checks 'pending' and the + * lease, so a row linked or refreshed since the pick is skipped rather than + * stripped of its object. */ export async function reapAbandonedUploads(): Promise { const expiredBatch = db @@ -282,6 +344,9 @@ export async function reapAbandonedUploads(): Promise Date: Thu, 22 Oct 2026 12:02:33 +0000 Subject: [PATCH 03/13] fix(web): link attachments before dispatch and require the client message id The personal and organization prepareSession and sendMessage handlers linked a message's attachment rows after the Worker had answered, inside the same try whose catch rethrows the Worker error. A ledger write failure therefore reported an accepted send as failed, and a retry of the reported failure had no stable id to replay against. Both handlers now link the rows before the dispatch, so no bookkeeping is left after the Worker accepts the turn, and they forward the client's message id instead of minting one. baseSendMessageNextSchema requires messageId, so a client that omits it is rejected before any Worker call. Comments in the composer hook describe the dispatch-time link they rely on. --- .../extension-agent-session-manager.test.ts | 2 + .../src/shared/extension-trpc-client.test.ts | 3 + .../hooks/useCloudAgentAttachmentUpload.ts | 8 +- .../routers/cloud-agent-next-router.test.ts | 157 +++++++++++++++++- .../src/routers/cloud-agent-next-router.ts | 60 +++---- .../src/routers/cloud-agent-next-schemas.ts | 8 +- ...ganization-cloud-agent-next-router.test.ts | 41 ++++- .../organization-cloud-agent-next-router.ts | 61 +++---- 8 files changed, 258 insertions(+), 82 deletions(-) diff --git a/apps/extension/src/shared/extension-agent-session-manager.test.ts b/apps/extension/src/shared/extension-agent-session-manager.test.ts index fa920f2069a..678fbb9897c 100644 --- a/apps/extension/src/shared/extension-agent-session-manager.test.ts +++ b/apps/extension/src/shared/extension-agent-session-manager.test.ts @@ -1254,6 +1254,7 @@ describe('createExtensionAgentSessionManager', () => { const trpc = setup(); await capturedConfig!.api.send({ attachments: { files: [], path: '/f', type: 'file' } as never, + messageId: 'msg_123456789abc123456789ABC04', payload: { mode: 'code', model: 'gpt', prompt: 'hi', type: 'prompt' }, sessionId: CLOUD_AGENT_ID as never, }); @@ -1325,6 +1326,7 @@ describe('createExtensionAgentSessionManager', () => { it('send calls org cloudAgentNext.sendMessage including organizationId', async () => { const trpc = setup(); await capturedConfig!.api.send({ + messageId: 'msg_123456789abc123456789ABC03', payload: { mode: 'code', model: 'gpt', prompt: 'hi', type: 'prompt' }, sessionId: CLOUD_AGENT_ID as never, }); diff --git a/apps/extension/src/shared/extension-trpc-client.test.ts b/apps/extension/src/shared/extension-trpc-client.test.ts index 37e20234391..48213cf0233 100644 --- a/apps/extension/src/shared/extension-trpc-client.test.ts +++ b/apps/extension/src/shared/extension-trpc-client.test.ts @@ -10,14 +10,17 @@ const FAKE_BASE_URL = 'https://api.example.com'; const FAKE_SESSION_ID = 'test-session-id'; const PROMPT_PAYLOAD = { cloudAgentSessionId: FAKE_SESSION_ID, + messageId: 'msg_123456789abc123456789ABCDE', payload: { mode: 'plan' as const, model: 'test-model', prompt: 'hello', type: 'prompt' as const }, }; const PROMPT_FIRST = { cloudAgentSessionId: FAKE_SESSION_ID, + messageId: 'msg_123456789abc123456789ABC01', payload: { mode: 'plan' as const, model: 'test-model', prompt: 'first', type: 'prompt' as const }, }; const PROMPT_SECOND = { cloudAgentSessionId: FAKE_SESSION_ID, + messageId: 'msg_123456789abc123456789ABC02', payload: { mode: 'plan' as const, model: 'test-model', diff --git a/apps/web/src/hooks/useCloudAgentAttachmentUpload.ts b/apps/web/src/hooks/useCloudAgentAttachmentUpload.ts index 83312af1554..5e8414bdee6 100644 --- a/apps/web/src/hooks/useCloudAgentAttachmentUpload.ts +++ b/apps/web/src/hooks/useCloudAgentAttachmentUpload.ts @@ -492,11 +492,9 @@ export function useCloudAgentAttachmentUpload( /** * Build the attachment payload for the caller to submit. The pending-ledger - * rows are NOT linked here: the server links authoritatively only after the - * send or prepare succeeds (both the personal and organization routers call - * linkPendingUploads post-send). Linking client-side before the send would - * leave a failed or abandoned send with 'linked' rows that the reaper never - * deletes. + * rows are linked by the server as part of the send or prepare dispatch, so + * this side must not link them earlier: a row this side linked before a send + * that never happened would keep its object out of the reaper's reach. */ const finalizeAttachments = useCallback(async (): Promise => { return getAttachmentsData(); diff --git a/apps/web/src/routers/cloud-agent-next-router.test.ts b/apps/web/src/routers/cloud-agent-next-router.test.ts index ea83baa15f4..ddc7155851e 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -1,7 +1,12 @@ import { describe, expect, it, jest, beforeAll, beforeEach, afterAll } from '@jest/globals'; import jwt from 'jsonwebtoken'; import { NEXTAUTH_SECRET } from '@kilocode/web-shared/lib/config.server'; -import { cli_sessions_v2, organizations, type User } from '@kilocode/db/schema'; +import { + cli_sessions_v2, + cloud_agent_pending_uploads, + organizations, + type User, +} from '@kilocode/db/schema'; import type { createWorktreeChat as CreateWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; import type * as CloudAgentClientModule from '@/lib/cloud-agent-next/cloud-agent-client'; import type { CloudAgentNextClient } from '@/lib/cloud-agent-next/cloud-agent-client'; @@ -38,6 +43,9 @@ import type * as StreamTicketModule from '@/lib/cloud-agent/stream-ticket'; type AttachmentReference = { path: string; files: string[] }; +/** A valid client message id; the send schema requires one. */ +const SEND_MESSAGE_ID = 'msg_123456789abc123456789ABCDE'; + const mockPrepareSession = jest.fn< (input: { githubRepo?: string; @@ -268,7 +276,7 @@ let createCaller: (ctx: { user: User; headersList?: Headers }) => { sendMessage: (input: { cloudAgentSessionId: string; expectedWorktreeId?: `worktree_${string}`; - messageId?: string; + messageId: string; payload: | { type: 'prompt'; prompt: string; mode: string; model: string } | { type: 'command'; command: string; arguments: string }; @@ -454,6 +462,7 @@ describe('cloudAgentNextRouter worktree changes access', () => { createCaller({ user: owner }).sendMessage({ cloudAgentSessionId: personalSessionId, expectedWorktreeId: worktreeId, + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'a'.repeat(100_001), @@ -576,6 +585,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await expect( caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Read PDF', mode: 'code', model: 'test' }, }) ).rejects.toThrow('Session not found or access denied'); @@ -592,6 +602,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Read PDF', mode: 'code', model: 'test' }, attachments, }); @@ -608,6 +619,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Read image', mode: 'code', model: 'test' }, images, }); @@ -625,6 +637,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Follow up on this', @@ -653,6 +666,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -672,6 +686,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -691,6 +706,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -708,6 +724,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -777,6 +794,141 @@ describe('cloudAgentNextRouter attachment forwarding', () => { }); }); +describe('cloudAgentNextRouter send attachment ledger', () => { + const userId = `oauth/ledger-send-${crypto.randomUUID()}`; + const messageUuid = '12345678-1234-4234-9234-123456789abc'; + const attachmentId = '87654321-4321-4321-8321-cba987654321'; + const file = `${attachmentId}.pdf`; + const attachments = { path: messageUuid, files: [file] }; + const objectKey = `${userId}/cloud-agent/${messageUuid}/${file}`; + const future = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString(); + const past = new Date(Date.now() - 60_000).toISOString(); + + async function insertLedgerRow( + status: 'pending' | 'linked' | 'reaped', + expiresAt = future + ): Promise { + await db.insert(cloud_agent_pending_uploads).values({ + id: crypto.randomUUID(), + kilo_user_id: userId, + object_key: objectKey, + message_uuid: messageUuid, + attachment_id: attachmentId, + byte_size: 42, + status, + expires_at: expiresAt, + }); + } + + async function readStatus(): Promise { + const rows = await db + .select({ status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.object_key, objectKey)); + return rows[0]?.status; + } + + beforeEach(async () => { + jest.clearAllMocks(); + mockSendMessage.mockReset(); + mockSendMessage.mockResolvedValue({ + cloudAgentSessionId: 'agent_123', + status: 'started', + streamUrl: '/stream', + messageId: SEND_MESSAGE_ID, + delivery: 'sent', + }); + mockVerifyUserOwnsSessionV2ByCloudAgentId.mockResolvedValue({ + kiloSessionId: 'ses_12345678901234567890123456', + }); + mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValue({ + isFree: false, + hasUserByokAvailable: false, + }); + mockGetSession.mockResolvedValue({ model: 'kilo/paid-model' }); + await db + .delete(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.object_key, objectKey)); + }); + + afterAll(async () => { + await db + .delete(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.object_key, objectKey)); + }); + + it('refuses a send whose attachment object the ledger already released', async () => { + await insertLedgerRow('reaped', past); + const caller = createCaller({ user: { id: userId, is_admin: false } as User }); + + await expect( + caller.sendMessage({ + cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, + payload: { type: 'prompt', prompt: 'Read the attachment', mode: 'code', model: 'test' }, + attachments, + }) + ).rejects.toMatchObject({ + code: 'BAD_REQUEST', + message: 'Attachment upload expired. Add the file again to send it.', + }); + + expect(mockSendMessage).not.toHaveBeenCalled(); + }); + + it('links the attachment rows before the worker accepts the send', async () => { + await insertLedgerRow('pending'); + const caller = createCaller({ user: { id: userId, is_admin: false } as User }); + + await expect( + caller.sendMessage({ + cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, + payload: { type: 'prompt', prompt: 'Read the attachment', mode: 'code', model: 'test' }, + attachments, + }) + ).resolves.toMatchObject({ messageId: SEND_MESSAGE_ID }); + + expect(mockSendMessage).toHaveBeenCalledWith(expect.objectContaining({ attachments })); + expect(await readStatus()).toBe('linked'); + }); + + it('keeps the attachment rows linked when the worker rejects the send', async () => { + await insertLedgerRow('pending'); + mockSendMessage.mockRejectedValueOnce(new TRPCClientError('worker unavailable')); + const caller = createCaller({ user: { id: userId, is_admin: false } as User }); + + await expect( + caller.sendMessage({ + cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, + payload: { type: 'prompt', prompt: 'Read the attachment', mode: 'code', model: 'test' }, + attachments, + }) + ).rejects.toThrow('worker unavailable'); + + expect(await readStatus()).toBe('linked'); + }); + + it('requires a client message id before calling the worker', async () => { + await insertLedgerRow('pending'); + const caller = createCaller({ user: { id: userId, is_admin: false } as User }); + + await expect( + caller.sendMessage( + // Omit the required message id to prove the input schema rejects it. + { + cloudAgentSessionId: 'agent_123', + payload: { type: 'prompt', prompt: 'No id', mode: 'code', model: 'test' }, + } as never + ) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + + expect(mockSendMessage).not.toHaveBeenCalled(); + expect(await readStatus()).toBe('pending'); + }); +}); + describe('cloudAgentNextRouter.cancelQueuedMessage', () => { beforeEach(() => { jest.clearAllMocks(); @@ -1460,6 +1612,7 @@ describe('cloudAgentNextRouter Worker error codes', () => { mockSendMessage.mockRejectedValueOnce(error); return caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Hello', mode: 'code', model: 'kilo/test-model' }, }); } diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index 52c3abe387c..6a0f2b34522 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -71,7 +71,11 @@ import { generateCloudAgentAttachmentDownloadUrl, generateImageUploadUrl, } from '@/lib/r2/cloud-agent-attachments'; -import { linkPendingUploads, releasePendingUploads } from '@/lib/r2/cloud-agent-pending-uploads'; +import { + linkAttachmentsForDispatch, + linkPendingUploads, + releasePendingUploads, +} from '@/lib/r2/cloud-agent-pending-uploads'; import * as z from 'zod'; import { PLATFORM } from '@/lib/integrations/core/constants'; import { signPreviewTicket, signStreamTicket } from '@/lib/cloud-agent/stream-ticket'; @@ -79,7 +83,6 @@ import { buildPreviewOpenUrl, buildPreviewUrlTemplate } from '@/lib/cloud-agent/ import { db } from '@kilocode/web-shared/lib/drizzle'; import { verifyUserOwnsSessionV2ByCloudAgentId } from '@/lib/cloud-agent/session-ownership'; import { TRPCError } from '@trpc/server'; -import { generateMessageId } from '@kilocode/cloud-agent-sdk/message-id'; import { getBalanceForUser } from '@kilocode/web-shared/lib/user/balance'; import { isMobileClient } from '@kilocode/web-shared/lib/trpc/min-version'; import { buildCloudAgentNextEligibility } from './cloud-agent-next-eligibility'; @@ -238,7 +241,15 @@ export const cloudAgentNextRouter = createTRPCRouter({ } try { - const result = await client.prepareSession({ + // New-session flows (mobile, continue-cloud-create) call prepareSession + // without a follow-up sendMessage, so the attachment rows must be + // linked before prepare can accept the session; otherwise the reaper + // deletes new-session objects. Only ATTACHMENT rows are admitted + // (images are never admitted; only + // generateCloudAgentAttachmentUploadUrl admits). + await linkAttachmentsForDispatch(ctx.user.id, attachments); + + return await client.prepareSession({ ...restInput, ...gitParams, attachments: attachments ?? images, @@ -246,21 +257,6 @@ export const cloudAgentNextRouter = createTRPCRouter({ githubAccessPurpose: 'agent', clientProvenance: isMobileClient(ctx.headersList) ? 'mobile' : 'browser', }); - - // New-session flows (mobile, continue-cloud-create) call prepareSession - // without a follow-up sendMessage, so the server must link the admitted - // attachment rows here too; otherwise the reaper deletes new-session - // objects. Only ATTACHMENT rows are admitted (images are never admitted; - // only generateCloudAgentAttachmentUploadUrl admits). - if (attachments && attachments.files.length > 0) { - await linkPendingUploads( - ctx.user.id, - attachments.path, - attachments.files.map(file => `${ctx.user.id}/cloud-agent/${attachments.path}/${file}`) - ); - } - - return result; } catch (error) { rethrowCloudAgentError(error); } @@ -350,29 +346,21 @@ export const cloudAgentNextRouter = createTRPCRouter({ // for GitHub, GIT_TOKEN_SERVICE for managed GitLab). try { const { attachments, images } = input; - const result = await client.sendMessage({ + // Link the attachment rows before the worker can accept the turn: the + // link is what keeps their objects out of the abandoned-upload reaper, + // and once the worker has answered there is no ledger work left that + // could report an accepted send as failed. Only ATTACHMENT rows are + // admitted (images are never admitted; only + // generateCloudAgentAttachmentUploadUrl admits). + await linkAttachmentsForDispatch(ctx.user.id, attachments); + + return await client.sendMessage({ cloudAgentSessionId: input.cloudAgentSessionId, payload: input.payload, autoCommit: input.autoCommit, attachments: attachments ?? images, - messageId: input.messageId ?? generateMessageId(), + messageId: input.messageId, }); - - // Compatibility: the server links authoritatively so any client whose - // send goes through sendMessage keeps its objects (mobile, old tabs). - // The client does not link pre-send, so a failed or abandoned send - // leaves rows 'pending' for the reaper. Only ATTACHMENT rows are - // admitted (images are never admitted; only - // generateCloudAgentAttachmentUploadUrl admits). - if (attachments && attachments.files.length > 0) { - await linkPendingUploads( - ctx.user.id, - attachments.path, - attachments.files.map(file => `${ctx.user.id}/cloud-agent/${attachments.path}/${file}`) - ); - } - - return result; } catch (error) { rethrowCloudAgentError(error); } diff --git a/apps/web/src/routers/cloud-agent-next-schemas.ts b/apps/web/src/routers/cloud-agent-next-schemas.ts index 5ef0b587bdd..3927c3e7842 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.ts @@ -589,7 +589,13 @@ export const baseSendMessageNextSchema = z payload: sendMessageNextSendPayloadSchema, expectedWorktreeId: cloudAgentWorktreeIdSchema.optional(), autoCommit: z.boolean().optional(), - messageId: messageIdNextSchema.nullish(), + /** + * Durable user-message identity. The control plane replays a repeated id + * instead of admitting a second turn, so a client retry after a failed + * response must carry the id of the turn it is retrying; the server cannot + * invent one on the client's behalf. + */ + messageId: messageIdNextSchema, attachments: cloudAgentAttachmentsSchema.optional(), images: cloudAgentImagesSchema, }) diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts index bcc20302ecf..9650c64acff 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts @@ -35,6 +35,7 @@ import { TRPCClientError } from '@trpc/client'; import { TRPCError } from '@trpc/server'; import { getHTTPStatusCodeFromError } from '@trpc/server/http'; import type { verifyOrgOwnsSessionV2ByCloudAgentId } from '@/lib/cloud-agent/session-ownership'; +import type * as PendingUploadsModule from '@/lib/r2/cloud-agent-pending-uploads'; import type * as StreamTicketModule from '@/lib/cloud-agent/stream-ticket'; import type { organizationPrepareSessionNextSchema, @@ -54,6 +55,9 @@ const ORGANIZATION_ID = '9a283301-b75d-4375-a1ba-e319a02e18b7'; type AttachmentReference = { path: string; files: string[] }; +/** A valid client message id; the send schema requires one. */ +const SEND_MESSAGE_ID = 'msg_123456789abc123456789ABCDE'; + const mockPrepareSession = jest.fn< (input: { githubRepo?: string; @@ -277,6 +281,7 @@ jest.mock('@/lib/r2/cloud-agent-attachments', () => ({ })); jest.mock('@/lib/r2/cloud-agent-pending-uploads', () => ({ + linkAttachmentsForDispatch: jest.fn(async () => {}), linkPendingUploads: jest.fn(), releasePendingUploads: jest.fn(), })); @@ -334,7 +339,7 @@ let createCaller: (ctx: { user: User; headersList?: Headers }) => { organizationId: string; cloudAgentSessionId: string; expectedWorktreeId?: `worktree_${string}`; - messageId?: string; + messageId: string; payload: | { type: 'prompt'; prompt: string; mode: string; model: string } | { type: 'command'; command: string; arguments: string }; @@ -965,6 +970,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Read notes', mode: 'code', model: 'test' }, }) ).rejects.toThrow('Organization does not own this session'); @@ -982,6 +988,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Read notes', mode: 'code', model: 'test' }, attachments, }); @@ -992,6 +999,31 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { ); }); + it('links the attachment rows before dispatching the send to the Worker', async () => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + const attachments = { + path: '12345678-1234-4234-9234-123456789abc', + files: ['87654321-4321-4321-8321-cba987654321.md'], + }; + const pendingUploads = jest.requireMock( + '@/lib/r2/cloud-agent-pending-uploads' + ); + const linkAttachmentsForDispatch = jest.mocked(pendingUploads.linkAttachmentsForDispatch); + + await caller.sendMessage({ + organizationId: ORGANIZATION_ID, + cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, + payload: { type: 'prompt', prompt: 'Read notes', mode: 'code', model: 'test' }, + attachments, + }); + + expect(linkAttachmentsForDispatch).toHaveBeenCalledWith('user-1', attachments); + expect(linkAttachmentsForDispatch.mock.invocationCallOrder[0]).toBeLessThan( + mockSendMessage.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY + ); + }); + it('normalizes legacy image requests to canonical Worker attachments', async () => { const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); const images = { @@ -1002,6 +1034,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Read image', mode: 'code', model: 'test' }, images, }); @@ -1020,6 +1053,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Follow up on this', @@ -1052,6 +1086,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -1075,6 +1110,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -1098,6 +1134,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -1116,6 +1153,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -2136,6 +2174,7 @@ describe('organizationCloudAgentNextRouter Worker error codes', () => { return caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Hello', mode: 'code', model: 'kilo/test-model' }, }); } diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index d506c102fd8..8a4a1b2e0b7 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -76,7 +76,11 @@ import { generateCloudAgentAttachmentUploadUrl, generateImageUploadUrl, } from '@/lib/r2/cloud-agent-attachments'; -import { linkPendingUploads, releasePendingUploads } from '@/lib/r2/cloud-agent-pending-uploads'; +import { + linkAttachmentsForDispatch, + linkPendingUploads, + releasePendingUploads, +} from '@/lib/r2/cloud-agent-pending-uploads'; import * as z from 'zod'; import { PLATFORM } from '@/lib/integrations/core/constants'; @@ -98,7 +102,6 @@ import { buildPreviewOpenUrl } from '@/lib/cloud-agent/preview-open-url'; import { db } from '@kilocode/web-shared/lib/drizzle'; import { verifyOrgOwnsSessionV2ByCloudAgentId } from '@/lib/cloud-agent/session-ownership'; import { TRPCError } from '@trpc/server'; -import { generateMessageId } from '@kilocode/cloud-agent-sdk/message-id'; import { getBalanceForOrganizationUser } from '@kilocode/web-shared/lib/organizations/organization-usage'; import { isMobileClient } from '@kilocode/web-shared/lib/trpc/min-version'; import { buildCloudAgentNextEligibility } from '../cloud-agent-next-eligibility'; @@ -384,7 +387,14 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ } try { - const result = await client.prepareSession({ + // New-session flows call prepareSession without a follow-up + // sendMessage, so the attachment rows must be linked before prepare + // can accept the session; otherwise the reaper deletes new-session + // objects. Only ATTACHMENT rows are admitted (images are never + // admitted; only generateCloudAgentAttachmentUploadUrl admits). + await linkAttachmentsForDispatch(ctx.user.id, attachments); + + return await client.prepareSession({ ...restInput, ...gitParams, attachments: attachments ?? images, @@ -393,21 +403,6 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ kilocodeOrganizationId: organizationId, clientProvenance: isMobileClient(ctx.headersList) ? 'mobile' : 'browser', }); - - // New-session flows call prepareSession without a follow-up - // sendMessage, so the server must link the admitted attachment rows - // here too; otherwise the reaper deletes new-session objects. Only - // ATTACHMENT rows are admitted (images are never admitted; only - // generateCloudAgentAttachmentUploadUrl admits). - if (attachments && attachments.files.length > 0) { - await linkPendingUploads( - ctx.user.id, - attachments.path, - attachments.files.map(file => `${ctx.user.id}/cloud-agent/${attachments.path}/${file}`) - ); - } - - return result; } catch (error) { rethrowCloudAgentError(error); } @@ -516,30 +511,22 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ // for GitHub, GIT_TOKEN_SERVICE for managed GitLab). organizationId is // consumed by the membership middleware; it is not forwarded. try { - const result = await client.sendMessage({ + // Link the attachment rows before the worker can accept the turn: the + // link is what keeps their objects out of the abandoned-upload reaper, + // and once the worker has answered there is no ledger work left that + // could report an accepted send as failed. Only ATTACHMENT rows are + // admitted (images are never admitted; only + // generateCloudAgentAttachmentUploadUrl admits). + const attachments = input.attachments; + await linkAttachmentsForDispatch(ctx.user.id, attachments); + + return await client.sendMessage({ cloudAgentSessionId: input.cloudAgentSessionId, payload: input.payload, autoCommit: input.autoCommit, - messageId: input.messageId ?? generateMessageId(), + messageId: input.messageId, attachments: input.attachments ?? input.images, }); - - // Compatibility: the server links authoritatively so any client whose - // send goes through sendMessage keeps its objects (mobile, old tabs). - // The client does not link pre-send, so a failed or abandoned send - // leaves rows 'pending' for the reaper. Only ATTACHMENT rows are - // admitted (images are never admitted; only - // generateCloudAgentAttachmentUploadUrl admits). - const attachments = input.attachments; - if (attachments && attachments.files.length > 0) { - await linkPendingUploads( - ctx.user.id, - attachments.path, - attachments.files.map(file => `${ctx.user.id}/cloud-agent/${attachments.path}/${file}`) - ); - } - - return result; } catch (error) { rethrowCloudAgentError(error); } From 22de599b43bfd2ba2eb5b5c27b502fd7b4fb4538 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:05:11 +0000 Subject: [PATCH 04/13] chore(kilo-mcp): refresh catalog.json and SKILL.md from CI --- services/kilo-mcp/catalog.json | 28 ++++++++-------------------- 1 file changed, 8 insertions(+), 20 deletions(-) diff --git a/services/kilo-mcp/catalog.json b/services/kilo-mcp/catalog.json index a346d0a1897..caee72a148e 100644 --- a/services/kilo-mcp/catalog.json +++ b/services/kilo-mcp/catalog.json @@ -8186,15 +8186,8 @@ "type": "boolean" }, "messageId": { - "anyOf": [ - { - "type": "string", - "pattern": "^msg_[0-9a-f]{12}[0-9A-Za-z]{14}$" - }, - { - "type": "null" - } - ] + "type": "string", + "pattern": "^msg_[0-9a-f]{12}[0-9A-Za-z]{14}$" }, "attachments": { "type": "object", @@ -8244,7 +8237,8 @@ }, "required": [ "cloudAgentSessionId", - "payload" + "payload", + "messageId" ] }, "tags": [ @@ -19990,15 +19984,8 @@ "type": "boolean" }, "messageId": { - "anyOf": [ - { - "type": "string", - "pattern": "^msg_[0-9a-f]{12}[0-9A-Za-z]{14}$" - }, - { - "type": "null" - } - ] + "type": "string", + "pattern": "^msg_[0-9a-f]{12}[0-9A-Za-z]{14}$" }, "attachments": { "type": "object", @@ -20049,7 +20036,8 @@ "required": [ "organizationId", "cloudAgentSessionId", - "payload" + "payload", + "messageId" ] }, "tags": [ From 7e4a4d4025ef8e2e4be97a5e87e10a6d63a4922e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:39:39 +0000 Subject: [PATCH 05/13] fix(web): keep a refused message's rows reapable and cover the schema fixture linkAttachmentsForDispatch now links and refuses in one transaction, so a message that cannot be dispatched leaves every row as it was instead of pinning the rows that could be linked. The link reports every requested key the message does not hold, which also covers a row a failed object delete returned to pending and therefore made reapable again. The schema test's send fixture carries the now-required message id, and the organization ordering assertion also requires the Worker call to have happened. --- .../r2/cloud-agent-pending-uploads.test.ts | 27 +++-- .../src/lib/r2/cloud-agent-pending-uploads.ts | 100 +++++++++++------- .../routers/cloud-agent-next-schemas.test.ts | 1 + ...ganization-cloud-agent-next-router.test.ts | 1 + 4 files changed, 84 insertions(+), 45 deletions(-) diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts index ccc60a6cbc4..b260e4f4c69 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts @@ -665,7 +665,7 @@ describe('cloud-agent pending-upload ledger', () => { }); const linked = await linkPendingUploads(kiloUserId, messageUuid, [objectKey]); - expect(linked).toEqual({ reapedObjectKeys: [] }); + expect(linked).toEqual({ unlinkedObjectKeys: [] }); const summary = await reapAbandonedUploads(); @@ -696,7 +696,7 @@ describe('cloud-agent pending-upload ledger', () => { await reapAbandonedUploads(); const linked = await linkPendingUploads(kiloUserId, messageUuid, [objectKey]); - expect(linked).toEqual({ reapedObjectKeys: [objectKey] }); + expect(linked).toEqual({ unlinkedObjectKeys: [objectKey] }); const [after] = await db .select({ status: cloud_agent_pending_uploads.status }) .from(cloud_agent_pending_uploads) @@ -753,7 +753,6 @@ describe('cloud-agent pending-upload ledger', () => { const messageUuid = crypto.randomUUID(); const attachment = crypto.randomUUID(); const legacyFile = `${crypto.randomUUID()}.pdf`; - const objectKey = buildObjectKey(kiloUserId, messageUuid, attachment); const future = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString(); const row = await insertPendingRow({ @@ -775,16 +774,24 @@ describe('cloud-agent pending-upload ledger', () => { expect(after?.status).toBe('linked'); }); - it('refuses a dispatch whose attachment row was already reaped', async () => { + it('leaves every row of a refused message reapable', async () => { const kiloUserId = `ca-user-${crypto.randomUUID()}`; const messageUuid = crypto.randomUUID(); - const attachment = crypto.randomUUID(); + const readyAttachment = crypto.randomUUID(); + const goneAttachment = crypto.randomUUID(); + const future = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString(); const past = new Date(Date.now() - 60_000).toISOString(); + const readyRow = await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: readyAttachment, + expiresAt: future, + }); await insertPendingRow({ kiloUserId, messageUuid, - attachmentId: attachment, + attachmentId: goneAttachment, expiresAt: past, status: 'reaped', }); @@ -792,8 +799,14 @@ describe('cloud-agent pending-upload ledger', () => { await expect( linkAttachmentsForDispatch(kiloUserId, { path: messageUuid, - files: [`${attachment}.bin`], + files: [`${readyAttachment}.bin`, `${goneAttachment}.bin`], }) ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + + const [after] = await db + .select({ status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.id, readyRow)); + expect(after?.status).toBe('pending'); }); }); diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts index cad3169acbf..ed8b6c8dac8 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts @@ -1,7 +1,7 @@ import { DeleteObjectCommand } from '@aws-sdk/client-s3'; import { TRPCError } from '@trpc/server'; import { randomUUID } from 'crypto'; -import { and, count, eq, inArray, like, not, sql, type SQL } from 'drizzle-orm'; +import { and, count, eq, inArray, like, ne, not, sql, type SQL } from 'drizzle-orm'; import { db } from '@kilocode/web-shared/lib/drizzle'; import { r2Client, r2CloudAgentAttachmentsBucketName } from '@kilocode/web-shared/lib/r2/client'; import { @@ -135,27 +135,34 @@ export async function admitPendingUpload({ export type LinkPendingUploadsResult = { /** - * Requested keys whose row is 'reaped', so the object behind them is gone or - * being deleted. A message referencing one of these must not be dispatched. + * Requested keys whose row the message does not hold: 'reaped', so the object + * behind them is gone or being deleted, or 'pending' again after a released + * delete put the row back for a later attempt. A message referencing one of + * these must not be dispatched. */ - reapedObjectKeys: string[]; + unlinkedObjectKeys: string[]; }; /** - * Flip pending ledger rows to 'linked' for a user's message in one guarded - * update. Rows whose object keys are not in the supplied set stay 'pending', - * so a partial key set leaves unmatched uploads to be reaped. A row already - * 'linked' belongs to an earlier dispatch of the same message and is reported - * as neither linked now nor reaped. + * Flip pending ledger rows to 'linked' for a user's message, and report the + * requested keys the message does not end up holding. The guarded update links + * the rows that are still 'pending'; the read-back that follows it in the same + * transaction finds every requested key whose row is not 'linked', which covers + * a row a concurrent release or claim took first and a row a failed delete put + * back for a later attempt. Rows whose object keys are not in the supplied set + * stay 'pending', so a partial key set leaves unmatched uploads to be reaped. A + * row already 'linked' belongs to an earlier dispatch of the same message and is + * not reported. */ -export async function linkPendingUploads( +async function linkPendingUploadsInTransaction( + tx: Parameters[0]>[0], kiloUserId: string, messageUuid: string, objectKeys: string[] ): Promise { - if (objectKeys.length === 0) return { reapedObjectKeys: [] }; + if (objectKeys.length === 0) return { unlinkedObjectKeys: [] }; - await db + await tx .update(cloud_agent_pending_uploads) .set({ status: 'linked' }) .where( @@ -167,25 +174,36 @@ export async function linkPendingUploads( ) ); - // Read back what the guarded update could not link. A row that a release or - // a reaper claim already took out of 'pending' is 'reaped', and its object is - // going away, so the caller must not dispatch a message that references it. - const rows = await db - .select({ - object_key: cloud_agent_pending_uploads.object_key, - status: cloud_agent_pending_uploads.status, - }) + // Read back every requested key the message does not hold. A row is left + // 'reaped' when a release or a claim took it, and is returned to 'pending' when + // its delete failed, so in either state the object is under the reaper's + // control and the caller must not dispatch a message that references it. + const rows = await tx + .select({ object_key: cloud_agent_pending_uploads.object_key }) .from(cloud_agent_pending_uploads) .where( and( eq(cloud_agent_pending_uploads.kilo_user_id, kiloUserId), eq(cloud_agent_pending_uploads.message_uuid, messageUuid), - eq(cloud_agent_pending_uploads.status, 'reaped'), + ne(cloud_agent_pending_uploads.status, 'linked'), inArray(cloud_agent_pending_uploads.object_key, objectKeys) ) ); - return { reapedObjectKeys: rows.map(row => row.object_key) }; + return { unlinkedObjectKeys: rows.map(row => row.object_key) }; +} + +/** Link a message's attachment rows and report the ones it does not hold. */ +export async function linkPendingUploads( + kiloUserId: string, + messageUuid: string, + objectKeys: string[] +): Promise { + if (objectKeys.length === 0) return { unlinkedObjectKeys: [] }; + + return db.transaction(tx => + linkPendingUploadsInTransaction(tx, kiloUserId, messageUuid, objectKeys) + ); } /** Message for a send or prepare whose attachment object is already gone. */ @@ -193,10 +211,12 @@ export const CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE = 'Attachment upload expired. Add the file again to send it.'; /** - * Link the ledger rows of one message's attachments before its dispatch, so - * the objects survive the reaper, and refuse the dispatch when a row was - * already reaped. Keys with no ledger row belong to an uploader that predates - * the ledger and are left alone. + * Link the ledger rows of one message's attachments before its dispatch, so the + * objects survive the reaper, and refuse the dispatch when the message does not + * hold one of its rows. The link and the refusal share a transaction, so a + * refused message keeps every row as it was and a retry after the user + * re-attaches the file links them all. Keys with no ledger row belong to an + * uploader that predates the ledger and are left alone. */ export async function linkAttachmentsForDispatch( kiloUserId: string, @@ -204,18 +224,22 @@ export async function linkAttachmentsForDispatch( ): Promise { if (!attachments || attachments.files.length === 0) return; - const { reapedObjectKeys } = await linkPendingUploads( - kiloUserId, - attachments.path, - attachments.files.map(file => `${kiloUserId}/cloud-agent/${attachments.path}/${file}`) - ); - - if (reapedObjectKeys.length > 0) { - throw new TRPCError({ - code: 'BAD_REQUEST', - message: CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE, - }); - } + await db.transaction(async tx => { + // The link and the refusal share one transaction, so a message whose + // dispatch is refused keeps all of its rows reapable. + const { unlinkedObjectKeys } = await linkPendingUploadsInTransaction( + tx, + kiloUserId, + attachments.path, + attachments.files.map(file => `${kiloUserId}/cloud-agent/${attachments.path}/${file}`) + ); + if (unlinkedObjectKeys.length > 0) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE, + }); + } + }); } export type ReapAbandonedUploadsSummary = { diff --git a/apps/web/src/routers/cloud-agent-next-schemas.test.ts b/apps/web/src/routers/cloud-agent-next-schemas.test.ts index 7354082d0bb..8b482a66a27 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.test.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.test.ts @@ -660,6 +660,7 @@ describe('send-time attachment filenames', () => { }; const sendInput = { cloudAgentSessionId: 'agent_12345678-1234-4234-9234-123456789abc', + messageId: 'msg_123456789abc123456789ABCDE', payload: { type: 'prompt', prompt: '', mode: 'code', model: 'kilo/test-model' }, }; const attachmentsWith = (extension: string) => ({ diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts index 9650c64acff..9874280c5f3 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts @@ -1019,6 +1019,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { }); expect(linkAttachmentsForDispatch).toHaveBeenCalledWith('user-1', attachments); + expect(mockSendMessage).toHaveBeenCalledTimes(1); expect(linkAttachmentsForDispatch.mock.invocationCallOrder[0]).toBeLessThan( mockSendMessage.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY ); From 55fa9cfa1293c9e683d833873d68a20146d04fd0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:45:59 +0000 Subject: [PATCH 06/13] test(web): give the send schema fixtures the required message id The dispatch schema now requires the client message id, so the send fixtures in cloud-agent-next-send-schema.test.ts carry one. --- apps/web/src/routers/cloud-agent-next-send-schema.test.ts | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/apps/web/src/routers/cloud-agent-next-send-schema.test.ts b/apps/web/src/routers/cloud-agent-next-send-schema.test.ts index 08092401c67..3cb4e3e8268 100644 --- a/apps/web/src/routers/cloud-agent-next-send-schema.test.ts +++ b/apps/web/src/routers/cloud-agent-next-send-schema.test.ts @@ -6,6 +6,7 @@ import { const MESSAGE_UUID = '12345678-1234-4234-9234-123456789abc'; const ATTACHMENT_ID = '87654321-4321-4321-8321-cba987654321'; +const MESSAGE_ID = 'msg_123456789abc123456789ABCDE'; const promptPayload = (prompt: string) => ({ type: 'prompt' as const, @@ -33,6 +34,7 @@ describe('baseSendMessageNextSchema', () => { it('accepts an empty prompt with an attachment', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: promptPayload(''), attachments: attachments(`${ATTACHMENT_ID}.png`), }); @@ -42,6 +44,7 @@ describe('baseSendMessageNextSchema', () => { it('accepts an empty prompt with an image', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: promptPayload(''), images: images(`${ATTACHMENT_ID}.png`), }); @@ -51,6 +54,7 @@ describe('baseSendMessageNextSchema', () => { it('rejects an empty prompt with no files', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: promptPayload(''), }); expect(result.success).toBe(false); @@ -59,6 +63,7 @@ describe('baseSendMessageNextSchema', () => { it('rejects a whitespace-only prompt with no files', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: promptPayload(' '), }); expect(result.success).toBe(false); @@ -67,6 +72,7 @@ describe('baseSendMessageNextSchema', () => { it('accepts a command payload with no prompt and no files', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: commandPayload(), }); expect(result.success).toBe(true); @@ -75,6 +81,7 @@ describe('baseSendMessageNextSchema', () => { it('still accepts a text-only prompt', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: promptPayload('hello'), }); expect(result.success).toBe(true); @@ -83,6 +90,7 @@ describe('baseSendMessageNextSchema', () => { it('still accepts a text prompt with an attachment', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: promptPayload('hello'), attachments: attachments(`${ATTACHMENT_ID}.png`), }); From 01616deb46e3ef6576713fee2c94c83b1848c121 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:45:29 +0000 Subject: [PATCH 07/13] test(cloud-agent): retain message ids in payload boundary cases --- apps/web/src/routers/cloud-agent-next-send-schema.test.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/apps/web/src/routers/cloud-agent-next-send-schema.test.ts b/apps/web/src/routers/cloud-agent-next-send-schema.test.ts index 7855b440f5c..5ee43bf5605 100644 --- a/apps/web/src/routers/cloud-agent-next-send-schema.test.ts +++ b/apps/web/src/routers/cloud-agent-next-send-schema.test.ts @@ -105,6 +105,7 @@ describe('baseSendMessageNextSchema', () => { it('rejects a follow-up prompt one character over the shared cap', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: promptPayload(overCap), }); expect(result.success).toBe(false); @@ -113,6 +114,7 @@ describe('baseSendMessageNextSchema', () => { it('accepts follow-up command arguments at the shared cap', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: commandPayload(atCap), }); expect(result.success).toBe(true); @@ -121,6 +123,7 @@ describe('baseSendMessageNextSchema', () => { it('rejects follow-up command arguments one character over the shared cap', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: commandPayload(overCap), }); expect(result.success).toBe(false); From 75a53cabc5b77ae06757900729fbd896e44a6b39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Sat, 24 Oct 2026 20:16:43 +0000 Subject: [PATCH 08/13] fix(web): drop ledger imports the dispatch helpers no longer use --- apps/web/src/routers/cloud-agent-next-router.ts | 1 - .../organizations/organization-cloud-agent-next-router.ts | 1 - 2 files changed, 2 deletions(-) diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index bfc7554c124..51caf4c87e6 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -73,7 +73,6 @@ import { } from '@/lib/r2/cloud-agent-attachments'; import { linkAttachmentsForDispatch, - linkPendingUploads, releasePendingUploads, } from '@/lib/r2/cloud-agent-pending-uploads'; import * as z from 'zod'; diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index 48c9265e828..0b973e0bcef 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -78,7 +78,6 @@ import { } from '@/lib/r2/cloud-agent-attachments'; import { linkAttachmentsForDispatch, - linkPendingUploads, releasePendingUploads, } from '@/lib/r2/cloud-agent-pending-uploads'; import * as z from 'zod'; From f4d376e42fc9078b3c57f41963bcf531124363e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Sat, 24 Oct 2026 21:43:03 +0000 Subject: [PATCH 09/13] fix(web): return a failed send's attachments when the worker did not admit it After a send fails, the router asks the worker whether that message exists. Only a message the worker reports as not found hands the attachment rows back to the abandoned-upload reaper; a message the worker holds, and an unavailable lookup, keep them linked. --- .../release-unadmitted-send.test.ts | 178 ++++++++++++++++++ .../release-unadmitted-send.ts | 71 +++++++ .../r2/cloud-agent-pending-uploads.test.ts | 72 ++++++- .../src/lib/r2/cloud-agent-pending-uploads.ts | 48 ++++- .../routers/cloud-agent-next-router.test.ts | 16 ++ .../src/routers/cloud-agent-next-router.ts | 17 +- ...ganization-cloud-agent-next-router.test.ts | 36 ++++ .../organization-cloud-agent-next-router.ts | 19 +- 8 files changed, 451 insertions(+), 6 deletions(-) create mode 100644 apps/web/src/lib/cloud-agent-next/release-unadmitted-send.test.ts create mode 100644 apps/web/src/lib/cloud-agent-next/release-unadmitted-send.ts diff --git a/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.test.ts b/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.test.ts new file mode 100644 index 00000000000..042853a5257 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.test.ts @@ -0,0 +1,178 @@ +import { beforeEach, describe, expect, it, jest } from '@jest/globals'; + +import type { DispatchAttachments } from '@/lib/r2/cloud-agent-pending-uploads'; + +import { releaseUnadmittedSend } from './release-unadmitted-send'; + +const mockReleaseUnadmittedUploads = jest.fn( + async (_kiloUserId: string, _attachments: DispatchAttachments) => 1 +); + +const messageId = 'msg_0123456789abABCDEFGHIJKLMN'; +const attachments = { path: 'f7c1f9a0-0000-4000-8000-000000000001', files: ['a.bin'] }; + +const sessionId = 'workspace_00000000-0000-4000-8000-000000000002'; + +function admittedResult() { + return { cloudAgentSessionId: sessionId, messageId, status: 'queued' as const }; +} + +function buildLookup(results: readonly (Error | null | 'admitted')[]) { + let calls = 0; + const getMessageResult = jest.fn( + async (_input: { cloudAgentSessionId: string; messageId: string }) => { + const result = results[Math.min(calls, results.length - 1)]; + calls += 1; + if (result instanceof Error) throw result; + return result === null ? null : admittedResult(); + } + ); + return { getMessageResult }; +} + +const dependencies = { delay: async () => {}, release: mockReleaseUnadmittedUploads }; + +describe('releaseUnadmittedSend', () => { + beforeEach(() => { + jest.restoreAllMocks(); + jest.clearAllMocks(); + jest.spyOn(console, 'error').mockImplementation(() => {}); + jest.spyOn(console, 'info').mockImplementation(() => {}); + }); + + it('returns the rows when both lookups report no such message', async () => { + const client = buildLookup([null, null]); + + await releaseUnadmittedSend( + { + attachments, + client, + cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', + kiloUserId: 'user-1', + messageId, + }, + dependencies + ); + + expect(client.getMessageResult).toHaveBeenCalledTimes(2); + expect(client.getMessageResult).toHaveBeenCalledWith({ + cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', + messageId, + }); + expect(mockReleaseUnadmittedUploads).toHaveBeenCalledWith('user-1', attachments); + }); + + it('keeps the rows when the message exists at either lookup', async () => { + const admittedFirst = buildLookup(['admitted']); + await releaseUnadmittedSend( + { + attachments, + client: admittedFirst, + cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', + kiloUserId: 'user-1', + messageId, + }, + dependencies + ); + expect(admittedFirst.getMessageResult).toHaveBeenCalledTimes(1); + + const admittedLater = buildLookup([null, 'admitted']); + await releaseUnadmittedSend( + { + attachments, + client: admittedLater, + cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', + kiloUserId: 'user-1', + messageId, + }, + dependencies + ); + expect(admittedLater.getMessageResult).toHaveBeenCalledTimes(2); + expect(mockReleaseUnadmittedUploads).not.toHaveBeenCalled(); + }); + + it('keeps the rows when a lookup fails', async () => { + const unavailable = buildLookup([new Error('worker unavailable')]); + + await expect( + releaseUnadmittedSend( + { + attachments, + client: unavailable, + cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', + kiloUserId: 'user-1', + messageId, + }, + dependencies + ) + ).resolves.toBeUndefined(); + + expect(jest.mocked(console.error)).toHaveBeenCalledWith( + '[cloud-agent] Could not confirm a failed send was not admitted', + expect.any(Error) + ); + expect(mockReleaseUnadmittedUploads).not.toHaveBeenCalled(); + }); + + it('keeps the rows when the second lookup is the one that fails', async () => { + const secondFails = buildLookup([null, new Error('worker unavailable')]); + + await releaseUnadmittedSend( + { + attachments, + client: secondFails, + cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', + kiloUserId: 'user-1', + messageId, + }, + dependencies + ); + + expect(secondFails.getMessageResult).toHaveBeenCalledTimes(2); + expect(mockReleaseUnadmittedUploads).not.toHaveBeenCalled(); + }); + + it('stays silent when the ledger cannot return the rows', async () => { + const client = buildLookup([null, null]); + const release = jest.fn(async () => { + throw new Error('database unavailable'); + }); + + await expect( + releaseUnadmittedSend( + { + attachments, + client, + cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', + kiloUserId: 'user-1', + messageId, + }, + { delay: async () => {}, release } + ) + ).resolves.toBeUndefined(); + + expect(release).toHaveBeenCalledTimes(1); + expect(jest.mocked(console.error)).toHaveBeenCalledWith( + '[cloud-agent] Could not return the attachments of an unadmitted send', + expect.any(Error) + ); + }); + + it('does nothing for a send without attachments', async () => { + const client = buildLookup([null]); + + await releaseUnadmittedSend( + { + attachments: undefined, + client, + cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', + kiloUserId: 'user-1', + messageId, + }, + dependencies + ); + + expect(client.getMessageResult).not.toHaveBeenCalled(); + expect(mockReleaseUnadmittedUploads).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.ts b/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.ts new file mode 100644 index 00000000000..198bcfa5887 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.ts @@ -0,0 +1,71 @@ +import type { CloudAgentNextClient } from './cloud-agent-client'; +import { + releaseUnadmittedUploads, + type DispatchAttachments, +} from '@/lib/r2/cloud-agent-pending-uploads'; + +/** Time between the first lookup and its confirmation. */ +const confirmationDelayMs = 1_000; + +type MessageResultLookup = Pick; + +export type UnadmittedSend = { + attachments: DispatchAttachments | undefined; + client: MessageResultLookup; + cloudAgentSessionId: string; + kiloUserId: string; + messageId: string; +}; + +type ReleaseDependencies = { + delay?: (milliseconds: number) => Promise; + release?: (kiloUserId: string, attachments: DispatchAttachments) => Promise; +}; + +const sleep = (milliseconds: number): Promise => { + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, milliseconds); + return promise; +}; + +/** + * Give a failed send's attachment rows back to the abandoned-upload lifecycle, + * but only when the worker states that it never admitted the turn. + * + * The worker records an admitted message before it answers the send, and its + * session object serves one request at a time, so a lookup that still reports + * no such message is a statement that the turn did not happen. The lookup runs + * twice, one confirmation apart, because a send that failed in transport can + * still be on its way. Every other answer keeps the rows linked: a message the + * worker holds must never lose its attachments, and an unavailable lookup is + * not proof that the turn was refused. + * + * The send's own error is what the caller reports; this helper never throws. + */ +export async function releaseUnadmittedSend( + input: UnadmittedSend, + { delay = sleep, release = releaseUnadmittedUploads }: ReleaseDependencies = {} +): Promise { + const { attachments, client, cloudAgentSessionId, kiloUserId, messageId } = input; + if (!attachments || attachments.files.length === 0) return; + + try { + for (let attempt = 0; attempt < 2; attempt += 1) { + if (attempt > 0) await delay(confirmationDelayMs); + const result = await client.getMessageResult({ cloudAgentSessionId, messageId }); + if (result !== null) return; + } + } catch (error) { + console.error('[cloud-agent] Could not confirm a failed send was not admitted', error); + return; + } + + try { + const released = await release(kiloUserId, attachments); + if (released > 0) { + console.info('[cloud-agent] Returned the attachments of an unadmitted send', { released }); + } + } catch (error) { + console.error('[cloud-agent] Could not return the attachments of an unadmitted send', error); + } +} diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts index 2593ee2898c..17b4c02a4c8 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts @@ -6,6 +6,7 @@ import type { purgeUserPendingUploads as PurgeUserPendingUploads, reapAbandonedUploads as ReapAbandonedUploads, releasePendingUploads as ReleasePendingUploads, + releaseUnadmittedUploads as ReleaseUnadmittedUploads, } from './cloud-agent-pending-uploads'; const mockSend = jest.fn<(command: { input: { Key: string }; name: string }) => Promise>(); @@ -22,7 +23,7 @@ jest.mock('@aws-sdk/client-s3', () => ({ DeleteObjectCommand: mockDeleteObjectCommand, })); -import { and, eq, sql } from 'drizzle-orm'; +import { and, eq, inArray, sql } from 'drizzle-orm'; import { CLOUD_AGENT_ATTACHMENT_MAX_COUNT, CLOUD_AGENT_ATTACHMENT_MAX_SIZE_BYTES, @@ -38,6 +39,7 @@ let linkPendingUploads: typeof LinkPendingUploads; let purgeUserPendingUploads: typeof PurgeUserPendingUploads; let reapAbandonedUploads: typeof ReapAbandonedUploads; let releasePendingUploads: typeof ReleasePendingUploads; +let releaseUnadmittedUploads: typeof ReleaseUnadmittedUploads; function buildObjectKey(kiloUserId: string, messageUuid: string, attachmentId: string): string { return `${kiloUserId}/cloud-agent/${messageUuid}/${attachmentId}.bin`; @@ -83,6 +85,7 @@ describe('cloud-agent pending-upload ledger', () => { purgeUserPendingUploads = pendingUploadsModule.purgeUserPendingUploads; reapAbandonedUploads = pendingUploadsModule.reapAbandonedUploads; releasePendingUploads = pendingUploadsModule.releasePendingUploads; + releaseUnadmittedUploads = pendingUploadsModule.releaseUnadmittedUploads; }); afterEach(async () => { @@ -919,6 +922,73 @@ describe('cloud-agent pending-upload ledger', () => { expect(after?.status).toBe('linked'); }); + it('returns only the unadmitted message rows to pending and leaves another message linked', async () => { + const kiloUserId = `ca-user-${crypto.randomUUID()}`; + const unadmittedMessage = crypto.randomUUID(); + const otherMessage = crypto.randomUUID(); + const unadmittedAttachment = crypto.randomUUID(); + const otherAttachment = crypto.randomUUID(); + const future = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString(); + + const unadmittedRow = await insertPendingRow({ + kiloUserId, + messageUuid: unadmittedMessage, + attachmentId: unadmittedAttachment, + expiresAt: future, + status: 'linked', + }); + const otherRow = await insertPendingRow({ + kiloUserId, + messageUuid: otherMessage, + attachmentId: otherAttachment, + expiresAt: future, + status: 'linked', + }); + + const released = await releaseUnadmittedUploads(kiloUserId, { + path: unadmittedMessage, + files: [`${unadmittedAttachment}.bin`], + }); + + expect(released).toBe(1); + const rows = await db + .select({ id: cloud_agent_pending_uploads.id, status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(inArray(cloud_agent_pending_uploads.id, [unadmittedRow, otherRow])); + const statusById = new Map(rows.map(row => [row.id, row.status])); + expect(statusById.get(unadmittedRow)).toBe('pending'); + expect(statusById.get(otherRow)).toBe('linked'); + }); + + it('lets the reaper reclaim a returned row once its lease lapses', async () => { + const kiloUserId = `ca-user-${crypto.randomUUID()}`; + const messageUuid = crypto.randomUUID(); + const attachment = crypto.randomUUID(); + const past = new Date(Date.now() - 60_000).toISOString(); + mockSend.mockResolvedValue({}); + + const row = await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: attachment, + expiresAt: past, + status: 'linked', + }); + + await releaseUnadmittedUploads(kiloUserId, { + path: messageUuid, + files: [`${attachment}.bin`], + }); + const summary = await reapAbandonedUploads(); + + expect(summary.reaped).toBe(1); + const [after] = await db + .select({ status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.id, row)); + expect(after?.status).toBe('reaped'); + }); + it('reports the attachment row a reaper claim took first and does not link it', async () => { const kiloUserId = `ca-user-${crypto.randomUUID()}`; const messageUuid = crypto.randomUUID(); diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts index 2887a0e4901..d0d9b4ca2c6 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts @@ -166,6 +166,17 @@ export type LinkPendingUploadsResult = { unlinkedObjectKeys: string[]; }; +/** The attachments of one message, as the send and prepare inputs carry them. */ +export type DispatchAttachments = { path: string; files: string[] }; + +/** Object keys of one message's attachments, as the presign helper admitted them. */ +export function attachmentObjectKeys( + kiloUserId: string, + attachments: DispatchAttachments +): string[] { + return attachments.files.map(file => `${kiloUserId}/cloud-agent/${attachments.path}/${file}`); +} + /** * Flip pending ledger rows to 'linked' for a user's message, and report the * requested keys the message does not end up holding. The guarded update links @@ -243,7 +254,7 @@ export const CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE = */ export async function linkAttachmentsForDispatch( kiloUserId: string, - attachments: { path: string; files: string[] } | undefined + attachments: DispatchAttachments | undefined ): Promise { if (!attachments || attachments.files.length === 0) return; @@ -254,7 +265,7 @@ export async function linkAttachmentsForDispatch( tx, kiloUserId, attachments.path, - attachments.files.map(file => `${kiloUserId}/cloud-agent/${attachments.path}/${file}`) + attachmentObjectKeys(kiloUserId, attachments) ); if (unlinkedObjectKeys.length > 0) { throw new TRPCError({ @@ -265,6 +276,39 @@ export async function linkAttachmentsForDispatch( }); } +/** + * Hand one message's attachment rows back to the abandoned-upload lifecycle + * after the worker stated it never admitted the turn. Only the rows this message + * still holds are moved back, so a row another message holds or the reaper + * already took keeps its state. The rows keep their original lease, so the + * reaper claims them on its next run once that lease lapses. Returns the number + * of rows handed back. + */ +export async function releaseUnadmittedUploads( + kiloUserId: string, + attachments: DispatchAttachments +): Promise { + if (attachments.files.length === 0) return 0; + + const released = await db + .update(cloud_agent_pending_uploads) + .set({ status: 'pending' }) + .where( + and( + eq(cloud_agent_pending_uploads.kilo_user_id, kiloUserId), + eq(cloud_agent_pending_uploads.message_uuid, attachments.path), + eq(cloud_agent_pending_uploads.status, 'linked'), + inArray( + cloud_agent_pending_uploads.object_key, + attachmentObjectKeys(kiloUserId, attachments) + ) + ) + ) + .returning({ id: cloud_agent_pending_uploads.id }); + + return released.length; +} + export type ReapAbandonedUploadsSummary = { reaped: number; }; diff --git a/apps/web/src/routers/cloud-agent-next-router.test.ts b/apps/web/src/routers/cloud-agent-next-router.test.ts index df63234c406..a721b1f2cac 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -41,6 +41,7 @@ import { TRPCError } from '@trpc/server'; import { getHTTPStatusCodeFromError } from '@trpc/server/http'; import type { verifyUserOwnsSessionV2ByCloudAgentId } from '@/lib/cloud-agent/session-ownership'; import type * as StreamTicketModule from '@/lib/cloud-agent/stream-ticket'; +import type { releaseUnadmittedSend } from '@/lib/cloud-agent-next/release-unadmitted-send'; type AttachmentReference = { path: string; files: string[] }; @@ -157,6 +158,7 @@ const mockIsFeatureFlagEnabledOrDevelopment = jest.fn<(flagName: string, distinctId: string) => Promise>(); const mockVerifyUserOwnsSessionV2ByCloudAgentId = jest.fn(); +const mockReleaseUnadmittedSend = jest.fn(async () => {}); const mockSignPreviewTicket = jest.fn(); const mockGetBalanceForUser = jest.fn<(user: User) => Promise<{ balance: number }>>(); const mockFetchGitHubRepositoriesForUser = jest.fn< @@ -210,6 +212,10 @@ jest.mock('@/lib/cloud-agent-next/worktree-chat', () => ({ createWorktreeChat: mockCreateWorktreeChat, })); +jest.mock('@/lib/cloud-agent-next/release-unadmitted-send', () => ({ + releaseUnadmittedSend: mockReleaseUnadmittedSend, +})); + jest.mock('@kilocode/web-shared/lib/trpc/min-version', () => ({ ...jest.requireActual('@kilocode/web-shared/lib/trpc/min-version'), getMinimumVersions: jest.fn(async () => ({ ios: '0.0.0', android: '0.0.0' })), @@ -883,6 +889,8 @@ describe('cloudAgentNextRouter send attachment ledger', () => { }); expect(mockSendMessage).not.toHaveBeenCalled(); + // The refusal rolls the link back, so there is nothing to hand back. + expect(mockReleaseUnadmittedSend).not.toHaveBeenCalled(); }); it('links the attachment rows before the worker accepts the send', async () => { @@ -899,6 +907,7 @@ describe('cloudAgentNextRouter send attachment ledger', () => { ).resolves.toMatchObject({ messageId: SEND_MESSAGE_ID }); expect(mockSendMessage).toHaveBeenCalledWith(expect.objectContaining({ attachments })); + expect(mockReleaseUnadmittedSend).not.toHaveBeenCalled(); expect(await readStatus()).toBe('linked'); }); @@ -916,6 +925,13 @@ describe('cloudAgentNextRouter send attachment ledger', () => { }) ).rejects.toThrow('worker unavailable'); + expect(mockReleaseUnadmittedSend).toHaveBeenCalledWith({ + attachments, + client: expect.anything(), + cloudAgentSessionId: 'agent_123', + kiloUserId: userId, + messageId: SEND_MESSAGE_ID, + }); expect(await readStatus()).toBe('linked'); }); diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index 51caf4c87e6..6ac6c64321f 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -8,6 +8,7 @@ import { } from '@/lib/cloud-agent-next/cloud-agent-client'; import { computeCloudAgentNextBalanceCheckEligibility } from '@/lib/cloud-agent-next/balance-check-eligibility'; import { isPublicRepoUrlEnabled } from '@/lib/cloud-agent-next/public-repo-url'; +import { releaseUnadmittedSend } from '@/lib/cloud-agent-next/release-unadmitted-send'; import { rethrowAsTerminalError } from '@/lib/cloud-agent-next/terminal-errors'; import { createWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; import { assertSessionWorktree } from '@/lib/cloud-agent-next/worktree-review-access'; @@ -369,8 +370,8 @@ export const cloudAgentNextRouter = createTRPCRouter({ // Tokens are refreshed inside cloud-agent-next (GitHub App installation // for GitHub, GIT_TOKEN_SERVICE for managed GitLab). + const { attachments, images } = input; try { - const { attachments, images } = input; // Link the attachment rows before the worker can accept the turn: the // link is what keeps their objects out of the abandoned-upload reaper, // and once the worker has answered there is no ledger work left that @@ -378,7 +379,11 @@ export const cloudAgentNextRouter = createTRPCRouter({ // admitted (images are never admitted; only // generateCloudAgentAttachmentUploadUrl admits). await linkAttachmentsForDispatch(ctx.user.id, attachments); + } catch (error) { + rethrowCloudAgentError(error); + } + try { return await client.sendMessage({ cloudAgentSessionId: input.cloudAgentSessionId, payload: input.payload, @@ -387,6 +392,16 @@ export const cloudAgentNextRouter = createTRPCRouter({ messageId: input.messageId, }); } catch (error) { + // A send can fail after the worker accepted the turn, so the rows are + // only returned to the reaper when the worker states that it never + // admitted this message. + await releaseUnadmittedSend({ + attachments, + client, + cloudAgentSessionId: input.cloudAgentSessionId, + kiloUserId: ctx.user.id, + messageId: input.messageId, + }); rethrowCloudAgentError(error); } }), diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts index 56bf3156df7..b9d27d92342 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts @@ -294,6 +294,10 @@ jest.mock('@/lib/r2/cloud-agent-pending-uploads', () => ({ releasePendingUploads: jest.fn(), })); +jest.mock('@/lib/cloud-agent-next/release-unadmitted-send', () => ({ + releaseUnadmittedSend: jest.fn(async () => {}), +})); + jest.mock('@kilocode/web-shared/routers/organizations/utils', () => { const trpcInit = jest.requireActual('@kilocode/web-shared/lib/trpc/init'); const zod = jest.requireActual('zod'); @@ -1033,6 +1037,38 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { ); }); + it('asks the Worker whether a failed send was admitted before keeping the rows', async () => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + const attachments = { + path: '12345678-1234-4234-9234-123456789abc', + files: ['87654321-4321-4321-8321-cba987654321.md'], + }; + const releaseUnadmittedSend = jest.mocked( + jest.requireMock<{ releaseUnadmittedSend: (input: unknown) => Promise }>( + '@/lib/cloud-agent-next/release-unadmitted-send' + ).releaseUnadmittedSend + ); + mockSendMessage.mockRejectedValueOnce(new TRPCClientError('worker unavailable')); + + await expect( + caller.sendMessage({ + organizationId: ORGANIZATION_ID, + cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, + payload: { type: 'prompt', prompt: 'Read notes', mode: 'code', model: 'test' }, + attachments, + }) + ).rejects.toThrow('worker unavailable'); + + expect(releaseUnadmittedSend).toHaveBeenCalledWith({ + attachments, + client: expect.anything(), + cloudAgentSessionId: 'agent_123', + kiloUserId: 'user-1', + messageId: SEND_MESSAGE_ID, + }); + }); + it('normalizes legacy image requests to canonical Worker attachments', async () => { const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); const images = { diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index 0b973e0bcef..569222a8550 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -8,6 +8,7 @@ import { } from '@/lib/cloud-agent-next/cloud-agent-client'; import { computeCloudAgentNextBalanceCheckEligibility } from '@/lib/cloud-agent-next/balance-check-eligibility'; import { isPublicRepoUrlEnabled } from '@/lib/cloud-agent-next/public-repo-url'; +import { releaseUnadmittedSend } from '@/lib/cloud-agent-next/release-unadmitted-send'; import { rethrowAsTerminalError } from '@/lib/cloud-agent-next/terminal-errors'; import { createWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; import { assertSessionWorktree } from '@/lib/cloud-agent-next/worktree-review-access'; @@ -522,6 +523,7 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ // Tokens are refreshed inside cloud-agent-next (GitHub App installation // for GitHub, GIT_TOKEN_SERVICE for managed GitLab). organizationId is // consumed by the membership middleware; it is not forwarded. + const attachments = input.attachments; try { // Link the attachment rows before the worker can accept the turn: the // link is what keeps their objects out of the abandoned-upload reaper, @@ -529,17 +531,30 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ // could report an accepted send as failed. Only ATTACHMENT rows are // admitted (images are never admitted; only // generateCloudAgentAttachmentUploadUrl admits). - const attachments = input.attachments; await linkAttachmentsForDispatch(ctx.user.id, attachments); + } catch (error) { + rethrowCloudAgentError(error); + } + try { return await client.sendMessage({ cloudAgentSessionId: input.cloudAgentSessionId, payload: input.payload, autoCommit: input.autoCommit, messageId: input.messageId, - attachments: input.attachments ?? input.images, + attachments: attachments ?? input.images, }); } catch (error) { + // A send can fail after the worker accepted the turn, so the rows are + // only returned to the reaper when the worker states that it never + // admitted this message. + await releaseUnadmittedSend({ + attachments, + client, + cloudAgentSessionId: input.cloudAgentSessionId, + kiloUserId: ctx.user.id, + messageId: input.messageId, + }); rethrowCloudAgentError(error); } }), From 18e0cf5bf448a3240e1c00408382ca3fc03d48ad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Sat, 24 Oct 2026 22:06:58 +0000 Subject: [PATCH 10/13] fix(web): hand back only the rows a refused send linked The link now reports the rows it moved to linked, and the failure path hands back exactly those keys and only when the worker itself refused the turn. A composer path reused by a later send therefore cannot free an accepted message's attachments, and a transport or server-side failure keeps them. --- .../release-refused-send.test.ts | 109 +++++++++++ .../cloud-agent-next/release-refused-send.ts | 59 ++++++ .../release-unadmitted-send.test.ts | 178 ------------------ .../release-unadmitted-send.ts | 71 ------- .../r2/cloud-agent-pending-uploads.test.ts | 54 ++++-- .../src/lib/r2/cloud-agent-pending-uploads.ts | 61 +++--- .../routers/cloud-agent-next-router.test.ts | 21 ++- .../src/routers/cloud-agent-next-router.ts | 18 +- ...ganization-cloud-agent-next-router.test.ts | 31 +-- .../organization-cloud-agent-next-router.ts | 18 +- 10 files changed, 293 insertions(+), 327 deletions(-) create mode 100644 apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts create mode 100644 apps/web/src/lib/cloud-agent-next/release-refused-send.ts delete mode 100644 apps/web/src/lib/cloud-agent-next/release-unadmitted-send.test.ts delete mode 100644 apps/web/src/lib/cloud-agent-next/release-unadmitted-send.ts diff --git a/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts b/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts new file mode 100644 index 00000000000..fa2b675d156 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts @@ -0,0 +1,109 @@ +import { beforeEach, describe, expect, it, jest } from '@jest/globals'; +import { TRPCClientError } from '@trpc/client'; + +import { releaseRefusedSend, workerRefusalStatus } from './release-refused-send'; + +const mockRelease = jest.fn( + async (_kiloUserId: string, _messageUuid: string, _objectKeys: string[]) => 1 +); + +const attachments = { path: 'f7c1f9a0-0000-4000-8000-000000000001', files: ['a.bin'] }; +const linkedObjectKeys = ['user-1/cloud-agent/f7c1f9a0-0000-4000-8000-000000000001/a.bin']; + +const refusal = (httpStatus: number) => + new TRPCClientError(`worker answered ${String(httpStatus)}`, { + result: { error: { code: -32600, message: 'refused', data: { httpStatus } } }, + }); + +const dependencies = { release: mockRelease }; + +describe('releaseRefusedSend', () => { + beforeEach(() => { + jest.restoreAllMocks(); + jest.clearAllMocks(); + jest.spyOn(console, 'error').mockImplementation(() => {}); + jest.spyOn(console, 'info').mockImplementation(() => {}); + }); + + it('returns the rows this dispatch linked when the worker refused the turn', async () => { + await releaseRefusedSend( + { attachments, error: refusal(400), kiloUserId: 'user-1', linkedObjectKeys }, + dependencies + ); + + expect(mockRelease).toHaveBeenCalledWith( + 'user-1', + 'f7c1f9a0-0000-4000-8000-000000000001', + linkedObjectKeys + ); + }); + + it('keeps the rows when the worker answered with a server-side status', async () => { + await releaseRefusedSend( + { attachments, error: refusal(503), kiloUserId: 'user-1', linkedObjectKeys }, + dependencies + ); + + expect(mockRelease).not.toHaveBeenCalled(); + }); + + it('keeps the rows when the send failed without a worker answer', async () => { + await releaseRefusedSend( + { + attachments, + error: new TRPCClientError('fetch failed'), + kiloUserId: 'user-1', + linkedObjectKeys, + }, + dependencies + ); + + expect(mockRelease).not.toHaveBeenCalled(); + }); + + it('keeps the rows when this dispatch linked none of them', async () => { + await releaseRefusedSend( + { attachments, error: refusal(400), kiloUserId: 'user-1', linkedObjectKeys: [] }, + dependencies + ); + + expect(mockRelease).not.toHaveBeenCalled(); + }); + + it('does nothing for a send without attachments', async () => { + await releaseRefusedSend( + { attachments: undefined, error: refusal(400), kiloUserId: 'user-1', linkedObjectKeys }, + dependencies + ); + + expect(mockRelease).not.toHaveBeenCalled(); + }); + + it('stays silent when the ledger cannot return the rows', async () => { + const release = jest.fn(async () => { + throw new Error('database unavailable'); + }); + + await expect( + releaseRefusedSend( + { attachments, error: refusal(400), kiloUserId: 'user-1', linkedObjectKeys }, + { release } + ) + ).resolves.toBeUndefined(); + + expect(release).toHaveBeenCalledTimes(1); + expect(jest.mocked(console.error)).toHaveBeenCalledWith( + '[cloud-agent] Could not return the attachments of a refused send', + expect.any(Error) + ); + }); +}); + +describe('workerRefusalStatus', () => { + it('reads the status the worker answered with, and nothing else', () => { + expect(workerRefusalStatus(refusal(400))).toBe(400); + expect(workerRefusalStatus(new TRPCClientError('fetch failed'))).toBeUndefined(); + expect(workerRefusalStatus(new Error('boom'))).toBeUndefined(); + expect(workerRefusalStatus(undefined)).toBeUndefined(); + }); +}); diff --git a/apps/web/src/lib/cloud-agent-next/release-refused-send.ts b/apps/web/src/lib/cloud-agent-next/release-refused-send.ts new file mode 100644 index 00000000000..e66bb3c6d1a --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/release-refused-send.ts @@ -0,0 +1,59 @@ +import { TRPCClientError } from '@trpc/client'; +import { + releaseUnadmittedUploads, + type DispatchAttachments, +} from '@/lib/r2/cloud-agent-pending-uploads'; + +/** The http status the worker answered a send with, when it answered at all. */ +export function workerRefusalStatus(error: unknown): number | undefined { + if (!(error instanceof TRPCClientError)) return undefined; + const status = (error.data as { httpStatus?: unknown } | undefined)?.httpStatus; + return typeof status === 'number' ? status : undefined; +} + +export type RefusedSend = { + attachments: DispatchAttachments | undefined; + error: unknown; + kiloUserId: string; + /** The object keys this dispatch's link moved to 'linked'. */ + linkedObjectKeys: string[]; +}; + +type ReleaseDependencies = { + release?: ( + kiloUserId: string, + messageUuid: string, + linkedObjectKeys: string[] + ) => Promise; +}; + +/** + * Give back the attachment rows of a send the worker refused. + * + * A refusal is the worker's own error response to this send, received here, so + * the turn was answered and not admitted: the worker refuses before it admits, + * and an admitted send answers with a queued result. A transport failure, a + * timeout or a server-side status proves nothing, so those keep the rows linked; + * the same holds for a send this request never linked anything for, which is the + * retry of a message an earlier dispatch already holds. + * + * The send's own error is what the caller reports; this helper never throws. + */ +export async function releaseRefusedSend( + input: RefusedSend, + { release = releaseUnadmittedUploads }: ReleaseDependencies = {} +): Promise { + const { attachments, error, kiloUserId, linkedObjectKeys } = input; + const status = workerRefusalStatus(error); + if (status === undefined || status < 400 || status >= 500) return; + if (!attachments || linkedObjectKeys.length === 0) return; + + try { + const released = await release(kiloUserId, attachments.path, linkedObjectKeys); + if (released > 0) { + console.info('[cloud-agent] Returned the attachments of a refused send', { released }); + } + } catch (error) { + console.error('[cloud-agent] Could not return the attachments of a refused send', error); + } +} diff --git a/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.test.ts b/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.test.ts deleted file mode 100644 index 042853a5257..00000000000 --- a/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.test.ts +++ /dev/null @@ -1,178 +0,0 @@ -import { beforeEach, describe, expect, it, jest } from '@jest/globals'; - -import type { DispatchAttachments } from '@/lib/r2/cloud-agent-pending-uploads'; - -import { releaseUnadmittedSend } from './release-unadmitted-send'; - -const mockReleaseUnadmittedUploads = jest.fn( - async (_kiloUserId: string, _attachments: DispatchAttachments) => 1 -); - -const messageId = 'msg_0123456789abABCDEFGHIJKLMN'; -const attachments = { path: 'f7c1f9a0-0000-4000-8000-000000000001', files: ['a.bin'] }; - -const sessionId = 'workspace_00000000-0000-4000-8000-000000000002'; - -function admittedResult() { - return { cloudAgentSessionId: sessionId, messageId, status: 'queued' as const }; -} - -function buildLookup(results: readonly (Error | null | 'admitted')[]) { - let calls = 0; - const getMessageResult = jest.fn( - async (_input: { cloudAgentSessionId: string; messageId: string }) => { - const result = results[Math.min(calls, results.length - 1)]; - calls += 1; - if (result instanceof Error) throw result; - return result === null ? null : admittedResult(); - } - ); - return { getMessageResult }; -} - -const dependencies = { delay: async () => {}, release: mockReleaseUnadmittedUploads }; - -describe('releaseUnadmittedSend', () => { - beforeEach(() => { - jest.restoreAllMocks(); - jest.clearAllMocks(); - jest.spyOn(console, 'error').mockImplementation(() => {}); - jest.spyOn(console, 'info').mockImplementation(() => {}); - }); - - it('returns the rows when both lookups report no such message', async () => { - const client = buildLookup([null, null]); - - await releaseUnadmittedSend( - { - attachments, - client, - cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', - kiloUserId: 'user-1', - messageId, - }, - dependencies - ); - - expect(client.getMessageResult).toHaveBeenCalledTimes(2); - expect(client.getMessageResult).toHaveBeenCalledWith({ - cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', - messageId, - }); - expect(mockReleaseUnadmittedUploads).toHaveBeenCalledWith('user-1', attachments); - }); - - it('keeps the rows when the message exists at either lookup', async () => { - const admittedFirst = buildLookup(['admitted']); - await releaseUnadmittedSend( - { - attachments, - client: admittedFirst, - cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', - kiloUserId: 'user-1', - messageId, - }, - dependencies - ); - expect(admittedFirst.getMessageResult).toHaveBeenCalledTimes(1); - - const admittedLater = buildLookup([null, 'admitted']); - await releaseUnadmittedSend( - { - attachments, - client: admittedLater, - cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', - kiloUserId: 'user-1', - messageId, - }, - dependencies - ); - expect(admittedLater.getMessageResult).toHaveBeenCalledTimes(2); - expect(mockReleaseUnadmittedUploads).not.toHaveBeenCalled(); - }); - - it('keeps the rows when a lookup fails', async () => { - const unavailable = buildLookup([new Error('worker unavailable')]); - - await expect( - releaseUnadmittedSend( - { - attachments, - client: unavailable, - cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', - kiloUserId: 'user-1', - messageId, - }, - dependencies - ) - ).resolves.toBeUndefined(); - - expect(jest.mocked(console.error)).toHaveBeenCalledWith( - '[cloud-agent] Could not confirm a failed send was not admitted', - expect.any(Error) - ); - expect(mockReleaseUnadmittedUploads).not.toHaveBeenCalled(); - }); - - it('keeps the rows when the second lookup is the one that fails', async () => { - const secondFails = buildLookup([null, new Error('worker unavailable')]); - - await releaseUnadmittedSend( - { - attachments, - client: secondFails, - cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', - kiloUserId: 'user-1', - messageId, - }, - dependencies - ); - - expect(secondFails.getMessageResult).toHaveBeenCalledTimes(2); - expect(mockReleaseUnadmittedUploads).not.toHaveBeenCalled(); - }); - - it('stays silent when the ledger cannot return the rows', async () => { - const client = buildLookup([null, null]); - const release = jest.fn(async () => { - throw new Error('database unavailable'); - }); - - await expect( - releaseUnadmittedSend( - { - attachments, - client, - cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', - kiloUserId: 'user-1', - messageId, - }, - { delay: async () => {}, release } - ) - ).resolves.toBeUndefined(); - - expect(release).toHaveBeenCalledTimes(1); - expect(jest.mocked(console.error)).toHaveBeenCalledWith( - '[cloud-agent] Could not return the attachments of an unadmitted send', - expect.any(Error) - ); - }); - - it('does nothing for a send without attachments', async () => { - const client = buildLookup([null]); - - await releaseUnadmittedSend( - { - attachments: undefined, - client, - cloudAgentSessionId: 'workspace_00000000-0000-4000-8000-000000000002', - kiloUserId: 'user-1', - messageId, - }, - dependencies - ); - - expect(client.getMessageResult).not.toHaveBeenCalled(); - expect(mockReleaseUnadmittedUploads).not.toHaveBeenCalled(); - }); -}); diff --git a/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.ts b/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.ts deleted file mode 100644 index 198bcfa5887..00000000000 --- a/apps/web/src/lib/cloud-agent-next/release-unadmitted-send.ts +++ /dev/null @@ -1,71 +0,0 @@ -import type { CloudAgentNextClient } from './cloud-agent-client'; -import { - releaseUnadmittedUploads, - type DispatchAttachments, -} from '@/lib/r2/cloud-agent-pending-uploads'; - -/** Time between the first lookup and its confirmation. */ -const confirmationDelayMs = 1_000; - -type MessageResultLookup = Pick; - -export type UnadmittedSend = { - attachments: DispatchAttachments | undefined; - client: MessageResultLookup; - cloudAgentSessionId: string; - kiloUserId: string; - messageId: string; -}; - -type ReleaseDependencies = { - delay?: (milliseconds: number) => Promise; - release?: (kiloUserId: string, attachments: DispatchAttachments) => Promise; -}; - -const sleep = (milliseconds: number): Promise => { - const { promise, resolve } = Promise.withResolvers(); - setTimeout(resolve, milliseconds); - return promise; -}; - -/** - * Give a failed send's attachment rows back to the abandoned-upload lifecycle, - * but only when the worker states that it never admitted the turn. - * - * The worker records an admitted message before it answers the send, and its - * session object serves one request at a time, so a lookup that still reports - * no such message is a statement that the turn did not happen. The lookup runs - * twice, one confirmation apart, because a send that failed in transport can - * still be on its way. Every other answer keeps the rows linked: a message the - * worker holds must never lose its attachments, and an unavailable lookup is - * not proof that the turn was refused. - * - * The send's own error is what the caller reports; this helper never throws. - */ -export async function releaseUnadmittedSend( - input: UnadmittedSend, - { delay = sleep, release = releaseUnadmittedUploads }: ReleaseDependencies = {} -): Promise { - const { attachments, client, cloudAgentSessionId, kiloUserId, messageId } = input; - if (!attachments || attachments.files.length === 0) return; - - try { - for (let attempt = 0; attempt < 2; attempt += 1) { - if (attempt > 0) await delay(confirmationDelayMs); - const result = await client.getMessageResult({ cloudAgentSessionId, messageId }); - if (result !== null) return; - } - } catch (error) { - console.error('[cloud-agent] Could not confirm a failed send was not admitted', error); - return; - } - - try { - const released = await release(kiloUserId, attachments); - if (released > 0) { - console.info('[cloud-agent] Returned the attachments of an unadmitted send', { released }); - } - } catch (error) { - console.error('[cloud-agent] Could not return the attachments of an unadmitted send', error); - } -} diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts index 17b4c02a4c8..649e11e5862 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts @@ -909,7 +909,7 @@ describe('cloud-agent pending-upload ledger', () => { }); const linked = await linkPendingUploads(kiloUserId, messageUuid, [objectKey]); - expect(linked).toEqual({ unlinkedObjectKeys: [] }); + expect(linked).toEqual({ linkedObjectKeys: [objectKey], unlinkedObjectKeys: [] }); const summary = await reapAbandonedUploads(); @@ -945,10 +945,11 @@ describe('cloud-agent pending-upload ledger', () => { status: 'linked', }); - const released = await releaseUnadmittedUploads(kiloUserId, { - path: unadmittedMessage, - files: [`${unadmittedAttachment}.bin`], - }); + const released = await releaseUnadmittedUploads( + kiloUserId, + unadmittedMessage, + [buildObjectKey(kiloUserId, unadmittedMessage, unadmittedAttachment)] + ); expect(released).toBe(1); const rows = await db @@ -975,10 +976,9 @@ describe('cloud-agent pending-upload ledger', () => { status: 'linked', }); - await releaseUnadmittedUploads(kiloUserId, { - path: messageUuid, - files: [`${attachment}.bin`], - }); + await releaseUnadmittedUploads(kiloUserId, messageUuid, [ + buildObjectKey(kiloUserId, messageUuid, attachment), + ]); const summary = await reapAbandonedUploads(); expect(summary.reaped).toBe(1); @@ -1007,7 +1007,7 @@ describe('cloud-agent pending-upload ledger', () => { await reapAbandonedUploads(); const linked = await linkPendingUploads(kiloUserId, messageUuid, [objectKey]); - expect(linked).toEqual({ unlinkedObjectKeys: [objectKey] }); + expect(linked).toEqual({ linkedObjectKeys: [], unlinkedObjectKeys: [objectKey] }); const [after] = await db .select({ status: cloud_agent_pending_uploads.status }) .from(cloud_agent_pending_uploads) @@ -1073,11 +1073,13 @@ describe('cloud-agent pending-upload ledger', () => { expiresAt: future, }); - await linkAttachmentsForDispatch(kiloUserId, { + const linkedKeys = await linkAttachmentsForDispatch(kiloUserId, { path: messageUuid, files: [`${attachment}.bin`, legacyFile], }); + expect(linkedKeys).toEqual([buildObjectKey(kiloUserId, messageUuid, attachment)]); + const [after] = await db .select({ status: cloud_agent_pending_uploads.status }) .from(cloud_agent_pending_uploads) @@ -1085,6 +1087,36 @@ describe('cloud-agent pending-upload ledger', () => { expect(after?.status).toBe('linked'); }); + it('links nothing for a dispatch whose rows an earlier send already holds', async () => { + const kiloUserId = `ca-user-${crypto.randomUUID()}`; + const messageUuid = crypto.randomUUID(); + const attachment = crypto.randomUUID(); + const future = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString(); + + const row = await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: attachment, + expiresAt: future, + status: 'linked', + }); + + // The composer reuses one attachment path across sends, so a second send + // finds the rows of the first one already linked and links none of them. + const linkedKeys = await linkAttachmentsForDispatch(kiloUserId, { + path: messageUuid, + files: [`${attachment}.bin`], + }); + + expect(linkedKeys).toEqual([]); + const [after] = await db + .select({ status: cloud_agent_pending_uploads.status }) + .from(cloud_agent_pending_uploads) + .where(eq(cloud_agent_pending_uploads.id, row)); + expect(after?.status).toBe('linked'); + expect(mockSend).not.toHaveBeenCalled(); + }); + it('leaves every row of a refused message reapable', async () => { const kiloUserId = `ca-user-${crypto.randomUUID()}`; const messageUuid = crypto.randomUUID(); diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts index d0d9b4ca2c6..c11e96c93b2 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.ts @@ -157,6 +157,8 @@ export async function admitPendingUpload({ } export type LinkPendingUploadsResult = { + /** Object keys this call moved from 'pending' to 'linked', one per row. */ + linkedObjectKeys: string[]; /** * Requested keys whose row the message does not hold: 'reaped', so the object * behind them is gone or being deleted, or 'pending' again after a released @@ -184,9 +186,10 @@ export function attachmentObjectKeys( * transaction finds every requested key whose row is not 'linked', which covers * a row a concurrent release or claim took first and a row a failed delete put * back for a later attempt. Rows whose object keys are not in the supplied set - * stay 'pending', so a partial key set leaves unmatched uploads to be reaped. A - * row already 'linked' belongs to an earlier dispatch of the same message and is - * not reported. + * stay 'pending', so a partial key set leaves unmatched uploads to be reaped. It + * also returns the keys this call linked, which is the set one dispatch may hand + * back if the worker later refuses its turn; a row already 'linked' belongs to an + * earlier dispatch and is not in either list. */ async function linkPendingUploadsInTransaction( tx: Parameters[0]>[0], @@ -194,9 +197,9 @@ async function linkPendingUploadsInTransaction( messageUuid: string, objectKeys: string[] ): Promise { - if (objectKeys.length === 0) return { unlinkedObjectKeys: [] }; + if (objectKeys.length === 0) return { linkedObjectKeys: [], unlinkedObjectKeys: [] }; - await tx + const linked = await tx .update(cloud_agent_pending_uploads) .set({ status: 'linked' }) .where( @@ -206,7 +209,8 @@ async function linkPendingUploadsInTransaction( eq(cloud_agent_pending_uploads.status, 'pending'), inArray(cloud_agent_pending_uploads.object_key, objectKeys) ) - ); + ) + .returning({ object_key: cloud_agent_pending_uploads.object_key }); // Read back every requested key the message does not hold. A row is left // 'reaped' when a release or a claim took it, and is returned to 'pending' when @@ -224,7 +228,10 @@ async function linkPendingUploadsInTransaction( ) ); - return { unlinkedObjectKeys: rows.map(row => row.object_key) }; + return { + linkedObjectKeys: linked.map(row => row.object_key), + unlinkedObjectKeys: rows.map(row => row.object_key), + }; } /** Link a message's attachment rows and report the ones it does not hold. */ @@ -233,7 +240,7 @@ export async function linkPendingUploads( messageUuid: string, objectKeys: string[] ): Promise { - if (objectKeys.length === 0) return { unlinkedObjectKeys: [] }; + if (objectKeys.length === 0) return { linkedObjectKeys: [], unlinkedObjectKeys: [] }; return db.transaction(tx => linkPendingUploadsInTransaction(tx, kiloUserId, messageUuid, objectKeys) @@ -251,17 +258,22 @@ export const CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE = * refused message keeps every row as it was and a retry after the user * re-attaches the file links them all. Keys with no ledger row belong to an * uploader that predates the ledger and are left alone. + * + * Returns the object keys this call moved to 'linked'. The composer's attachment + * path is stable across sends from one page, so a later dispatch that finds the + * rows already linked links nothing; this list is what lets one dispatch hand + * back exactly its own rows. */ export async function linkAttachmentsForDispatch( kiloUserId: string, attachments: DispatchAttachments | undefined -): Promise { - if (!attachments || attachments.files.length === 0) return; +): Promise { + if (!attachments || attachments.files.length === 0) return []; - await db.transaction(async tx => { + return db.transaction(async tx => { // The link and the refusal share one transaction, so a message whose // dispatch is refused keeps all of its rows reapable. - const { unlinkedObjectKeys } = await linkPendingUploadsInTransaction( + const { linkedObjectKeys, unlinkedObjectKeys } = await linkPendingUploadsInTransaction( tx, kiloUserId, attachments.path, @@ -273,22 +285,24 @@ export async function linkAttachmentsForDispatch( message: CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE, }); } + return linkedObjectKeys; }); } /** - * Hand one message's attachment rows back to the abandoned-upload lifecycle - * after the worker stated it never admitted the turn. Only the rows this message - * still holds are moved back, so a row another message holds or the reaper - * already took keeps its state. The rows keep their original lease, so the - * reaper claims them on its next run once that lease lapses. Returns the number - * of rows handed back. + * Hand back exactly the rows one dispatch linked, after the worker refused that + * turn. The caller supplies the keys the link returned, so an attachment an + * earlier accepted send holds keeps its state even though the composer reuses + * the path. The rows keep their original lease, so the abandoned-upload reaper + * claims them on its next run once that lease lapses. Returns the number of rows + * handed back. */ export async function releaseUnadmittedUploads( kiloUserId: string, - attachments: DispatchAttachments + messageUuid: string, + linkedObjectKeys: string[] ): Promise { - if (attachments.files.length === 0) return 0; + if (linkedObjectKeys.length === 0) return 0; const released = await db .update(cloud_agent_pending_uploads) @@ -296,12 +310,9 @@ export async function releaseUnadmittedUploads( .where( and( eq(cloud_agent_pending_uploads.kilo_user_id, kiloUserId), - eq(cloud_agent_pending_uploads.message_uuid, attachments.path), + eq(cloud_agent_pending_uploads.message_uuid, messageUuid), eq(cloud_agent_pending_uploads.status, 'linked'), - inArray( - cloud_agent_pending_uploads.object_key, - attachmentObjectKeys(kiloUserId, attachments) - ) + inArray(cloud_agent_pending_uploads.object_key, linkedObjectKeys) ) ) .returning({ id: cloud_agent_pending_uploads.id }); diff --git a/apps/web/src/routers/cloud-agent-next-router.test.ts b/apps/web/src/routers/cloud-agent-next-router.test.ts index a721b1f2cac..e3d2d379c85 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -41,7 +41,7 @@ import { TRPCError } from '@trpc/server'; import { getHTTPStatusCodeFromError } from '@trpc/server/http'; import type { verifyUserOwnsSessionV2ByCloudAgentId } from '@/lib/cloud-agent/session-ownership'; import type * as StreamTicketModule from '@/lib/cloud-agent/stream-ticket'; -import type { releaseUnadmittedSend } from '@/lib/cloud-agent-next/release-unadmitted-send'; +import type { releaseRefusedSend } from '@/lib/cloud-agent-next/release-refused-send'; type AttachmentReference = { path: string; files: string[] }; @@ -158,7 +158,7 @@ const mockIsFeatureFlagEnabledOrDevelopment = jest.fn<(flagName: string, distinctId: string) => Promise>(); const mockVerifyUserOwnsSessionV2ByCloudAgentId = jest.fn(); -const mockReleaseUnadmittedSend = jest.fn(async () => {}); +const mockReleaseRefusedSend = jest.fn(async () => {}); const mockSignPreviewTicket = jest.fn(); const mockGetBalanceForUser = jest.fn<(user: User) => Promise<{ balance: number }>>(); const mockFetchGitHubRepositoriesForUser = jest.fn< @@ -212,8 +212,8 @@ jest.mock('@/lib/cloud-agent-next/worktree-chat', () => ({ createWorktreeChat: mockCreateWorktreeChat, })); -jest.mock('@/lib/cloud-agent-next/release-unadmitted-send', () => ({ - releaseUnadmittedSend: mockReleaseUnadmittedSend, +jest.mock('@/lib/cloud-agent-next/release-refused-send', () => ({ + releaseRefusedSend: mockReleaseRefusedSend, })); jest.mock('@kilocode/web-shared/lib/trpc/min-version', () => ({ @@ -890,7 +890,7 @@ describe('cloudAgentNextRouter send attachment ledger', () => { expect(mockSendMessage).not.toHaveBeenCalled(); // The refusal rolls the link back, so there is nothing to hand back. - expect(mockReleaseUnadmittedSend).not.toHaveBeenCalled(); + expect(mockReleaseRefusedSend).not.toHaveBeenCalled(); }); it('links the attachment rows before the worker accepts the send', async () => { @@ -907,7 +907,7 @@ describe('cloudAgentNextRouter send attachment ledger', () => { ).resolves.toMatchObject({ messageId: SEND_MESSAGE_ID }); expect(mockSendMessage).toHaveBeenCalledWith(expect.objectContaining({ attachments })); - expect(mockReleaseUnadmittedSend).not.toHaveBeenCalled(); + expect(mockReleaseRefusedSend).not.toHaveBeenCalled(); expect(await readStatus()).toBe('linked'); }); @@ -925,12 +925,13 @@ describe('cloudAgentNextRouter send attachment ledger', () => { }) ).rejects.toThrow('worker unavailable'); - expect(mockReleaseUnadmittedSend).toHaveBeenCalledWith({ + expect(mockReleaseRefusedSend).toHaveBeenCalledWith({ attachments, - client: expect.anything(), - cloudAgentSessionId: 'agent_123', + error: expect.any(TRPCClientError), kiloUserId: userId, - messageId: SEND_MESSAGE_ID, + linkedObjectKeys: attachments.files.map( + file => `${userId}/cloud-agent/${attachments.path}/${file}` + ), }); expect(await readStatus()).toBe('linked'); }); diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index 6ac6c64321f..ed160bd8cc2 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -8,7 +8,7 @@ import { } from '@/lib/cloud-agent-next/cloud-agent-client'; import { computeCloudAgentNextBalanceCheckEligibility } from '@/lib/cloud-agent-next/balance-check-eligibility'; import { isPublicRepoUrlEnabled } from '@/lib/cloud-agent-next/public-repo-url'; -import { releaseUnadmittedSend } from '@/lib/cloud-agent-next/release-unadmitted-send'; +import { releaseRefusedSend } from '@/lib/cloud-agent-next/release-refused-send'; import { rethrowAsTerminalError } from '@/lib/cloud-agent-next/terminal-errors'; import { createWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; import { assertSessionWorktree } from '@/lib/cloud-agent-next/worktree-review-access'; @@ -371,6 +371,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ // Tokens are refreshed inside cloud-agent-next (GitHub App installation // for GitHub, GIT_TOKEN_SERVICE for managed GitLab). const { attachments, images } = input; + let linkedObjectKeys: string[] = []; try { // Link the attachment rows before the worker can accept the turn: the // link is what keeps their objects out of the abandoned-upload reaper, @@ -378,7 +379,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ // could report an accepted send as failed. Only ATTACHMENT rows are // admitted (images are never admitted; only // generateCloudAgentAttachmentUploadUrl admits). - await linkAttachmentsForDispatch(ctx.user.id, attachments); + linkedObjectKeys = await linkAttachmentsForDispatch(ctx.user.id, attachments); } catch (error) { rethrowCloudAgentError(error); } @@ -392,15 +393,14 @@ export const cloudAgentNextRouter = createTRPCRouter({ messageId: input.messageId, }); } catch (error) { - // A send can fail after the worker accepted the turn, so the rows are - // only returned to the reaper when the worker states that it never - // admitted this message. - await releaseUnadmittedSend({ + // A send can fail after the worker accepted the turn, so only the rows + // this dispatch linked are handed back, and only when the worker itself + // refused the turn. + await releaseRefusedSend({ attachments, - client, - cloudAgentSessionId: input.cloudAgentSessionId, + error, kiloUserId: ctx.user.id, - messageId: input.messageId, + linkedObjectKeys, }); rethrowCloudAgentError(error); } diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts index b9d27d92342..a49d948621b 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts @@ -289,13 +289,13 @@ jest.mock('@/lib/r2/cloud-agent-attachments', () => ({ })); jest.mock('@/lib/r2/cloud-agent-pending-uploads', () => ({ - linkAttachmentsForDispatch: jest.fn(async () => {}), + linkAttachmentsForDispatch: jest.fn(async () => ['user-1/cloud-agent/linked/a.md']), linkPendingUploads: jest.fn(), releasePendingUploads: jest.fn(), })); -jest.mock('@/lib/cloud-agent-next/release-unadmitted-send', () => ({ - releaseUnadmittedSend: jest.fn(async () => {}), +jest.mock('@/lib/cloud-agent-next/release-refused-send', () => ({ + releaseRefusedSend: jest.fn(async () => {}), })); jest.mock('@kilocode/web-shared/routers/organizations/utils', () => { @@ -1037,18 +1037,22 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { ); }); - it('asks the Worker whether a failed send was admitted before keeping the rows', async () => { + it('hands a refused send the rows this dispatch linked', async () => { const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); const attachments = { path: '12345678-1234-4234-9234-123456789abc', files: ['87654321-4321-4321-8321-cba987654321.md'], }; - const releaseUnadmittedSend = jest.mocked( - jest.requireMock<{ releaseUnadmittedSend: (input: unknown) => Promise }>( - '@/lib/cloud-agent-next/release-unadmitted-send' - ).releaseUnadmittedSend + const releaseRefusedSend = jest.mocked( + jest.requireMock<{ releaseRefusedSend: (input: unknown) => Promise }>( + '@/lib/cloud-agent-next/release-refused-send' + ).releaseRefusedSend + ); + mockSendMessage.mockRejectedValueOnce( + new TRPCClientError('refused', { + result: { error: { code: -32600, message: 'refused', data: { httpStatus: 400 } } }, + }) ); - mockSendMessage.mockRejectedValueOnce(new TRPCClientError('worker unavailable')); await expect( caller.sendMessage({ @@ -1058,14 +1062,13 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { payload: { type: 'prompt', prompt: 'Read notes', mode: 'code', model: 'test' }, attachments, }) - ).rejects.toThrow('worker unavailable'); + ).rejects.toThrow('refused'); - expect(releaseUnadmittedSend).toHaveBeenCalledWith({ + expect(releaseRefusedSend).toHaveBeenCalledWith({ attachments, - client: expect.anything(), - cloudAgentSessionId: 'agent_123', + error: expect.any(TRPCClientError), kiloUserId: 'user-1', - messageId: SEND_MESSAGE_ID, + linkedObjectKeys: ['user-1/cloud-agent/linked/a.md'], }); }); diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index 569222a8550..62ba321cf1f 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -8,7 +8,7 @@ import { } from '@/lib/cloud-agent-next/cloud-agent-client'; import { computeCloudAgentNextBalanceCheckEligibility } from '@/lib/cloud-agent-next/balance-check-eligibility'; import { isPublicRepoUrlEnabled } from '@/lib/cloud-agent-next/public-repo-url'; -import { releaseUnadmittedSend } from '@/lib/cloud-agent-next/release-unadmitted-send'; +import { releaseRefusedSend } from '@/lib/cloud-agent-next/release-refused-send'; import { rethrowAsTerminalError } from '@/lib/cloud-agent-next/terminal-errors'; import { createWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; import { assertSessionWorktree } from '@/lib/cloud-agent-next/worktree-review-access'; @@ -524,6 +524,7 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ // for GitHub, GIT_TOKEN_SERVICE for managed GitLab). organizationId is // consumed by the membership middleware; it is not forwarded. const attachments = input.attachments; + let linkedObjectKeys: string[] = []; try { // Link the attachment rows before the worker can accept the turn: the // link is what keeps their objects out of the abandoned-upload reaper, @@ -531,7 +532,7 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ // could report an accepted send as failed. Only ATTACHMENT rows are // admitted (images are never admitted; only // generateCloudAgentAttachmentUploadUrl admits). - await linkAttachmentsForDispatch(ctx.user.id, attachments); + linkedObjectKeys = await linkAttachmentsForDispatch(ctx.user.id, attachments); } catch (error) { rethrowCloudAgentError(error); } @@ -545,15 +546,14 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ attachments: attachments ?? input.images, }); } catch (error) { - // A send can fail after the worker accepted the turn, so the rows are - // only returned to the reaper when the worker states that it never - // admitted this message. - await releaseUnadmittedSend({ + // A send can fail after the worker accepted the turn, so only the rows + // this dispatch linked are handed back, and only when the worker itself + // refused the turn. + await releaseRefusedSend({ attachments, - client, - cloudAgentSessionId: input.cloudAgentSessionId, + error, kiloUserId: ctx.user.id, - messageId: input.messageId, + linkedObjectKeys, }); rethrowCloudAgentError(error); } From 2ce47034bcd3911d59398307249ac35e12011f7f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Sat, 24 Oct 2026 22:10:20 +0000 Subject: [PATCH 11/13] style(web): format the ledger test --- apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts index 649e11e5862..1efcf44124c 100644 --- a/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts +++ b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts @@ -945,11 +945,9 @@ describe('cloud-agent pending-upload ledger', () => { status: 'linked', }); - const released = await releaseUnadmittedUploads( - kiloUserId, - unadmittedMessage, - [buildObjectKey(kiloUserId, unadmittedMessage, unadmittedAttachment)] - ); + const released = await releaseUnadmittedUploads(kiloUserId, unadmittedMessage, [ + buildObjectKey(kiloUserId, unadmittedMessage, unadmittedAttachment), + ]); expect(released).toBe(1); const rows = await db From a04fa791b100e5e2946724b2002ce298dbc07e6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Sat, 24 Oct 2026 22:26:26 +0000 Subject: [PATCH 12/13] fix(web): read the refusal status from the errors the client returns The send path rewrites some worker answers into its own error classes, so the release now reads the status from any error that carries one and keeps the rows for a transport failure and a server-side status. --- .../release-refused-send.test.ts | 46 +++++++++++++++++++ .../cloud-agent-next/release-refused-send.ts | 34 ++++++++++---- 2 files changed, 70 insertions(+), 10 deletions(-) diff --git a/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts b/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts index fa2b675d156..79bcd777649 100644 --- a/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts +++ b/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts @@ -1,6 +1,10 @@ import { beforeEach, describe, expect, it, jest } from '@jest/globals'; import { TRPCClientError } from '@trpc/client'; +import { + CloudAgentBillingError, + InsufficientCreditsError, +} from '@/lib/cloud-agent-next/cloud-agent-client'; import { releaseRefusedSend, workerRefusalStatus } from './release-refused-send'; const mockRelease = jest.fn( @@ -47,6 +51,44 @@ describe('releaseRefusedSend', () => { expect(mockRelease).not.toHaveBeenCalled(); }); + it('returns the rows when the client rewrote a credits refusal', async () => { + await releaseRefusedSend( + { + attachments, + error: new InsufficientCreditsError(), + kiloUserId: 'user-1', + linkedObjectKeys, + }, + dependencies + ); + + expect(mockRelease).toHaveBeenCalledTimes(1); + }); + + it('returns the rows for a billing refusal but keeps them for a billing outage', async () => { + await releaseRefusedSend( + { + attachments, + error: new CloudAgentBillingError({ kind: 'billing' } as never, 402), + kiloUserId: 'user-1', + linkedObjectKeys, + }, + dependencies + ); + expect(mockRelease).toHaveBeenCalledTimes(1); + + await releaseRefusedSend( + { + attachments, + error: new CloudAgentBillingError({ kind: 'billing' } as never, 503), + kiloUserId: 'user-1', + linkedObjectKeys, + }, + dependencies + ); + expect(mockRelease).toHaveBeenCalledTimes(1); + }); + it('keeps the rows when the send failed without a worker answer', async () => { await releaseRefusedSend( { @@ -102,6 +144,10 @@ describe('releaseRefusedSend', () => { describe('workerRefusalStatus', () => { it('reads the status the worker answered with, and nothing else', () => { expect(workerRefusalStatus(refusal(400))).toBe(400); + expect(workerRefusalStatus(new InsufficientCreditsError())).toBe(402); + expect(workerRefusalStatus(new CloudAgentBillingError({ kind: 'billing' } as never, 409))).toBe( + 409 + ); expect(workerRefusalStatus(new TRPCClientError('fetch failed'))).toBeUndefined(); expect(workerRefusalStatus(new Error('boom'))).toBeUndefined(); expect(workerRefusalStatus(undefined)).toBeUndefined(); diff --git a/apps/web/src/lib/cloud-agent-next/release-refused-send.ts b/apps/web/src/lib/cloud-agent-next/release-refused-send.ts index e66bb3c6d1a..45aa87d177a 100644 --- a/apps/web/src/lib/cloud-agent-next/release-refused-send.ts +++ b/apps/web/src/lib/cloud-agent-next/release-refused-send.ts @@ -4,11 +4,26 @@ import { type DispatchAttachments, } from '@/lib/r2/cloud-agent-pending-uploads'; -/** The http status the worker answered a send with, when it answered at all. */ +/** + * The http status a send failure carries, when the worker answered it. + * + * `client.sendMessage` rewrites several worker answers before this sees them: a + * credits refusal becomes `InsufficientCreditsError`, a billing refusal becomes + * `CloudAgentBillingError`, and those carry `httpStatus` themselves instead of + * the tRPC shape. + */ export function workerRefusalStatus(error: unknown): number | undefined { - if (!(error instanceof TRPCClientError)) return undefined; - const status = (error.data as { httpStatus?: unknown } | undefined)?.httpStatus; - return typeof status === 'number' ? status : undefined; + const direct = (error as { httpStatus?: unknown } | null | undefined)?.httpStatus; + if (typeof direct === 'number') return direct; + + if (error instanceof TRPCClientError) { + const data = error.data as { httpStatus?: unknown } | undefined; + const shapeData = error.shape?.data as { httpStatus?: unknown } | undefined; + const status = data?.httpStatus ?? shapeData?.httpStatus; + return typeof status === 'number' ? status : undefined; + } + + return undefined; } export type RefusedSend = { @@ -30,12 +45,11 @@ type ReleaseDependencies = { /** * Give back the attachment rows of a send the worker refused. * - * A refusal is the worker's own error response to this send, received here, so - * the turn was answered and not admitted: the worker refuses before it admits, - * and an admitted send answers with a queued result. A transport failure, a - * timeout or a server-side status proves nothing, so those keep the rows linked; - * the same holds for a send this request never linked anything for, which is the - * retry of a message an earlier dispatch already holds. + * A refusal is a 4xx answer to this send: the worker refuses before it admits, + * and an admitted send answers with a queued result, so a refused turn never + * existed. A transport failure, a timeout or a server-side status proves + * nothing, so those keep the rows linked; the same holds for a send that linked + * no rows here, which is the retry of an attachment an earlier dispatch holds. * * The send's own error is what the caller reports; this helper never throws. */ From f0cca6b36eae96eb7499033e8c8bdd403ce2d4a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= <4062236+iscekic@users.noreply.github.com> Date: Sat, 24 Oct 2026 22:34:18 +0000 Subject: [PATCH 13/13] fix(web): follow the status the client kept on the cause A non-JSON answer is wrapped with the original error as its cause, so the release reads the status along that chain too and still keeps the rows when no status is present. --- .../release-refused-send.test.ts | 17 +++++++++++ .../cloud-agent-next/release-refused-send.ts | 29 ++++++++++++------- 2 files changed, 36 insertions(+), 10 deletions(-) diff --git a/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts b/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts index 79bcd777649..d323dc5d8d1 100644 --- a/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts +++ b/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts @@ -89,6 +89,19 @@ describe('releaseRefusedSend', () => { expect(mockRelease).toHaveBeenCalledTimes(1); }); + it('returns the rows for a status the client kept only on the cause', async () => { + const normalized = new Error('Cloud agent returned a non-JSON error response', { + cause: refusal(413), + }); + + await releaseRefusedSend( + { attachments, error: normalized, kiloUserId: 'user-1', linkedObjectKeys }, + dependencies + ); + + expect(mockRelease).toHaveBeenCalledTimes(1); + }); + it('keeps the rows when the send failed without a worker answer', async () => { await releaseRefusedSend( { @@ -148,6 +161,10 @@ describe('workerRefusalStatus', () => { expect(workerRefusalStatus(new CloudAgentBillingError({ kind: 'billing' } as never, 409))).toBe( 409 ); + expect(workerRefusalStatus(new Error('non-JSON', { cause: refusal(429) }))).toBe(429); + expect( + workerRefusalStatus(new Error('non-JSON', { cause: new TRPCClientError('fetch failed') })) + ).toBeUndefined(); expect(workerRefusalStatus(new TRPCClientError('fetch failed'))).toBeUndefined(); expect(workerRefusalStatus(new Error('boom'))).toBeUndefined(); expect(workerRefusalStatus(undefined)).toBeUndefined(); diff --git a/apps/web/src/lib/cloud-agent-next/release-refused-send.ts b/apps/web/src/lib/cloud-agent-next/release-refused-send.ts index 45aa87d177a..c4bbdaa3088 100644 --- a/apps/web/src/lib/cloud-agent-next/release-refused-send.ts +++ b/apps/web/src/lib/cloud-agent-next/release-refused-send.ts @@ -9,20 +9,29 @@ import { * * `client.sendMessage` rewrites several worker answers before this sees them: a * credits refusal becomes `InsufficientCreditsError`, a billing refusal becomes - * `CloudAgentBillingError`, and those carry `httpStatus` themselves instead of - * the tRPC shape. + * `CloudAgentBillingError`, and a non-JSON answer becomes a plain error whose + * cause is the original. Those carry the status themselves or keep it on the + * cause, so the status is read along that chain as well as from the tRPC shape. */ export function workerRefusalStatus(error: unknown): number | undefined { - const direct = (error as { httpStatus?: unknown } | null | undefined)?.httpStatus; - if (typeof direct === 'number') return direct; + // The client wraps some answers and keeps the original as the cause, so the + // status may sit one or two links down the chain. + let current: unknown = error; + for (let depth = 0; depth < 4; depth += 1) { + if (current === null || typeof current !== 'object') return undefined; - if (error instanceof TRPCClientError) { - const data = error.data as { httpStatus?: unknown } | undefined; - const shapeData = error.shape?.data as { httpStatus?: unknown } | undefined; - const status = data?.httpStatus ?? shapeData?.httpStatus; - return typeof status === 'number' ? status : undefined; - } + const direct = (current as { httpStatus?: unknown }).httpStatus; + if (typeof direct === 'number') return direct; + + if (current instanceof TRPCClientError) { + const data = current.data as { httpStatus?: unknown } | undefined; + const shapeData = current.shape?.data as { httpStatus?: unknown } | undefined; + const status = data?.httpStatus ?? shapeData?.httpStatus; + if (typeof status === 'number') return status; + } + current = (current as { cause?: unknown }).cause; + } return undefined; }