diff --git a/services/session-ingest/src/dos/SessionIngestDO.test.ts b/services/session-ingest/src/dos/SessionIngestDO.test.ts index 50b33a51b4..75588abda7 100644 --- a/services/session-ingest/src/dos/SessionIngestDO.test.ts +++ b/services/session-ingest/src/dos/SessionIngestDO.test.ts @@ -159,7 +159,10 @@ describe('SessionIngestDO ingest ordering', () => { ); await Promise.all(waitUntilPromises); - expect(result.changes).toEqual([{ name: 'title', value: 'Hello' }]); + expect(result).toMatchObject({ + accepted: true, + changes: [{ name: 'title', value: 'Hello' }], + }); expect(deleteObject).toHaveBeenCalledWith(['items/old']); expect(operations.indexOf('meta:title:Hello')).toBeLessThan(operations.indexOf('r2Delete')); expect(metaValues.get('title')).toBe('Newer'); @@ -314,6 +317,30 @@ describe('SessionIngestDO session-ready push', () => { expect(sendSessionReadyNotification).not.toHaveBeenCalled(); }); + it('reports a deleted ingest and cleans up caller R2 references', async () => { + const { durableObject, rows } = makeHarness(); + rows.set('deleted', { value: 'true' }); + const deleteObject = vi.mocked( + ( + durableObject as unknown as { + env: { SESSION_INGEST_R2: { delete: ReturnType } }; + } + ).env.SESSION_INGEST_R2.delete + ); + + const result = await durableObject.ingest( + [{ type: 'message', data: { id: 'msg_deleted' } }], + 'usr_deleted', + 'ses_deleted', + 1, + 1, + { 'message/msg_deleted': 'items/deleted' } + ); + + expect(result).toEqual({ accepted: false, reason: 'deleted', changes: [] }); + expect(deleteObject).toHaveBeenCalledWith(['items/deleted']); + }); + it('does not push from ingest, even for a parentless session record', async () => { const { durableObject, sendSessionReadyNotification, settle } = makeHarness(); diff --git a/services/session-ingest/src/dos/SessionIngestDO.ts b/services/session-ingest/src/dos/SessionIngestDO.ts index 7c19a8f7e8..7adc3a7b13 100644 --- a/services/session-ingest/src/dos/SessionIngestDO.ts +++ b/services/session-ingest/src/dos/SessionIngestDO.ts @@ -87,6 +87,10 @@ const INGEST_META_EXTRACTORS: Array<{ type Changes = Array<{ name: ExtractableMetaKey; value: string | null }>; +export type IngestResult = + | { accepted: true; changes: Changes } + | { accepted: false; reason: 'deleted'; changes: never[] }; + type IngestLifecycleEvent = | { type: 'session_open' } | { @@ -136,9 +140,7 @@ export class SessionIngestDO extends DurableObject { ingestVersion = 0, ingestedAt?: number, r2References?: Record - ): Promise<{ - changes: Changes; - }> { + ): Promise { const deletedRow = this.db .select({ value: ingestMeta.value }) .from(ingestMeta) @@ -152,7 +154,7 @@ export class SessionIngestDO extends DurableObject { await this.env.SESSION_INGEST_R2.delete(keys); } } - return { changes: [] }; + return { accepted: false, reason: 'deleted', changes: [] }; } writeIngestMetaIfChanged(this.db, { key: 'kiloUserId', incomingValue: kiloUserId }); @@ -301,6 +303,7 @@ export class SessionIngestDO extends DurableObject { } return { + accepted: true, changes, }; } diff --git a/services/session-ingest/src/ingest/item-extractor.ts b/services/session-ingest/src/ingest/item-extractor.ts new file mode 100644 index 0000000000..a32de77319 --- /dev/null +++ b/services/session-ingest/src/ingest/item-extractor.ts @@ -0,0 +1,252 @@ +import { Tokenizer, TokenParser, TokenType } from '@streamparser/json'; + +import { MAX_SINGLE_ITEM_BYTES } from '../util/ingest-limits'; + +type ItemExtractorOptions = { + logErrors?: boolean; + logOversizedItems?: boolean; + validateStructure?: boolean; +}; + +type ContainerState = + | { type: 'object'; state: 'key_or_end' | 'key' | 'colon' | 'value' | 'comma_or_end' } + | { type: 'array'; state: 'value_or_end' | 'value' | 'comma_or_end' }; + +export function createItemExtractor(r2Key: string, options: ItemExtractorOptions = {}) { + const pending: Record[] = []; + let parseError: Error | null = null; + let skippedItemCount = 0; + let oversizedItemCount = 0; + let rootState: 'value' | 'complete' = 'value'; + const containers: ContainerState[] = []; + + // Depth: 0=before root, 1=root object, 2=$.data array, 3+=inside an item + let depth = 0; + let pendingKey: string | undefined; + let foundDataArray = false; + let dataArray: 'present' | 'missing' | 'wrong_type' = 'missing'; + let itemStartOffset = 0; + let skippingItem = false; + let itemParser: TokenParser | null = null; + + function startItemParser() { + itemParser = new TokenParser({ paths: ['$'], keepStack: false }); + itemParser.onValue = ({ value, stack }) => { + if (stack.length === 0 && value != null) { + pending.push(value as Record); + } + }; + itemParser.onError = (err: Error) => { + if (options.logErrors !== false) { + console.error('TokenParser error in queue consumer', { r2Key, error: err.message }); + } + }; + } + + function setParseError(message: string) { + if (parseError) return; + parseError = new Error(message); + if (options.logErrors !== false) { + console.error('Tokenizer error in queue consumer', { r2Key, error: message }); + } + } + + function consumeValue() { + const parent = containers.at(-1); + if (!parent) { + if (rootState !== 'value') return false; + rootState = 'complete'; + return true; + } + if (parent.type === 'object') { + if (parent.state !== 'value') return false; + parent.state = 'comma_or_end'; + return true; + } + if (parent.state !== 'value' && parent.state !== 'value_or_end') return false; + parent.state = 'comma_or_end'; + return true; + } + + function validateToken(token: TokenType): boolean { + const current = containers.at(-1); + + if (current?.type === 'object') { + if ( + (current.state === 'key_or_end' || current.state === 'key') && + token === TokenType.STRING + ) { + current.state = 'colon'; + return true; + } + if (current.state === 'colon' && token === TokenType.COLON) { + current.state = 'value'; + return true; + } + if (current.state === 'comma_or_end' && token === TokenType.COMMA) { + current.state = 'key'; + return true; + } + if ( + (current.state === 'key_or_end' || current.state === 'comma_or_end') && + token === TokenType.RIGHT_BRACE + ) { + containers.pop(); + return true; + } + } else if (current?.type === 'array') { + if (current.state === 'comma_or_end' && token === TokenType.COMMA) { + current.state = 'value'; + return true; + } + if ( + (current.state === 'value_or_end' || current.state === 'comma_or_end') && + token === TokenType.RIGHT_BRACKET + ) { + containers.pop(); + return true; + } + } + + const isScalar = + token === TokenType.STRING || + token === TokenType.NUMBER || + token === TokenType.TRUE || + token === TokenType.FALSE || + token === TokenType.NULL; + if (isScalar) return consumeValue(); + + if (token === TokenType.LEFT_BRACE || token === TokenType.LEFT_BRACKET) { + if (!consumeValue()) return false; + containers.push( + token === TokenType.LEFT_BRACE + ? { type: 'object', state: 'key_or_end' } + : { type: 'array', state: 'value_or_end' } + ); + return true; + } + + return false; + } + + const tokenizer = new Tokenizer(); + tokenizer.onToken = ({ token, value, offset }) => { + if (options.validateStructure && !validateToken(token)) { + setParseError(`Unexpected JSON token ${TokenType[token]}`); + return; + } + + const isOpen = token === TokenType.LEFT_BRACE || token === TokenType.LEFT_BRACKET; + const isClose = token === TokenType.RIGHT_BRACE || token === TokenType.RIGHT_BRACKET; + + if (skippingItem) { + if (isOpen) depth++; + if (isClose) { + depth--; + if (depth === 2) skippingItem = false; + } + return; + } + + if (foundDataArray && depth >= 3) { + if (offset - itemStartOffset > MAX_SINGLE_ITEM_BYTES) { + skippedItemCount += 1; + oversizedItemCount += 1; + if (options.logOversizedItems !== false) { + console.warn('Skipping oversized item in queue consumer (byte budget exceeded)', { + r2Key, + bytesConsumed: offset - itemStartOffset, + maxBytes: MAX_SINGLE_ITEM_BYTES, + }); + } + skippingItem = true; + itemParser = null; + if (isOpen) depth++; + if (isClose) depth--; + if (depth === 2) skippingItem = false; + return; + } + + itemParser?.write({ token, value }); + if (isOpen) depth++; + if (isClose) { + depth--; + if (depth === 2) itemParser = null; + } + return; + } + + if (isOpen) { + const opensDataArray = + depth === 1 && token === TokenType.LEFT_BRACKET && pendingKey === 'data'; + if (foundDataArray && depth === 2 && token === TokenType.LEFT_BRACKET) { + skippedItemCount += 1; + skippingItem = true; + depth++; + return; + } + if (depth === 1 && pendingKey === 'data') { + dataArray = opensDataArray ? 'present' : 'wrong_type'; + } + depth++; + + if (foundDataArray && depth === 3 && token === TokenType.LEFT_BRACE) { + itemStartOffset = offset; + startItemParser(); + itemParser?.write({ token, value }); + return; + } + + if (opensDataArray) { + foundDataArray = true; + pendingKey = undefined; + return; + } + + pendingKey = undefined; + return; + } + + if (isClose) { + if (foundDataArray && depth === 2 && token === TokenType.RIGHT_BRACKET) { + foundDataArray = false; + } + depth--; + return; + } + + if (foundDataArray && depth === 2 && token !== TokenType.COMMA && token !== TokenType.COLON) { + skippedItemCount += 1; + } + + if (depth === 1 && pendingKey === 'data' && token !== TokenType.COLON) { + dataArray = 'wrong_type'; + } + if (depth === 1 && token === TokenType.STRING) { + pendingKey = value as string; + } else if (token !== TokenType.COLON) { + pendingKey = undefined; + } + }; + + tokenizer.onError = (err: Error) => { + if (options.logErrors !== false) { + console.error('Tokenizer error in queue consumer', { r2Key, error: err.message }); + } + parseError = err; + }; + + return { + tokenizer, + pending, + getParseError: () => parseError, + getDataArray: () => dataArray, + getSkippedItemCount: () => skippedItemCount, + getOversizedItemCount: () => oversizedItemCount, + isComplete: () => + (!options.validateStructure || (rootState === 'complete' && containers.length === 0)) && + depth === 0 && + itemParser === null && + !skippingItem, + }; +} diff --git a/services/session-ingest/src/ingest/metadata.ts b/services/session-ingest/src/ingest/metadata.ts new file mode 100644 index 0000000000..044e0ab1c0 --- /dev/null +++ b/services/session-ingest/src/ingest/metadata.ts @@ -0,0 +1,229 @@ +import { and, eq, sql } from 'drizzle-orm'; +import { getWorkerDb } from '@kilocode/db/client'; +import { cli_sessions_v2 } from '@kilocode/db/schema'; +import { normalizeGitUrl } from '@kilocode/worker-utils'; + +import type { Env } from '../env'; +import { mapSessionEventRow, notifyUserSessionEvent } from '../session-events'; +import { SessionStatusSchema } from '../types/user-connection-protocol'; + +type SessionMetadataUpdates = Partial< + Pick< + typeof cli_sessions_v2.$inferInsert, + | 'title' + | 'created_on_platform' + | 'organization_id' + | 'git_url' + | 'git_branch' + | 'status' + | 'status_updated_at' + > +>; + +export function computeSessionMetadataUpdates( + mergedChanges: Map, + now: () => string = () => new Date().toISOString() +): SessionMetadataUpdates { + const updates: SessionMetadataUpdates = {}; + + if (mergedChanges.has('title')) updates.title = mergedChanges.get('title') ?? null; + if (mergedChanges.has('platform')) { + const platform = mergedChanges.get('platform') ?? null; + if (platform !== null) updates.created_on_platform = platform; + } + if (mergedChanges.has('orgId')) updates.organization_id = mergedChanges.get('orgId') ?? null; + if (mergedChanges.has('gitUrl')) { + const gitUrl = mergedChanges.get('gitUrl') ?? null; + updates.git_url = gitUrl === null ? null : normalizeGitUrl(gitUrl); + } + if (mergedChanges.has('gitBranch')) updates.git_branch = mergedChanges.get('gitBranch') ?? null; + if (mergedChanges.has('status')) { + updates.status = mergedChanges.get('status') ?? null; + updates.status_updated_at = now(); + } + + return updates; +} + +export async function applyMetadataChanges( + env: Env, + kiloUserId: string, + sessionId: string, + mergedChanges: Map, + ctx: ExecutionContext +): Promise { + if (mergedChanges.size === 0) return; + + const db = getWorkerDb(env.HYPERDRIVE.connectionString); + const status = mergedChanges.has('status') ? (mergedChanges.get('status') ?? null) : undefined; + const updates = computeSessionMetadataUpdates(mergedChanges); + const parentSessionId = mergedChanges.has('parentId') + ? (mergedChanges.get('parentId') ?? null) + : undefined; + const changedNonStatus = + mergedChanges.has('title') || + mergedChanges.has('platform') || + mergedChanges.has('orgId') || + mergedChanges.has('gitUrl') || + mergedChanges.has('gitBranch') || + parentSessionId !== undefined; + + const notification = await db.transaction(async tx => { + const statusChange = + status === undefined + ? { changed: false, previousStatus: null } + : await (async () => { + const [statusRow] = await tx + .select({ status: cli_sessions_v2.status }) + .from(cli_sessions_v2) + .where( + and( + eq(cli_sessions_v2.session_id, sessionId), + eq(cli_sessions_v2.kilo_user_id, kiloUserId) + ) + ) + .limit(1) + .for('update'); + if (!statusRow) return null; + const previousStatus = SessionStatusSchema.nullable().parse(statusRow.status); + return { changed: status !== previousStatus, previousStatus }; + })(); + + if (!statusChange) return null; + + if (Object.keys(updates).length > 0) { + await tx + .update(cli_sessions_v2) + .set(updates) + .where( + and( + eq(cli_sessions_v2.session_id, sessionId), + eq(cli_sessions_v2.kilo_user_id, kiloUserId) + ) + ); + } + + if (parentSessionId !== undefined) { + if (parentSessionId && parentSessionId !== sessionId) { + const parentRows = await tx + .select({ session_id: cli_sessions_v2.session_id }) + .from(cli_sessions_v2) + .where( + and( + eq(cli_sessions_v2.session_id, parentSessionId), + eq(cli_sessions_v2.kilo_user_id, kiloUserId) + ) + ) + .limit(1); + + if (parentRows[0]) { + await tx + .update(cli_sessions_v2) + .set({ parent_session_id: parentSessionId }) + .where( + and( + eq(cli_sessions_v2.session_id, sessionId), + eq(cli_sessions_v2.kilo_user_id, kiloUserId), + sql`${cli_sessions_v2.parent_session_id} IS DISTINCT FROM ${parentSessionId}` + ) + ); + } + } else if (parentSessionId === null) { + await tx + .update(cli_sessions_v2) + .set({ parent_session_id: null }) + .where( + and( + eq(cli_sessions_v2.session_id, sessionId), + eq(cli_sessions_v2.kilo_user_id, kiloUserId), + sql`${cli_sessions_v2.parent_session_id} IS DISTINCT FROM ${parentSessionId}` + ) + ); + } + } + + if (!changedNonStatus && !statusChange.changed) return null; + + const [persistedRow] = await tx + .select({ + session_id: cli_sessions_v2.session_id, + created_at: cli_sessions_v2.created_at, + updated_at: cli_sessions_v2.updated_at, + title: cli_sessions_v2.title, + created_on_platform: cli_sessions_v2.created_on_platform, + organization_id: cli_sessions_v2.organization_id, + git_url: cli_sessions_v2.git_url, + git_branch: cli_sessions_v2.git_branch, + parent_session_id: cli_sessions_v2.parent_session_id, + status: cli_sessions_v2.status, + status_updated_at: cli_sessions_v2.status_updated_at, + }) + .from(cli_sessions_v2) + .where( + and(eq(cli_sessions_v2.session_id, sessionId), eq(cli_sessions_v2.kilo_user_id, kiloUserId)) + ) + .limit(1); + + if (!persistedRow) return null; + + return { + changedNonStatus, + changedStatus: statusChange.changed, + previousStatus: statusChange.previousStatus, + session: mapSessionEventRow(persistedRow), + }; + }); + if (!notification) return; + + if (notification.changedNonStatus) { + notifyUserSessionEvent( + env, + kiloUserId, + { + type: 'session.updated', + data: { + source: 'v2', + session: notification.session, + changedAt: notification.session.updatedAt, + }, + }, + ctx + ); + } + if (notification.changedStatus) { + notifyUserSessionEvent( + env, + kiloUserId, + { + type: 'session.status.updated', + data: { + source: 'v2', + session: notification.session, + previousStatus: notification.previousStatus, + status: notification.session.status, + statusUpdatedAt: notification.session.statusUpdatedAt, + changedAt: notification.session.updatedAt, + }, + }, + ctx + ); + } +} + +export async function flushPartialMetadataChanges( + env: Env, + params: { r2Key: string; kiloUserId: string; sessionId: string }, + mergedChanges: Map, + ctx: ExecutionContext +): Promise { + if (mergedChanges.size === 0) return; + try { + await applyMetadataChanges(env, params.kiloUserId, params.sessionId, mergedChanges, ctx); + } catch (err) { + console.error('Failed to flush partial metadata changes after ingest error', { + r2Key: params.r2Key, + sessionId: params.sessionId, + error: err instanceof Error ? err.message : String(err), + }); + } +} diff --git a/services/session-ingest/src/ingest/stage-and-enqueue.ts b/services/session-ingest/src/ingest/stage-and-enqueue.ts new file mode 100644 index 0000000000..c12ab48268 --- /dev/null +++ b/services/session-ingest/src/ingest/stage-and-enqueue.ts @@ -0,0 +1,27 @@ +import type { Env } from '../env'; +import type { IngestQueueMessage } from '../queue-consumer'; + +type StageAndEnqueueParams = Omit & { + r2Key: string; + ingestedAt?: number; +}; + +export async function stageAndEnqueue( + env: Env, + params: StageAndEnqueueParams, + body: ReadableStream | Uint8Array +): Promise { + await env.SESSION_INGEST_R2.put(params.r2Key, body); + + const message: IngestQueueMessage = { + ...params, + ingestedAt: params.ingestedAt ?? Date.now(), + }; + + try { + await env.INGEST_QUEUE.send(message); + } catch (error) { + await env.SESSION_INGEST_R2.delete(params.r2Key).catch(() => {}); + throw error; + } +} diff --git a/services/session-ingest/src/ingest/validate.test.ts b/services/session-ingest/src/ingest/validate.test.ts new file mode 100644 index 0000000000..88474d23d8 --- /dev/null +++ b/services/session-ingest/src/ingest/validate.test.ts @@ -0,0 +1,174 @@ +import { describe, expect, it } from 'vitest'; + +import { + INGEST_CHUNK_MAX_BYTES, + INGEST_CHUNK_MAX_ITEMS, + MAX_SINGLE_ITEM_BYTES, +} from '../util/ingest-limits'; +import { validateAndParseIngestPayload } from './validate'; + +const encoder = new TextEncoder(); + +function validate(value: unknown) { + return validateAndParseIngestPayload(encoder.encode(JSON.stringify(value))); +} + +describe('validateAndParseIngestPayload', () => { + it('returns parsed valid items and their serialized data sizes', () => { + const items = [ + { type: 'session', data: { title: 'Hello' } }, + { type: 'message', data: { id: 'msg_1', text: 'hi' } }, + ]; + + expect(validate({ data: items })).toEqual({ + ok: true, + items, + dataArray: 'present', + validItemCount: 2, + skippedItemCount: 0, + totalValidItemBytes: + encoder.encode(JSON.stringify(items[0].data)).byteLength + + encoder.encode(JSON.stringify(items[1].data)).byteLength, + maxValidItemBytes: encoder.encode(JSON.stringify(items[1].data)).byteLength, + }); + }); + + it.each([ + [ + 'valid prefix with malformed tail', + '{"data":[{"type":"message","data":{"id":"msg_1"}},broken', + ], + ['truncated body', '{"data":[{"type":"message","data":{"id":"msg_1"}}'], + ])('rejects %s', (_name, body) => { + expect(validateAndParseIngestPayload(encoder.encode(body))).toEqual({ + ok: false, + error: 'malformed_json', + }); + }); + + it.each(['', ' ', '{]', '[}', '{}{}', '{"data":[]} null'])( + 'rejects structurally invalid JSON %j', + body => { + expect(validateAndParseIngestPayload(encoder.encode(body))).toEqual({ + ok: false, + error: 'malformed_json', + }); + } + ); + + it.each([ + ['missing', {}, 'missing'], + ['wrong-shaped', { data: {} }, 'wrong_type'], + ['empty', { data: [] }, 'present'], + ] as const)('distinguishes %s data arrays', (_name, payload, dataArray) => { + expect(validate(payload)).toEqual({ + ok: true, + items: [], + dataArray, + validItemCount: 0, + skippedItemCount: 0, + totalValidItemBytes: 0, + maxValidItemBytes: 0, + }); + }); + + it('skips invalid items without charging count or byte budgets', () => { + const validItem = { type: 'message', data: { id: 'msg_valid' } }; + const result = validate({ + data: [ + validItem, + { type: 'message', data: {} }, + { type: 'unknown', data: { content: 'x'.repeat(1000) } }, + ], + }); + + expect(result).toMatchObject({ + ok: true, + items: [validItem], + validItemCount: 1, + skippedItemCount: 2, + totalValidItemBytes: encoder.encode(JSON.stringify(validItem.data)).byteLength, + }); + }); + + it('counts non-object data array entries as skipped', () => { + expect( + validate({ + data: [null, 42, 'invalid', true, [], { type: 'message', data: { id: 'msg_valid' } }], + }) + ).toMatchObject({ + ok: true, + validItemCount: 1, + skippedItemCount: 5, + }); + }); + + it('counts only valid items when raw item count exceeds the RPC item budget', () => { + const validItems = Array.from({ length: INGEST_CHUNK_MAX_ITEMS }, (_, index) => ({ + type: 'message', + data: { id: `msg_${index}` }, + })); + const result = validate({ data: [{ type: 'message', data: {} }, ...validItems] }); + + expect(result).toMatchObject({ + ok: true, + validItemCount: INGEST_CHUNK_MAX_ITEMS, + skippedItemCount: 1, + }); + }); + + it('reports more than one RPC chunk worth of valid items', () => { + const items = Array.from({ length: INGEST_CHUNK_MAX_ITEMS + 1 }, (_, index) => ({ + type: 'message', + data: { id: `msg_${index}` }, + })); + + expect(validate({ data: items })).toMatchObject({ + ok: true, + validItemCount: INGEST_CHUNK_MAX_ITEMS + 1, + skippedItemCount: 0, + }); + }); + + it('reports an oversized valid item', () => { + const data = { id: 'msg_large', content: 'x'.repeat(2 * 1024 * 1024) }; + const dataBytes = encoder.encode(JSON.stringify(data)).byteLength; + + expect(validate({ data: [{ type: 'message', data }] })).toMatchObject({ + ok: true, + validItemCount: 1, + totalValidItemBytes: dataBytes, + maxValidItemBytes: dataBytes, + }); + }); + + it('reports parser-skipped oversized items as ineligible', () => { + const result = validate({ + data: [ + { type: 'message', data: { id: 'msg_huge', content: 'x'.repeat(MAX_SINGLE_ITEM_BYTES) } }, + ], + }); + + expect(result).toMatchObject({ + ok: true, + validItemCount: 0, + skippedItemCount: 1, + maxValidItemBytes: MAX_SINGLE_ITEM_BYTES + 1, + }); + }); + + it.each([ + ['at', INGEST_CHUNK_MAX_BYTES], + ['over', INGEST_CHUNK_MAX_BYTES + 1], + ])('reports valid-item totals %s the RPC byte budget', (_name, expectedBytes) => { + const prefix = encoder.encode(JSON.stringify({ id: 'msg_boundary', content: '' })).byteLength; + const data = { id: 'msg_boundary', content: 'x'.repeat(expectedBytes - prefix) }; + + expect(validate({ data: [{ type: 'message', data }] })).toMatchObject({ + ok: true, + validItemCount: 1, + totalValidItemBytes: expectedBytes, + maxValidItemBytes: expectedBytes, + }); + }); +}); diff --git a/services/session-ingest/src/ingest/validate.ts b/services/session-ingest/src/ingest/validate.ts new file mode 100644 index 0000000000..490d885c38 --- /dev/null +++ b/services/session-ingest/src/ingest/validate.ts @@ -0,0 +1,61 @@ +import { SessionItemSchema, type SessionDataItem } from '../types/session-sync'; +import { MAX_SINGLE_ITEM_BYTES } from '../util/ingest-limits'; +import { createItemExtractor } from './item-extractor'; + +type DataArrayState = 'present' | 'missing' | 'wrong_type'; + +export type IngestPayloadValidationResult = + | { + ok: true; + items: SessionDataItem[]; + dataArray: DataArrayState; + validItemCount: number; + skippedItemCount: number; + totalValidItemBytes: number; + maxValidItemBytes: number; + } + | { ok: false; error: 'malformed_json' }; + +export function validateAndParseIngestPayload(bytes: Uint8Array): IngestPayloadValidationResult { + const extractor = createItemExtractor('buffered-ingest-validation', { + logErrors: false, + logOversizedItems: false, + validateStructure: true, + }); + extractor.tokenizer.write(bytes); + extractor.tokenizer.end(); + + if (extractor.getParseError() || !extractor.isComplete()) { + return { ok: false, error: 'malformed_json' }; + } + + const encoder = new TextEncoder(); + const items: SessionDataItem[] = []; + let skippedItemCount = extractor.getSkippedItemCount(); + let totalValidItemBytes = 0; + let maxValidItemBytes = 0; + + for (const rawItem of extractor.pending) { + const parsed = SessionItemSchema.safeParse(rawItem); + if (!parsed.success) { + skippedItemCount += 1; + continue; + } + + const itemBytes = encoder.encode(JSON.stringify(parsed.data.data)).byteLength; + items.push(parsed.data); + totalValidItemBytes += itemBytes; + maxValidItemBytes = Math.max(maxValidItemBytes, itemBytes); + } + + return { + ok: true, + items, + dataArray: extractor.getDataArray(), + validItemCount: items.length, + skippedItemCount, + totalValidItemBytes, + maxValidItemBytes: + extractor.getOversizedItemCount() > 0 ? MAX_SINGLE_ITEM_BYTES + 1 : maxValidItemBytes, + }; +} diff --git a/services/session-ingest/src/queue-consumer.test.ts b/services/session-ingest/src/queue-consumer.test.ts index 7fc80c8fcb..3ef5c19345 100644 --- a/services/session-ingest/src/queue-consumer.test.ts +++ b/services/session-ingest/src/queue-consumer.test.ts @@ -37,6 +37,8 @@ vi.mock('./session-events', async importOriginal => { // Mock ingest-limits so we can exercise both streaming and SQLite-row compaction thresholds. vi.mock('./util/ingest-limits', () => ({ + INGEST_CHUNK_MAX_BYTES: 4 * 1024 * 1024, + INGEST_CHUNK_MAX_ITEMS: 128, MAX_INGEST_ITEM_BYTES: 100, MAX_SINGLE_ITEM_BYTES: 500, })); @@ -44,12 +46,8 @@ vi.mock('./util/ingest-limits', () => ({ import { getWorkerDb } from '@kilocode/db/client'; import { getSessionIngestDO } from './dos/SessionIngestDO'; import { notifyUserSessionEvent } from './session-events'; -import { - QUEUE_RETRY_DELAY_SECONDS, - computeSessionMetadataUpdates, - createItemExtractor, - queue, -} from './queue-consumer'; +import { QUEUE_RETRY_DELAY_SECONDS, createItemExtractor, queue } from './queue-consumer'; +import { computeSessionMetadataUpdates } from './ingest/metadata'; const encoder = new TextEncoder(); @@ -132,6 +130,27 @@ describe('createItemExtractor', () => { expect(ext.getParseError()).toBeInstanceOf(Error); }); + it('preserves lexical-only parsing for concatenated roots', () => { + const ext = createItemExtractor('test-key'); + feedAll(ext, '{}{}'); + + expect(ext.getParseError()).toBeNull(); + }); + + it('counts a large array entry once without treating it as an oversized object', () => { + const ext = createItemExtractor('test-key', { logOversizedItems: false }); + feedAll( + ext, + JSON.stringify({ + data: [Array.from({ length: 600 }, () => 'x'), { type: 'message', data: { id: 'msg_1' } }], + }) + ); + + expect(ext.getSkippedItemCount()).toBe(1); + expect(ext.getOversizedItemCount()).toBe(0); + expect(ext.pending).toEqual([{ type: 'message', data: { id: 'msg_1' } }]); + }); + it('ignores non-data top-level keys', () => { const ext = createItemExtractor('test-key'); const payload = JSON.stringify({ @@ -372,6 +391,165 @@ describe('queue', () => { expect(retry).not.toHaveBeenCalled(); }); + it('deletes and acknowledges staging when the session DO is tombstoned', async () => { + const ingest = vi.fn( + async () => ({ accepted: false, reason: 'deleted', changes: [] }) as const + ); + vi.mocked(getSessionIngestDO).mockReturnValue({ ingest } as never); + const limit = vi.fn(async () => [{ session_id: 'ses_tombstoned' }]); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + const transaction = vi.fn(); + vi.mocked(getWorkerDb).mockReturnValue({ + select: vi.fn(() => ({ from })), + transaction, + } as never); + + const deleteObject = vi.fn(async () => undefined); + const env = { + HYPERDRIVE: { connectionString: 'postgres://unused' }, + SESSION_INGEST_R2: { + get: vi.fn( + async () => + new Response( + JSON.stringify({ data: [{ type: 'message', data: { id: 'msg_tombstoned' } }] }) + ) + ), + put: vi.fn(async () => undefined), + delete: deleteObject, + }, + } as never; + const ack = vi.fn(); + const retry = vi.fn(); + + await queue( + { + messages: [ + { + body: { + r2Key: 'staging/tombstoned', + kiloUserId: 'usr_tombstoned', + sessionId: 'ses_tombstoned', + ingestVersion: 1, + ingestedAt: 1, + }, + ack, + retry, + }, + ], + } as never, + env, + { waitUntil: vi.fn() } as unknown as ExecutionContext + ); + + expect(transaction).not.toHaveBeenCalled(); + expect(deleteObject).toHaveBeenCalledWith('staging/tombstoned'); + expect(ack).toHaveBeenCalledTimes(1); + expect(retry).not.toHaveBeenCalled(); + }); + + it('acknowledges a tombstone without parsing a malformed trailing suffix', async () => { + const ingest = vi.fn( + async () => ({ accepted: false, reason: 'deleted', changes: [] }) as const + ); + vi.mocked(getSessionIngestDO).mockReturnValue({ ingest } as never); + const limit = vi.fn(async () => [{ session_id: 'ses_tombstoned_tail' }]); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + vi.mocked(getWorkerDb).mockReturnValue({ select: vi.fn(() => ({ from })) } as never); + + const items = Array.from({ length: 128 }, (_, index) => ({ + type: 'message', + data: { id: `msg_${index}` }, + })); + const body = `{"data":[${items.map(item => JSON.stringify(item)).join(',')},broken`; + const deleteObject = vi.fn(async () => undefined); + const env = { + HYPERDRIVE: { connectionString: 'postgres://unused' }, + SESSION_INGEST_R2: { + get: vi.fn(async () => new Response(body)), + put: vi.fn(async () => undefined), + delete: deleteObject, + }, + } as never; + const ack = vi.fn(); + const retry = vi.fn(); + + await queue( + { + messages: [ + { + body: { + r2Key: 'staging/tombstoned-tail', + kiloUserId: 'usr_tombstoned_tail', + sessionId: 'ses_tombstoned_tail', + ingestVersion: 1, + ingestedAt: 1, + }, + ack, + retry, + }, + ], + } as never, + env, + { waitUntil: vi.fn() } as unknown as ExecutionContext + ); + + expect(ingest).toHaveBeenCalledTimes(1); + expect(deleteObject).toHaveBeenCalledWith('staging/tombstoned-tail'); + expect(ack).toHaveBeenCalledTimes(1); + expect(retry).not.toHaveBeenCalled(); + }); + + it('does not stage a duplicate item after its preceding chunk is tombstoned', async () => { + const ingest = vi.fn( + async () => ({ accepted: false, reason: 'deleted', changes: [] }) as const + ); + vi.mocked(getSessionIngestDO).mockReturnValue({ ingest } as never); + const limit = vi.fn(async () => [{ session_id: 'ses_tombstoned_duplicate' }]); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + vi.mocked(getWorkerDb).mockReturnValue({ select: vi.fn(() => ({ from })) } as never); + + const item = { type: 'message', data: { id: 'msg_duplicate' } }; + const deleteObject = vi.fn(async () => undefined); + const env = { + HYPERDRIVE: { connectionString: 'postgres://unused' }, + SESSION_INGEST_R2: { + get: vi.fn(async () => new Response(JSON.stringify({ data: [item, item] }))), + put: vi.fn(async () => undefined), + delete: deleteObject, + }, + } as never; + const ack = vi.fn(); + const retry = vi.fn(); + + await queue( + { + messages: [ + { + body: { + r2Key: 'staging/tombstoned-duplicate', + kiloUserId: 'usr_tombstoned_duplicate', + sessionId: 'ses_tombstoned_duplicate', + ingestVersion: 1, + ingestedAt: 1, + }, + ack, + retry, + }, + ], + } as never, + env, + { waitUntil: vi.fn() } as unknown as ExecutionContext + ); + + expect(ingest).toHaveBeenCalledTimes(1); + expect(deleteObject).toHaveBeenCalledWith('staging/tombstoned-duplicate'); + expect(ack).toHaveBeenCalledTimes(1); + expect(retry).not.toHaveBeenCalled(); + }); + it('splits a message past the chunk cap into ordered DO ingest calls', async () => { // 129 items exceed INGEST_CHUNK_MAX_ITEMS (128) -> two chunks: 128 then 1. // Both chunks must commit in order; non-empty changes trigger the final metadata flush. diff --git a/services/session-ingest/src/queue-consumer.ts b/services/session-ingest/src/queue-consumer.ts index a56e8d0273..642d753869 100644 --- a/services/session-ingest/src/queue-consumer.ts +++ b/services/session-ingest/src/queue-consumer.ts @@ -1,16 +1,20 @@ -import { eq, and, sql } from 'drizzle-orm'; +import { eq, and } from 'drizzle-orm'; import { getWorkerDb } from '@kilocode/db/client'; import { cli_sessions_v2 } from '@kilocode/db/schema'; -import { Tokenizer, TokenParser, TokenType } from '@streamparser/json'; import type { Env } from './env'; import { SessionItemSchema, type SessionDataItem } from './types/session-sync'; import { getItemIdentity } from './util/compaction'; -import { MAX_INGEST_ITEM_BYTES, MAX_SINGLE_ITEM_BYTES } from './util/ingest-limits'; -import { getSessionIngestDO } from './dos/SessionIngestDO'; -import { withDORetry, normalizeGitUrl } from '@kilocode/worker-utils'; -import { mapSessionEventRow, notifyUserSessionEvent } from './session-events'; -import { SessionStatusSchema } from './types/user-connection-protocol'; +import { + INGEST_CHUNK_MAX_BYTES, + INGEST_CHUNK_MAX_ITEMS, + MAX_INGEST_ITEM_BYTES, +} from './util/ingest-limits'; +import { getSessionIngestDO, type IngestResult } from './dos/SessionIngestDO'; +import { withDORetry } from '@kilocode/worker-utils'; +import { applyMetadataChanges, flushPartialMetadataChanges } from './ingest/metadata'; +export { createItemExtractor } from './ingest/item-extractor'; +import { createItemExtractor } from './ingest/item-extractor'; export interface IngestQueueMessage { r2Key: string; @@ -27,139 +31,6 @@ export const QUEUE_RETRY_DELAY_SECONDS = 5 * 60; // ingest() RPCs to that session's DO instead of one RPC per item. Prod snapshot // (2026-06-03): p99 is ~13 items / ~1.7 MiB, so ~99% fit in one chunk; only // ~0.3% (>4 MiB) split. These caps bound memory and RPC size. -const INGEST_CHUNK_MAX_BYTES = 4 * 1024 * 1024; -const INGEST_CHUNK_MAX_ITEMS = 128; - -/** - * Creates a streaming item extractor that uses a low-level Tokenizer to parse - * items from `$.data[]` one at a time, with a per-item byte budget. - * - * Items within budget get their tokens fed to a fresh TokenParser that builds - * the JS object. Oversized items have their tokens discarded without ever - * materializing a JS object. - * - * Peak memory: one R2 chunk + one parsed item (bounded by MAX_SINGLE_ITEM_BYTES). - */ -export function createItemExtractor(r2Key: string) { - const pending: Record[] = []; - let parseError: Error | null = null; - - // Depth: 0=before root, 1=root object, 2=$.data array, 3+=inside an item - let depth = 0; - let pendingKey: string | undefined; - let foundDataArray = false; - let itemStartOffset = 0; - let skippingItem = false; - let itemParser: TokenParser | null = null; - - function startItemParser() { - itemParser = new TokenParser({ paths: ['$'], keepStack: false }); - itemParser.onValue = ({ value, stack }) => { - if (stack.length === 0 && value != null) { - pending.push(value as Record); - } - }; - itemParser.onError = (err: Error) => { - console.error('TokenParser error in queue consumer', { r2Key, error: err.message }); - }; - } - - const tokenizer = new Tokenizer(); - tokenizer.onToken = ({ token, value, offset }) => { - const isOpen = token === TokenType.LEFT_BRACE || token === TokenType.LEFT_BRACKET; - const isClose = token === TokenType.RIGHT_BRACE || token === TokenType.RIGHT_BRACKET; - - // --- Skipping an oversized item: just track depth to find closing brace --- - if (skippingItem) { - if (isOpen) depth++; - if (isClose) { - depth--; - if (depth === 2) { - skippingItem = false; - } - } - return; - } - - // --- Inside an item (depth >= 3): feed tokens to item parser with byte budget --- - if (foundDataArray && depth >= 3) { - if (offset - itemStartOffset > MAX_SINGLE_ITEM_BYTES) { - console.warn('Skipping oversized item in queue consumer (byte budget exceeded)', { - r2Key, - bytesConsumed: offset - itemStartOffset, - maxBytes: MAX_SINGLE_ITEM_BYTES, - }); - skippingItem = true; - itemParser = null; - if (isOpen) depth++; - if (isClose) depth--; - if (depth === 2) skippingItem = false; // item ended on the trigger token - return; - } - - itemParser?.write({ token, value }); - if (isOpen) depth++; - if (isClose) { - depth--; - if (depth === 2) { - // Item complete — onValue already fired, clean up - itemParser = null; - } - } - return; - } - - // --- Structural tokens outside items --- - if (isOpen) { - depth++; - - // depth just became 3 inside $.data[] with { → item start - if (foundDataArray && depth === 3 && token === TokenType.LEFT_BRACE) { - itemStartOffset = offset; - startItemParser(); - itemParser?.write({ token, value }); - return; - } - - // depth just became 2 with [ after "data" key → found $.data array - if (depth === 2 && token === TokenType.LEFT_BRACKET && pendingKey === 'data') { - foundDataArray = true; - pendingKey = undefined; - return; - } - - pendingKey = undefined; - return; - } - - if (isClose) { - if (foundDataArray && depth === 2 && token === TokenType.RIGHT_BRACKET) { - foundDataArray = false; - } - depth--; - return; - } - - // Track keys at depth 1 (root object properties) to detect "data" - if (depth === 1 && token === TokenType.STRING) { - pendingKey = value as string; - } else if (token !== TokenType.COLON) { - pendingKey = undefined; - } - }; - - tokenizer.onError = (err: Error) => { - console.error('Tokenizer error in queue consumer', { r2Key, error: err.message }); - parseError = err; - }; - - return { - tokenizer, - pending, - getParseError: () => parseError, - }; -} - async function processMessage( env: Env, msg: IngestQueueMessage, @@ -171,7 +42,10 @@ async function processMessage( const mergedChanges = new Map(); try { - await ingestStagedSessionItems(env, msg, body, mergedChanges); + const accepted = await ingestStagedSessionItems(env, msg, body, mergedChanges); + if (accepted) { + await applyMetadataChanges(env, msg.kiloUserId, msg.sessionId, mergedChanges, ctx); + } } catch (err) { // An earlier chunk may have committed to the DO before a later chunk (or the // JSON parse) failed. The DO reports a metadata change only when its stored @@ -182,28 +56,9 @@ async function processMessage( throw err; } - await applyMetadataChanges(env, msg.kiloUserId, msg.sessionId, mergedChanges, ctx); await env.SESSION_INGEST_R2.delete(msg.r2Key); } -async function flushPartialMetadataChanges( - env: Env, - msg: IngestQueueMessage, - mergedChanges: Map, - ctx: ExecutionContext -): Promise { - if (mergedChanges.size === 0) return; - try { - await applyMetadataChanges(env, msg.kiloUserId, msg.sessionId, mergedChanges, ctx); - } catch (err) { - console.error('Failed to flush partial metadata changes after ingest error', { - r2Key: msg.r2Key, - sessionId: msg.sessionId, - error: err instanceof Error ? err.message : String(err), - }); - } -} - async function deleteStagingObjectIfSessionMissing( env: Env, msg: IngestQueueMessage @@ -241,7 +96,7 @@ async function ingestStagedSessionItems( msg: IngestQueueMessage, body: ReadableStream, mergedChanges: Map -): Promise { +): Promise { const chunker = createIngestChunker(env, msg, mergedChanges); const parseError = await streamSessionItems(msg.r2Key, body, rawItem => chunker.stage(rawItem)); @@ -251,12 +106,13 @@ async function ingestStagedSessionItems( // Handle any remaining items not flushed yet. await chunker.flushChunkToSessionDO(); + return chunker.wasAccepted(); } async function streamSessionItems( r2Key: string, body: ReadableStream, - onItem: (rawItem: Record) => Promise + onItem: (rawItem: Record) => Promise ): Promise { const { tokenizer, pending, getParseError } = createItemExtractor(r2Key); const reader = body.getReader(); @@ -275,7 +131,7 @@ async function streamSessionItems( while (pending.length > 0) { const rawItem = pending.shift(); if (!rawItem) break; - await onItem(rawItem); + if (!(await onItem(rawItem))) return null; } if (result.done) break; @@ -330,6 +186,7 @@ function createIngestChunker( const chunkItemIds = new Set(); let chunkR2References: Record = {}; let chunkBytes = 0; + let accepted = true; const flushChunkToSessionDO = async (): Promise => { if (chunk.length === 0) return; @@ -339,17 +196,23 @@ function createIngestChunker( chunkR2References = {}; chunkBytes = 0; - const ingestResult = await withDORetry( + const ingestResult = await withDORetry, IngestResult>( () => getSessionIngestDO(env, { kiloUserId, sessionId }), - stub => stub.ingest(items, kiloUserId, sessionId, ingestVersion, ingestedAt, r2References), + async stub => + stub.ingest(items, kiloUserId, sessionId, ingestVersion, ingestedAt, r2References), 'SessionIngestDO.ingest' ); + if (ingestResult.accepted === false) { + accepted = false; + return; + } for (const change of ingestResult.changes) { mergedChanges.set(change.name, change.value); } }; - const stage = async (rawItem: Record): Promise => { + const stage = async (rawItem: Record): Promise => { + if (!accepted) return false; const parsed = SessionItemSchema.safeParse(rawItem); if (!parsed.success) { console.warn('Skipping invalid item in queue consumer', { @@ -357,7 +220,7 @@ function createIngestChunker( type: rawItem['type'], errors: parsed.error.issues.map(i => i.message), }); - return; + return true; } const item = parsed.data; @@ -368,6 +231,7 @@ function createIngestChunker( if (chunkItemIds.has(item_id)) { await flushChunkToSessionDO(); + if (!accepted) return false; } // Offload data above the DO SQLite row limit to R2; the DO stores a @@ -390,226 +254,10 @@ function createIngestChunker( if (chunk.length >= INGEST_CHUNK_MAX_ITEMS || chunkBytes >= INGEST_CHUNK_MAX_BYTES) { await flushChunkToSessionDO(); } + return accepted; }; - return { stage, flushChunkToSessionDO }; -} - -type SessionMetadataUpdates = Partial< - Pick< - typeof cli_sessions_v2.$inferInsert, - | 'title' - | 'created_on_platform' - | 'organization_id' - | 'git_url' - | 'git_branch' - | 'status' - | 'status_updated_at' - > ->; - -/** - * Build the `cli_sessions_v2` partial update from a set of metadata changes. - * - * `git_url` is passed through `normalizeGitUrl` on write so that the - * `github_branch_pull_requests` cache (keyed on the canonical form) can - * match new sessions without per-read normalization. Status bumps carry - * `status_updated_at = now()`. - */ -export function computeSessionMetadataUpdates( - mergedChanges: Map, - now: () => string = () => new Date().toISOString() -): SessionMetadataUpdates { - const updates: SessionMetadataUpdates = {}; - - if (mergedChanges.has('title')) { - updates.title = mergedChanges.get('title') ?? null; - } - if (mergedChanges.has('platform')) { - const platform = mergedChanges.get('platform') ?? null; - if (platform !== null) updates.created_on_platform = platform; - } - if (mergedChanges.has('orgId')) { - updates.organization_id = mergedChanges.get('orgId') ?? null; - } - if (mergedChanges.has('gitUrl')) { - const gitUrl = mergedChanges.get('gitUrl') ?? null; - updates.git_url = gitUrl === null ? null : normalizeGitUrl(gitUrl); - } - if (mergedChanges.has('gitBranch')) { - updates.git_branch = mergedChanges.get('gitBranch') ?? null; - } - if (mergedChanges.has('status')) { - updates.status = mergedChanges.get('status') ?? null; - updates.status_updated_at = now(); - } - - return updates; -} - -async function applyMetadataChanges( - env: Env, - kiloUserId: string, - sessionId: string, - mergedChanges: Map, - ctx: ExecutionContext -): Promise { - if (mergedChanges.size === 0) return; - - const db = getWorkerDb(env.HYPERDRIVE.connectionString); - const status = mergedChanges.has('status') ? (mergedChanges.get('status') ?? null) : undefined; - const updates = computeSessionMetadataUpdates(mergedChanges); - const parentSessionId = mergedChanges.has('parentId') - ? (mergedChanges.get('parentId') ?? null) - : undefined; - const changedNonStatus = - mergedChanges.has('title') || - mergedChanges.has('platform') || - mergedChanges.has('orgId') || - mergedChanges.has('gitUrl') || - mergedChanges.has('gitBranch') || - parentSessionId !== undefined; - - const notification = await db.transaction(async tx => { - const statusChange = - status === undefined - ? { changed: false, previousStatus: null } - : await (async () => { - const [statusRow] = await tx - .select({ status: cli_sessions_v2.status }) - .from(cli_sessions_v2) - .where( - and( - eq(cli_sessions_v2.session_id, sessionId), - eq(cli_sessions_v2.kilo_user_id, kiloUserId) - ) - ) - .limit(1) - .for('update'); - if (!statusRow) return null; - const previousStatus = SessionStatusSchema.nullable().parse(statusRow.status); - return { changed: status !== previousStatus, previousStatus }; - })(); - - if (!statusChange) return null; - - if (Object.keys(updates).length > 0) { - await tx - .update(cli_sessions_v2) - .set(updates) - .where( - and( - eq(cli_sessions_v2.session_id, sessionId), - eq(cli_sessions_v2.kilo_user_id, kiloUserId) - ) - ); - } - - if (parentSessionId !== undefined) { - if (parentSessionId && parentSessionId !== sessionId) { - const parentRows = await tx - .select({ session_id: cli_sessions_v2.session_id }) - .from(cli_sessions_v2) - .where( - and( - eq(cli_sessions_v2.session_id, parentSessionId), - eq(cli_sessions_v2.kilo_user_id, kiloUserId) - ) - ) - .limit(1); - - if (parentRows[0]) { - await tx - .update(cli_sessions_v2) - .set({ parent_session_id: parentSessionId }) - .where( - and( - eq(cli_sessions_v2.session_id, sessionId), - eq(cli_sessions_v2.kilo_user_id, kiloUserId), - sql`${cli_sessions_v2.parent_session_id} IS DISTINCT FROM ${parentSessionId}` - ) - ); - } - } else if (parentSessionId === null) { - await tx - .update(cli_sessions_v2) - .set({ parent_session_id: null }) - .where( - and( - eq(cli_sessions_v2.session_id, sessionId), - eq(cli_sessions_v2.kilo_user_id, kiloUserId), - sql`${cli_sessions_v2.parent_session_id} IS DISTINCT FROM ${parentSessionId}` - ) - ); - } - } - - if (!changedNonStatus && !statusChange.changed) return null; - - const [persistedRow] = await tx - .select({ - session_id: cli_sessions_v2.session_id, - created_at: cli_sessions_v2.created_at, - updated_at: cli_sessions_v2.updated_at, - title: cli_sessions_v2.title, - created_on_platform: cli_sessions_v2.created_on_platform, - organization_id: cli_sessions_v2.organization_id, - git_url: cli_sessions_v2.git_url, - git_branch: cli_sessions_v2.git_branch, - parent_session_id: cli_sessions_v2.parent_session_id, - status: cli_sessions_v2.status, - status_updated_at: cli_sessions_v2.status_updated_at, - }) - .from(cli_sessions_v2) - .where( - and(eq(cli_sessions_v2.session_id, sessionId), eq(cli_sessions_v2.kilo_user_id, kiloUserId)) - ) - .limit(1); - - if (!persistedRow) return null; - - return { - changedNonStatus, - changedStatus: statusChange.changed, - previousStatus: statusChange.previousStatus, - session: mapSessionEventRow(persistedRow), - }; - }); - if (!notification) return; - - if (notification.changedNonStatus) { - notifyUserSessionEvent( - env, - kiloUserId, - { - type: 'session.updated', - data: { - source: 'v2', - session: notification.session, - changedAt: notification.session.updatedAt, - }, - }, - ctx - ); - } - if (notification.changedStatus) { - notifyUserSessionEvent( - env, - kiloUserId, - { - type: 'session.status.updated', - data: { - source: 'v2', - session: notification.session, - previousStatus: notification.previousStatus, - status: notification.session.status, - statusUpdatedAt: notification.session.statusUpdatedAt, - changedAt: notification.session.updatedAt, - }, - }, - ctx - ); - } + return { stage, flushChunkToSessionDO, wasAccepted: () => accepted }; } export async function queue( diff --git a/services/session-ingest/src/routes/api.test.ts b/services/session-ingest/src/routes/api.test.ts index 902b339c46..9d3556d7d7 100644 --- a/services/session-ingest/src/routes/api.test.ts +++ b/services/session-ingest/src/routes/api.test.ts @@ -40,7 +40,7 @@ type HyperdriveBinding = { connectionString: string }; type TestBindings = { HYPERDRIVE: HyperdriveBinding; - SESSION_INGEST_R2: { put: ReturnType }; + SESSION_INGEST_R2: { put: ReturnType; delete: ReturnType }; INGEST_QUEUE: { send: ReturnType }; NOTIFICATIONS: { sendSessionReadyNotification: ReturnType }; }; @@ -48,7 +48,10 @@ type TestBindings = { function makeTestEnv(): TestBindings { return { HYPERDRIVE: { connectionString: 'postgres://test' }, - SESSION_INGEST_R2: { put: vi.fn(async () => undefined) }, + SESSION_INGEST_R2: { + put: vi.fn(async () => undefined), + delete: vi.fn(async () => undefined), + }, INGEST_QUEUE: { send: vi.fn(async () => undefined) }, NOTIFICATIONS: { sendSessionReadyNotification: vi.fn(async () => ({ dispatched: true })) }, }; @@ -375,6 +378,61 @@ describe('api routes', () => { expect(typeof queueMsg['ingestedAt']).toBe('number'); }); + it('captures ingestedAt after staging completes', async () => { + const { db } = makeDbFakes(); + vi.mocked(getWorkerDb).mockReturnValue(db); + vi.mocked(getSessionAccessCacheDO).mockReturnValue({ has: vi.fn(async () => true) } as never); + + const app = makeApiApp(); + const env = makeTestEnv(); + const operations: string[] = []; + env.SESSION_INGEST_R2.put.mockImplementationOnce(async () => { + operations.push('put'); + }); + const now = vi.spyOn(Date, 'now').mockImplementation(() => { + operations.push('now'); + return 123; + }); + + const response = await app.fetch( + new Request('http://local/session/ses_12345678901234567890123456/ingest', { + method: 'POST', + body: JSON.stringify({ data: [] }), + }), + env + ); + now.mockRestore(); + + expect(response.status).toBe(200); + expect(operations).toEqual(['put', 'now']); + expect(env.INGEST_QUEUE.send).toHaveBeenCalledWith( + expect.objectContaining({ ingestedAt: 123 }) + ); + }); + + it('deletes the staged object when queue enqueue fails', async () => { + const { db } = makeDbFakes(); + vi.mocked(getWorkerDb).mockReturnValue(db); + vi.mocked(getSessionAccessCacheDO).mockReturnValue({ has: vi.fn(async () => true) } as never); + + const app = makeApiApp(); + const env = makeTestEnv(); + env.INGEST_QUEUE.send.mockRejectedValueOnce(new Error('queue unavailable')); + + const response = await app.fetch( + new Request('http://local/session/ses_12345678901234567890123456/ingest', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ data: [] }), + }), + env + ); + + expect(response.status).toBe(500); + const r2Key = env.SESSION_INGEST_R2.put.mock.calls[0][0]; + expect(env.SESSION_INGEST_R2.delete).toHaveBeenCalledWith(r2Key); + }); + it('POST /session/:sessionId/ingest returns 404 on cache miss + missing session', async () => { const { db, fns } = makeDbFakes(); vi.mocked(getWorkerDb).mockReturnValue(db); diff --git a/services/session-ingest/src/routes/api.ts b/services/session-ingest/src/routes/api.ts index 44b021878d..079c473a4c 100644 --- a/services/session-ingest/src/routes/api.ts +++ b/services/session-ingest/src/routes/api.ts @@ -11,7 +11,7 @@ import { getSessionAccessCacheDO } from '../dos/SessionAccessCacheDO'; import { getUserConnectionDO } from '../dos/UserConnectionDO'; import { getSessionExport } from '../services/session-export'; import { mapSessionEventRow, notifyUserSessionEvent } from '../session-events'; -import type { IngestQueueMessage } from '../queue-consumer'; +import { stageAndEnqueue } from '../ingest/stage-and-enqueue'; export type ApiContext = { Bindings: Env; @@ -252,23 +252,16 @@ api.post('/session/:sessionId/ingest', async c => { // Stream request body directly to R2 (zero memory) const r2Key = `ingest/${kiloUserId}/${sessionId}/${crypto.randomUUID()}`; - await c.env.SESSION_INGEST_R2.put(r2Key, c.req.raw.body); - - // Enqueue for async processing - const queueMessage: IngestQueueMessage = { - r2Key, - kiloUserId, - sessionId, - ingestVersion, - ingestedAt: Date.now(), - }; - try { - await c.env.INGEST_QUEUE.send(queueMessage); - } catch (err) { - // Clean up staging R2 object to prevent orphaned blobs - await c.env.SESSION_INGEST_R2.delete(r2Key).catch(() => {}); - throw err; - } + await stageAndEnqueue( + c.env, + { + r2Key, + kiloUserId, + sessionId, + ingestVersion, + }, + c.req.raw.body ?? new Uint8Array() + ); return c.json({ success: true }, 200); }); diff --git a/services/session-ingest/src/util/ingest-limits.ts b/services/session-ingest/src/util/ingest-limits.ts index 2bcafba086..ef27dbda5c 100644 --- a/services/session-ingest/src/util/ingest-limits.ts +++ b/services/session-ingest/src/util/ingest-limits.ts @@ -6,3 +6,7 @@ export const MAX_INGEST_ITEM_BYTES = 2 * 1024 * 1024 - 64 * 1024; // Items above this byte count are skipped during queue processing. // Tracked incrementally during streaming parse — item is aborted as soon as it exceeds this. export const MAX_SINGLE_ITEM_BYTES = 50 * 1024 * 1024; + +// Shared policy for keeping one ingest RPC within conservative payload bounds. +export const INGEST_CHUNK_MAX_BYTES = 4 * 1024 * 1024; +export const INGEST_CHUNK_MAX_ITEMS = 128;