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 9d15716b28f..74c4fdcf8e0 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/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..d323dc5d8d1 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/release-refused-send.test.ts @@ -0,0 +1,172 @@ +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( + 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('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('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( + { + 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 InsufficientCreditsError())).toBe(402); + 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 new file mode 100644 index 00000000000..c4bbdaa3088 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/release-refused-send.ts @@ -0,0 +1,82 @@ +import { TRPCClientError } from '@trpc/client'; +import { + releaseUnadmittedUploads, + type DispatchAttachments, +} from '@/lib/r2/cloud-agent-pending-uploads'; + +/** + * 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 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 { + // 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; + + 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; +} + +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 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. + */ +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/r2/cloud-agent-pending-uploads.test.ts b/apps/web/src/lib/r2/cloud-agent-pending-uploads.test.ts index 5571bea3184..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 @@ -1,10 +1,12 @@ import { afterEach, beforeAll, describe, expect, it, jest } from '@jest/globals'; import type { admitPendingUpload as AdmitPendingUpload, + linkAttachmentsForDispatch as LinkAttachmentsForDispatch, linkPendingUploads as LinkPendingUploads, 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>(); @@ -21,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, @@ -32,10 +34,12 @@ 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 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`; @@ -76,10 +80,12 @@ 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; purgeUserPendingUploads = pendingUploadsModule.purgeUserPendingUploads; reapAbandonedUploads = pendingUploadsModule.reapAbandonedUploads; releasePendingUploads = pendingUploadsModule.releasePendingUploads; + releaseUnadmittedUploads = pendingUploadsModule.releaseUnadmittedUploads; }); afterEach(async () => { @@ -794,7 +800,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(); @@ -855,7 +861,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(); @@ -869,9 +875,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 @@ -886,4 +892,262 @@ 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({ linkedObjectKeys: [objectKey], unlinkedObjectKeys: [] }); + + 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('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, unadmittedMessage, [ + buildObjectKey(kiloUserId, unadmittedMessage, unadmittedAttachment), + ]); + + 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, messageUuid, [ + buildObjectKey(kiloUserId, messageUuid, attachment), + ]); + 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(); + 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({ linkedObjectKeys: [], unlinkedObjectKeys: [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 future = new Date(Date.now() + 24 * 60 * 60 * 1000).toISOString(); + + const row = await insertPendingRow({ + kiloUserId, + messageUuid, + attachmentId: attachment, + expiresAt: future, + }); + + 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) + .where(eq(cloud_agent_pending_uploads.id, row)); + 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(); + 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: goneAttachment, + expiresAt: past, + status: 'reaped', + }); + + await expect( + linkAttachmentsForDispatch(kiloUserId, { + path: messageUuid, + 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 afd8037e0ee..c11e96c93b2 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, asc, count, eq, gt, inArray, like, not, sql, type SQL } from 'drizzle-orm'; +import { and, asc, count, eq, gt, 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 { @@ -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, @@ -156,33 +156,170 @@ 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 + * delete put the row back for a later attempt. A message referencing one of + * these must not be dispatched. + */ + 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 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, 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. 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], + kiloUserId: string, + messageUuid: string, + objectKeys: string[] +): Promise { + if (objectKeys.length === 0) return { linkedObjectKeys: [], unlinkedObjectKeys: [] }; + + const linked = 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) + ) + ) + .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 + // 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), + ne(cloud_agent_pending_uploads.status, 'linked'), + inArray(cloud_agent_pending_uploads.object_key, objectKeys) + ) + ); + + 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. */ export async function linkPendingUploads( kiloUserId: string, messageUuid: string, objectKeys: string[] -): Promise { - if (objectKeys.length === 0) return; +): Promise { + if (objectKeys.length === 0) return { linkedObjectKeys: [], unlinkedObjectKeys: [] }; - 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) - ) - ); + return db.transaction(tx => + linkPendingUploadsInTransaction(tx, kiloUserId, messageUuid, objectKeys) + ); +} + +/** 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 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. + * + * 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 []; + + 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 { linkedObjectKeys, unlinkedObjectKeys } = await linkPendingUploadsInTransaction( + tx, + kiloUserId, + attachments.path, + attachmentObjectKeys(kiloUserId, attachments) + ); + if (unlinkedObjectKeys.length > 0) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: CLOUD_AGENT_ATTACHMENT_EXPIRED_MESSAGE, + }); + } + return linkedObjectKeys; }); } +/** + * 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, + messageUuid: string, + linkedObjectKeys: string[] +): Promise { + if (linkedObjectKeys.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, messageUuid), + eq(cloud_agent_pending_uploads.status, 'linked'), + inArray(cloud_agent_pending_uploads.object_key, linkedObjectKeys) + ) + ) + .returning({ id: cloud_agent_pending_uploads.id }); + + return released.length; +} + export type ReapAbandonedUploadsSummary = { reaped: number; }; @@ -216,10 +353,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 rows claimed and 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 rows + * claimed and the number of objects actually deleted. */ async function claimAndDeleteUploadObjects( claim: SQL | undefined @@ -331,8 +471,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 @@ -349,6 +490,9 @@ export async function reapAbandonedUploads(): Promise Promise>(); const mockVerifyUserOwnsSessionV2ByCloudAgentId = jest.fn(); +const mockReleaseRefusedSend = jest.fn(async () => {}); const mockSignPreviewTicket = jest.fn(); const mockGetBalanceForUser = jest.fn<(user: User) => Promise<{ balance: number }>>(); const mockFetchGitHubRepositoriesForUser = jest.fn< @@ -202,6 +212,10 @@ jest.mock('@/lib/cloud-agent-next/worktree-chat', () => ({ createWorktreeChat: mockCreateWorktreeChat, })); +jest.mock('@/lib/cloud-agent-next/release-refused-send', () => ({ + releaseRefusedSend: mockReleaseRefusedSend, +})); + 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' })), @@ -277,7 +291,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 }; @@ -463,6 +477,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), @@ -585,6 +600,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'); @@ -601,6 +617,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, }); @@ -617,6 +634,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, }); @@ -634,6 +652,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'prompt', prompt: 'Follow up on this', @@ -662,6 +681,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -681,6 +701,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -700,6 +721,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -717,6 +739,7 @@ describe('cloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -786,6 +809,152 @@ 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(); + // The refusal rolls the link back, so there is nothing to hand back. + expect(mockReleaseRefusedSend).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(mockReleaseRefusedSend).not.toHaveBeenCalled(); + 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(mockReleaseRefusedSend).toHaveBeenCalledWith({ + attachments, + error: expect.any(TRPCClientError), + kiloUserId: userId, + linkedObjectKeys: attachments.files.map( + file => `${userId}/cloud-agent/${attachments.path}/${file}` + ), + }); + 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(); @@ -1545,6 +1714,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 bfb68d42cc8..ed160bd8cc2 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 { 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'; @@ -71,7 +72,10 @@ import { generateCloudAgentAttachmentDownloadUrl, generateImageUploadUrl, } from '@/lib/r2/cloud-agent-attachments'; -import { linkPendingUploads, releasePendingUploads } from '@/lib/r2/cloud-agent-pending-uploads'; +import { + linkAttachmentsForDispatch, + 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'; @@ -264,7 +267,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, @@ -272,21 +283,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); } @@ -374,32 +370,38 @@ 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 { - 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). + linkedObjectKeys = await linkAttachmentsForDispatch(ctx.user.id, attachments); + } catch (error) { + rethrowCloudAgentError(error); + } + + try { + 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) { + // 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, + error, + kiloUserId: ctx.user.id, + linkedObjectKeys, + }); rethrowCloudAgentError(error); } }), 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 784f6b74d15..a41a76e16ed 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.test.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.test.ts @@ -672,6 +672,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/cloud-agent-next-schemas.ts b/apps/web/src/routers/cloud-agent-next-schemas.ts index deb51f89aff..e9622dbf323 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.ts @@ -583,7 +583,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/cloud-agent-next-send-schema.test.ts b/apps/web/src/routers/cloud-agent-next-send-schema.test.ts index 4a98df0d4c5..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 @@ -7,6 +7,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, @@ -38,6 +39,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`), }); @@ -47,6 +49,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`), }); @@ -56,6 +59,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); @@ -64,6 +68,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); @@ -72,6 +77,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); @@ -80,6 +86,7 @@ describe('baseSendMessageNextSchema', () => { it('accepts a text-only prompt', () => { const result = baseSendMessageNextSchema.safeParse({ cloudAgentSessionId: 'cloud-session-1', + messageId: MESSAGE_ID, payload: promptPayload(atCap), }); expect(result.success).toBe(true); @@ -88,6 +95,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`), }); @@ -97,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); @@ -105,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); @@ -113,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); 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 a6d11b92aeb..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 @@ -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; @@ -285,10 +289,15 @@ jest.mock('@/lib/r2/cloud-agent-attachments', () => ({ })); jest.mock('@/lib/r2/cloud-agent-pending-uploads', () => ({ + linkAttachmentsForDispatch: jest.fn(async () => ['user-1/cloud-agent/linked/a.md']), linkPendingUploads: jest.fn(), releasePendingUploads: jest.fn(), })); +jest.mock('@/lib/cloud-agent-next/release-refused-send', () => ({ + releaseRefusedSend: 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'); @@ -342,7 +351,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 }; @@ -973,6 +982,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'); @@ -990,6 +1000,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, }); @@ -1000,6 +1011,67 @@ 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(mockSendMessage).toHaveBeenCalledTimes(1); + expect(linkAttachmentsForDispatch.mock.invocationCallOrder[0]).toBeLessThan( + mockSendMessage.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY + ); + }); + + 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 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 } } }, + }) + ); + + 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('refused'); + + expect(releaseRefusedSend).toHaveBeenCalledWith({ + attachments, + error: expect.any(TRPCClientError), + kiloUserId: 'user-1', + linkedObjectKeys: ['user-1/cloud-agent/linked/a.md'], + }); + }); + it('normalizes legacy image requests to canonical Worker attachments', async () => { const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); const images = { @@ -1010,6 +1082,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, }); @@ -1028,6 +1101,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', @@ -1060,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: '' }, }); @@ -1083,6 +1158,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -1106,6 +1182,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -1124,6 +1201,7 @@ describe('organizationCloudAgentNextRouter attachment forwarding', () => { await caller.sendMessage({ organizationId: ORGANIZATION_ID, cloudAgentSessionId: 'agent_123', + messageId: SEND_MESSAGE_ID, payload: { type: 'command', command: 'review', arguments: '' }, }); @@ -2201,6 +2279,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 a34d4d0fd83..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,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 { 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'; @@ -76,7 +77,10 @@ import { generateCloudAgentAttachmentUploadUrl, generateImageUploadUrl, } from '@/lib/r2/cloud-agent-attachments'; -import { linkPendingUploads, releasePendingUploads } from '@/lib/r2/cloud-agent-pending-uploads'; +import { + linkAttachmentsForDispatch, + 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'; @@ -397,7 +400,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, @@ -406,21 +416,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); } @@ -528,32 +523,38 @@ 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; + let linkedObjectKeys: string[] = []; 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). + linkedObjectKeys = 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 ?? generateMessageId(), - attachments: input.attachments ?? input.images, + messageId: input.messageId, + attachments: 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) { + // 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, + error, + kiloUserId: ctx.user.id, + linkedObjectKeys, + }); rethrowCloudAgentError(error); } }), 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 72a6f0dfea0..a4e9172c9a4 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 800777a26bd..470e0b327fb 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts @@ -481,16 +481,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 dcb1e88b4e8..1e2058438d4 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 be56e6b8880..c825bfed19c 100644 --- a/packages/cloud-agent-sdk/src/transport.ts +++ b/packages/cloud-agent-sdk/src/transport.ts @@ -193,7 +193,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; diff --git a/services/kilo-mcp/catalog.json b/services/kilo-mcp/catalog.json index 70ef08ecb26..e543df024e8 100644 --- a/services/kilo-mcp/catalog.json +++ b/services/kilo-mcp/catalog.json @@ -7653,15 +7653,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", @@ -7711,7 +7704,8 @@ }, "required": [ "cloudAgentSessionId", - "payload" + "payload", + "messageId" ] }, "tags": [ @@ -18668,15 +18662,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", @@ -18727,7 +18714,8 @@ "required": [ "organizationId", "cloudAgentSessionId", - "payload" + "payload", + "messageId" ] }, "tags": [