From 403f72887c1d5fa14455b3c17679a0580ffd24fc Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Tue, 2 Jun 2026 15:47:49 +0200 Subject: [PATCH 1/6] feat(cloud-agent): add Kilo SDK session facade --- .kilo/rules/tools.md | 8 +- .../session-ingest-contracts/package.json | 22 + .../session-ingest-contracts/src/index.ts | 1 + .../src/rpc-contract.ts | 593 ++++ .../session-ingest-contracts/tsconfig.json | 17 + pnpm-lock.yaml | 22 + services/cloud-agent-next/package.json | 2 + services/cloud-agent-next/src/index.ts | 1 + .../src/kilo-facade/basic-prompt.ts | 46 + .../cloud-agent-extension-events.test.ts | 131 + .../cloud-agent-extension-events.ts | 179 + .../kilo-facade/public-sdk-projection.test.ts | 451 +++ .../src/kilo-facade/public-sdk-projection.ts | 449 +++ .../src/kilo-facade/session-proxy.ts | 85 + .../src/kilo-facade/user-kilo-facade.test.ts | 3072 +++++++++++++++++ .../src/kilo-facade/user-kilo-facade.ts | 1634 +++++++++ .../src/persistence/CloudAgentSession.ts | 65 + services/cloud-agent-next/src/router.test.ts | 5 + .../src/router/handlers/session-execution.ts | 29 +- .../src/router/handlers/session-send.ts | 17 +- services/cloud-agent-next/src/server.test.ts | 181 + services/cloud-agent-next/src/server.ts | 144 + .../src/session-ingest-binding.ts | 31 +- .../src/session-service.test.ts | 8 +- .../cloud-agent-next/src/session-service.ts | 9 +- .../src/session/queue-message.ts | 38 +- .../wrapper-global-feed-validation.test.ts | 63 + .../session/wrapper-global-feed-validation.ts | 45 + .../cloud-agent-next/src/shared/http-proxy.ts | 52 + services/cloud-agent-next/src/types.ts | 3 + services/cloud-agent-next/test/e2e/README.md | 27 +- services/cloud-agent-next/test/e2e/auth.ts | 12 +- .../test/e2e/sandbox-control.ts | 99 +- .../test/e2e/sdk-basic-chat.ts | 577 ++++ .../integration/kilo-facade-runtime.test.ts | 356 ++ services/cloud-agent-next/test/test-worker.ts | 53 +- .../test/unit/sandbox-control.test.ts | 112 + .../test/unit/wrapper/server.test.ts | 133 + services/cloud-agent-next/wrangler.jsonc | 12 + services/cloud-agent-next/wrangler.test.jsonc | 8 + .../wrapper/src/global-feed-manager.test.ts | 128 + .../wrapper/src/global-feed-manager.ts | 58 + .../wrapper/src/global-feed.test.ts | 304 ++ .../wrapper/src/global-feed.ts | 299 ++ services/cloud-agent-next/wrapper/src/main.ts | 37 +- .../wrapper/src/restore-session.test.ts | 67 + .../wrapper/src/restore-session.ts | 36 +- .../wrapper/src/server.test.ts | 195 ++ .../cloud-agent-next/wrapper/src/server.ts | 67 +- .../cloud-agent-next/wrapper/src/utils.ts | 28 +- services/session-ingest/package.json | 1 + .../session-ingest/src/dos/SessionIngestDO.ts | 23 +- .../src/dos/kilo-sdk-materialization.ts | 829 +++++ .../dos/session-ingest-history-budget.test.ts | 224 ++ .../session-ingest/src/queue-consumer.test.ts | 81 +- .../src/session-ingest-rpc.test.ts | 822 +++++ .../session-ingest/src/session-ingest-rpc.ts | 226 +- .../src/types/session-sync.test.ts | 54 + .../session-ingest/src/types/session-sync.ts | 13 +- .../src/util/compaction.test.ts | 11 +- .../session-ingest/src/util/compaction.ts | 5 + .../integration/session-ingest-do.test.ts | 1845 +++++++++- services/session-ingest/test/test-worker.ts | 7 + services/session-ingest/wrangler.test.jsonc | 8 +- 64 files changed, 13989 insertions(+), 171 deletions(-) create mode 100644 packages/session-ingest-contracts/package.json create mode 100644 packages/session-ingest-contracts/src/index.ts create mode 100644 packages/session-ingest-contracts/src/rpc-contract.ts create mode 100644 packages/session-ingest-contracts/tsconfig.json create mode 100644 services/cloud-agent-next/src/kilo-facade/basic-prompt.ts create mode 100644 services/cloud-agent-next/src/kilo-facade/cloud-agent-extension-events.test.ts create mode 100644 services/cloud-agent-next/src/kilo-facade/cloud-agent-extension-events.ts create mode 100644 services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts create mode 100644 services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts create mode 100644 services/cloud-agent-next/src/kilo-facade/session-proxy.ts create mode 100644 services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts create mode 100644 services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts create mode 100644 services/cloud-agent-next/src/session/wrapper-global-feed-validation.test.ts create mode 100644 services/cloud-agent-next/src/session/wrapper-global-feed-validation.ts create mode 100644 services/cloud-agent-next/src/shared/http-proxy.ts create mode 100644 services/cloud-agent-next/test/e2e/sdk-basic-chat.ts create mode 100644 services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts create mode 100644 services/cloud-agent-next/test/unit/sandbox-control.test.ts create mode 100644 services/cloud-agent-next/wrapper/src/global-feed-manager.test.ts create mode 100644 services/cloud-agent-next/wrapper/src/global-feed-manager.ts create mode 100644 services/cloud-agent-next/wrapper/src/global-feed.test.ts create mode 100644 services/cloud-agent-next/wrapper/src/global-feed.ts create mode 100644 services/session-ingest/src/dos/kilo-sdk-materialization.ts create mode 100644 services/session-ingest/src/dos/session-ingest-history-budget.test.ts create mode 100644 services/session-ingest/src/session-ingest-rpc.test.ts create mode 100644 services/session-ingest/src/types/session-sync.test.ts create mode 100644 services/session-ingest/test/test-worker.ts diff --git a/.kilo/rules/tools.md b/.kilo/rules/tools.md index 5225e6e476..77621423ab 100644 --- a/.kilo/rules/tools.md +++ b/.kilo/rules/tools.md @@ -7,9 +7,11 @@ - Date handling: date-fns - Database: drizzle-orm with pg - For React components, ALWAYS use React hooks with @tanstack/react-query for data fetching, caching, and server state management. Do not implement custom data fetching logic when react-query can handle it. -- The webserver should already be running. Read the dev server port from the `.dev-port` file in the project root (e.g., `cat .dev-port`), then access it via `http://localhost:/`. If `.dev-port` does not exist, start the dev server by running `pnpm dev:start &` in the background, then wait for `.dev-port` to appear before reading the port from it. -- If you need to log in as a fake user, open http://localhost:/users/sign_in?fakeUser= where is read from `.dev-port` (falling back to 3000) and is constructed from "kilo-", my username (based on homedir), and the time (include seconds) and then '@example.com'. If you need an admin account, the email address must end in @admin.example.com. After logging in, wait for the creating your account spinner to complete before proceeding. The admin panels can be accessed from your profile via the account icon in the top-right corner, which opens a drop-down, allowing access to the admin panel. -- be sure to add the callbackPath url parameter to go directly to the page after logging in! +- Before accessing local app or service endpoints, check this worktree's running services with `pnpm dev:status --json`. If local services are already active, reuse that session and its ports; do not start a competing stack. +- If no local dev services are running for this worktree and the task requires them, start only the minimum required group or named services with `KILO_PORT_OFFSET=auto pnpm dev:start `. If generated local endpoint configuration is required before startup, first run `KILO_PORT_OFFSET=auto pnpm dev:env ` (or the matching documented selector), then start the same selection with `KILO_PORT_OFFSET=auto`, retaining the selected offset/session. +- After startup or when reusing services, obtain actual ports from `.dev-port`, `pnpm dev:status --json`, or `dev/logs/manifest.json`; never assume default ports. Access the web app at `http://localhost:/` using its reported port. +- If you need to log in as a fake user, open `http://localhost:/users/sign_in?fakeUser=&callbackPath=` where `` is the reported web app port and `` is constructed from `"kilo-"`, my username (based on homedir), and the time (include seconds) and then `'@example.com'`. If you need an admin account, the email address must end in `@admin.example.com`. After logging in, wait for the creating your account spinner to complete before proceeding. The admin panels can be accessed from your profile via the account icon in the top-right corner, which opens a drop-down, allowing access to the admin panel. +- Always set the `callbackPath` URL parameter to go directly to the page after logging in. - Dev services status and management: - `cat dev/logs/manifest.json` — static snapshot of started services and ports (written on `dev:start`) - `pnpm dev:status` — live status of running services with ports diff --git a/packages/session-ingest-contracts/package.json b/packages/session-ingest-contracts/package.json new file mode 100644 index 0000000000..74206aaf6e --- /dev/null +++ b/packages/session-ingest-contracts/package.json @@ -0,0 +1,22 @@ +{ + "name": "@kilocode/session-ingest-contracts", + "version": "0.0.1", + "private": true, + "type": "module", + "main": "./src/index.ts", + "types": "./src/index.ts", + "exports": { + ".": "./src/index.ts" + }, + "scripts": { + "typecheck": "tsgo --noEmit", + "lint": "pnpm -w exec oxlint --config .oxlintrc.json packages/session-ingest-contracts/src" + }, + "dependencies": { + "zod": "catalog:" + }, + "devDependencies": { + "@typescript/native-preview": "catalog:", + "typescript": "catalog:" + } +} diff --git a/packages/session-ingest-contracts/src/index.ts b/packages/session-ingest-contracts/src/index.ts new file mode 100644 index 0000000000..6cedf0e64f --- /dev/null +++ b/packages/session-ingest-contracts/src/index.ts @@ -0,0 +1 @@ +export * from './rpc-contract'; diff --git a/packages/session-ingest-contracts/src/rpc-contract.ts b/packages/session-ingest-contracts/src/rpc-contract.ts new file mode 100644 index 0000000000..3e5653a2d8 --- /dev/null +++ b/packages/session-ingest-contracts/src/rpc-contract.ts @@ -0,0 +1,593 @@ +import { z } from 'zod'; + +export const sessionIdSchema = z.string().startsWith('ses_').length(30); +export const messageIdSchema = z + .string() + .startsWith('msg') + .refine(id => !id.includes('/') && !id.includes('\u0000'), { + message: 'message IDs must not contain / or U+0000', + }); +export const partIdSchema = z + .string() + .startsWith('prt') + .refine(id => !id.includes('/') && !id.includes('\u0000'), { + message: 'part IDs must not contain / or U+0000', + }); +const sdkMetadataSchema = z.record(z.string(), z.unknown()); + +export const createSessionForCloudAgentSchema = z.object({ + sessionId: sessionIdSchema, + kiloUserId: z.string().min(1), + cloudAgentSessionId: z.string().min(1), + organizationId: z.string().optional(), + createdOnPlatform: z.string().min(1), + title: z.string().optional(), +}); +export type CreateSessionForCloudAgentParams = z.input; + +export const deleteSessionForCloudAgentSchema = z.object({ + sessionId: sessionIdSchema, + kiloUserId: z.string().min(1), + onlyIfEmpty: z.boolean().optional(), +}); +export type DeleteSessionForCloudAgentParams = z.input; + +export const resolveCloudAgentRootSessionSchema = z.object({ + kiloUserId: z.string().min(1), + kiloSessionId: sessionIdSchema, +}); +export type ResolveCloudAgentRootSessionForKiloSessionParams = z.input< + typeof resolveCloudAgentRootSessionSchema +>; +export type ResolveCloudAgentRootSessionForKiloSessionResult = { + cloudAgentSessionId: string; +} | null; + +export const kiloSdkSessionInfoSchema = z.object({ + id: sessionIdSchema, + slug: z.string(), + projectID: z.string(), + workspaceID: z.string().optional(), + directory: z.string(), + path: z.string().optional(), + parentID: z.string().optional(), + summary: z + .object({ + additions: z.number(), + deletions: z.number(), + files: z.number(), + diffs: z + .array( + z.object({ + file: z.string(), + additions: z.number(), + deletions: z.number(), + status: z.enum(['added', 'deleted', 'modified']).optional(), + }) + ) + .optional(), + }) + .optional(), + share: z.object({ url: z.string() }).optional(), + title: z.string(), + agent: z.string().optional(), + model: z + .object({ + id: z.string(), + providerID: z.string(), + variant: z.string().optional(), + }) + .optional(), + version: z.string(), + time: z.object({ + created: z.number(), + updated: z.number(), + compacting: z.number().optional(), + archived: z.number().optional(), + }), + permission: z + .array( + z.object({ + permission: z.string(), + pattern: z.string(), + action: z.enum(['allow', 'deny', 'ask']), + }) + ) + .optional(), + revert: z + .object({ + messageID: z.string(), + partID: z.string().optional(), + snapshot: z.string().optional(), + diff: z.string().optional(), + }) + .optional(), +}); +export type KiloSdkSessionInfo = z.infer; + +export const kiloSdkSessionSnapshotOutcomeSchema = z.discriminatedUnion('kind', [ + z.object({ kind: z.literal('pending') }), + z.object({ + kind: z.literal('value'), + info: kiloSdkSessionInfoSchema, + byteLength: z.number().int().nonnegative(), + }), + z.object({ kind: z.literal('too_large'), maximumBytes: z.number().int().positive() }), + z.object({ kind: z.literal('retryable_failure') }), + z.object({ kind: z.literal('invalid_data') }), +]); +export type KiloSdkSessionSnapshotRead = z.infer; +export type KiloSdkSessionSnapshotOutcome = KiloSdkSessionSnapshotRead; + +export type CloudAgentRootSessionSnapshot = { + kiloSessionId: string; + cloudAgentSessionId: string; + snapshot: KiloSdkSessionSnapshotRead; +}; +export type GetCloudAgentRootSessionSnapshotParams = + ResolveCloudAgentRootSessionForKiloSessionParams; +export type GetCloudAgentRootSessionSnapshotResult = CloudAgentRootSessionSnapshot | null; + +export const MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE = 100; + +export const listCloudAgentRootSessionsSchema = z.object({ + kiloUserId: z.string().min(1), + limit: z.number().int().min(1).max(100).optional().default(100), + start: z.number().int().nonnegative().max(8_640_000_000_000_000).optional(), +}); +export type ListCloudAgentRootSessionsParams = z.input; +export type CloudAgentRootSessionSummary = { + kiloSessionId: string; + cloudAgentSessionId: string; + title: string | null; + created: number; + updated: number; +}; + +export const sdkApiErrorSchema = z.object({ + name: z.literal('APIError'), + data: z.object({ + message: z.string(), + statusCode: z.number().optional(), + isRetryable: z.boolean(), + responseHeaders: z.record(z.string(), z.string()).optional(), + responseBody: z.string().optional(), + metadata: z.record(z.string(), z.string()).optional(), + }), +}); +export type KiloSdkApiError = z.infer; + +export const sdkAssistantErrorSchema = z.discriminatedUnion('name', [ + z.object({ + name: z.literal('ProviderAuthError'), + data: z.object({ providerID: z.string(), message: z.string() }), + }), + z.object({ name: z.literal('UnknownError'), data: z.object({ message: z.string() }) }), + z.object({ name: z.literal('MessageOutputLengthError'), data: sdkMetadataSchema }), + z.object({ name: z.literal('MessageAbortedError'), data: z.object({ message: z.string() }) }), + z.object({ + name: z.literal('StructuredOutputError'), + data: z.object({ message: z.string(), retries: z.number() }), + }), + z.object({ + name: z.literal('ContextOverflowError'), + data: z.object({ message: z.string(), responseBody: z.string().optional() }), + }), + sdkApiErrorSchema, +]); +export type KiloSdkAssistantError = z.infer; + +const sdkSnapshotFileDiffBaseSchema = z.object({ + file: z.string(), + additions: z.number(), + deletions: z.number(), + status: z.enum(['added', 'deleted', 'modified']).optional(), +}); +export const sdkSnapshotFileDiffSchema = sdkSnapshotFileDiffBaseSchema + .extend({ patch: z.string() }) + .strict(); +const historicalSnapshotFileDiffSchema = sdkSnapshotFileDiffBaseSchema + .extend({ + before: z.string(), + after: z.string(), + }) + .strict(); +export const persistedSnapshotFileDiffsSchema = z + .array(z.union([sdkSnapshotFileDiffSchema, historicalSnapshotFileDiffSchema])) + .transform(diffs => diffs.filter(diff => 'patch' in diff)); + +const kiloSdkUserMessageBaseShape = { + id: messageIdSchema, + sessionID: sessionIdSchema, + role: z.literal('user'), + time: z.object({ created: z.number() }), + format: z + .discriminatedUnion('type', [ + z.object({ type: z.literal('text') }), + z.object({ + type: z.literal('json_schema'), + schema: sdkMetadataSchema, + retryCount: z.number().optional(), + }), + ]) + .optional(), + agent: z.string(), + model: z.object({ + providerID: z.string(), + modelID: z.string(), + variant: z.string().optional(), + }), + system: z.string().optional(), + tools: z.record(z.string(), z.boolean()).optional(), + editorContext: z + .object({ + visibleFiles: z.array(z.string()).optional(), + openTabs: z.array(z.string()).optional(), + activeFile: z.string().optional(), + shell: z.string().optional(), + }) + .optional(), +}; +const sdkUserMessageSummaryBaseShape = { + title: z.string().optional(), + body: z.string().optional(), +}; +export const kiloSdkUserMessageSchema = z.object({ + ...kiloSdkUserMessageBaseShape, + summary: z + .object({ + ...sdkUserMessageSummaryBaseShape, + diffs: z.array(sdkSnapshotFileDiffSchema), + }) + .optional(), +}); +export type KiloSdkUserMessage = z.infer; + +export const kiloSdkAssistantMessageSchema = z.object({ + id: messageIdSchema, + sessionID: sessionIdSchema, + role: z.literal('assistant'), + time: z.object({ created: z.number(), completed: z.number().optional() }), + error: sdkAssistantErrorSchema.optional(), + parentID: messageIdSchema, + modelID: z.string(), + providerID: z.string(), + mode: z.string(), + agent: z.string(), + path: z.object({ cwd: z.string(), root: z.string() }), + summary: z.boolean().optional(), + cost: z.number(), + tokens: z.object({ + total: z.number().optional(), + input: z.number(), + output: z.number(), + reasoning: z.number(), + cache: z.object({ read: z.number(), write: z.number() }), + }), + structured: z.unknown().optional(), + variant: z.string().optional(), + finish: z.string().optional(), +}); +export type KiloSdkAssistantMessage = z.infer; + +export const kiloSdkMessageSchema = z.discriminatedUnion('role', [ + kiloSdkUserMessageSchema, + kiloSdkAssistantMessageSchema, +]); +export type KiloSdkMessageInfo = z.infer; + +const sdkPartBaseShape = { + id: partIdSchema, + sessionID: sessionIdSchema, + messageID: messageIdSchema, +}; +const sdkPartBaseSchema = z.object(sdkPartBaseShape); +export type KiloSdkPartBase = z.infer; + +export const sdkFilePartSchema = z.object({ + ...sdkPartBaseShape, + type: z.literal('file'), + mime: z.string(), + filename: z.string().optional(), + url: z.string(), + source: z + .discriminatedUnion('type', [ + z.object({ + type: z.literal('file'), + text: z.object({ value: z.string(), start: z.number(), end: z.number() }), + path: z.string(), + }), + z.object({ + type: z.literal('symbol'), + text: z.object({ value: z.string(), start: z.number(), end: z.number() }), + path: z.string(), + range: z.object({ + start: z.object({ line: z.number(), character: z.number() }), + end: z.object({ line: z.number(), character: z.number() }), + }), + name: z.string(), + kind: z.number(), + }), + z.object({ + type: z.literal('resource'), + text: z.object({ value: z.string(), start: z.number(), end: z.number() }), + clientName: z.string(), + uri: z.string(), + }), + ]) + .optional(), +}); +export type KiloSdkFilePart = z.infer; + +export const sdkToolStateSchema = z.discriminatedUnion('status', [ + z.object({ status: z.literal('pending'), input: sdkMetadataSchema, raw: z.string() }), + z.object({ + status: z.literal('running'), + input: sdkMetadataSchema, + title: z.string().optional(), + metadata: sdkMetadataSchema.optional(), + time: z.object({ start: z.number() }), + }), + z.object({ + status: z.literal('completed'), + input: sdkMetadataSchema, + output: z.string(), + title: z.string(), + metadata: sdkMetadataSchema, + time: z.object({ start: z.number(), end: z.number(), compacted: z.number().optional() }), + attachments: z.array(sdkFilePartSchema).optional(), + }), + z.object({ + status: z.literal('error'), + input: sdkMetadataSchema, + error: z.string(), + metadata: sdkMetadataSchema.optional(), + time: z.object({ start: z.number(), end: z.number() }), + }), +]); +export type KiloSdkToolState = z.infer; + +export const kiloSdkPartSchema = z.discriminatedUnion('type', [ + z.object({ + ...sdkPartBaseShape, + type: z.literal('text'), + text: z.string(), + synthetic: z.boolean().optional(), + ignored: z.boolean().optional(), + time: z.object({ start: z.number(), end: z.number().optional() }).optional(), + metadata: sdkMetadataSchema.optional(), + }), + z.object({ + ...sdkPartBaseShape, + type: z.literal('subtask'), + prompt: z.string(), + description: z.string(), + agent: z.string(), + model: z.object({ providerID: z.string(), modelID: z.string() }).optional(), + command: z.string().optional(), + }), + z.object({ + ...sdkPartBaseShape, + type: z.literal('reasoning'), + text: z.string(), + metadata: sdkMetadataSchema.optional(), + time: z.object({ start: z.number(), end: z.number().optional() }), + }), + sdkFilePartSchema, + z.object({ + ...sdkPartBaseShape, + type: z.literal('tool'), + callID: z.string(), + tool: z.string(), + state: sdkToolStateSchema, + metadata: sdkMetadataSchema.optional(), + }), + z.object({ ...sdkPartBaseShape, type: z.literal('step-start'), snapshot: z.string().optional() }), + z.object({ + ...sdkPartBaseShape, + type: z.literal('step-finish'), + reason: z.string(), + snapshot: z.string().optional(), + cost: z.number(), + tokens: z.object({ + total: z.number().optional(), + input: z.number(), + output: z.number(), + reasoning: z.number(), + cache: z.object({ read: z.number(), write: z.number() }), + }), + }), + z.object({ ...sdkPartBaseShape, type: z.literal('snapshot'), snapshot: z.string() }), + z.object({ + ...sdkPartBaseShape, + type: z.literal('patch'), + hash: z.string(), + files: z.array(z.string()), + }), + z.object({ + ...sdkPartBaseShape, + type: z.literal('agent'), + name: z.string(), + source: z.object({ value: z.string(), start: z.number(), end: z.number() }).optional(), + }), + z.object({ + ...sdkPartBaseShape, + type: z.literal('retry'), + attempt: z.number(), + error: sdkApiErrorSchema, + time: z.object({ created: z.number() }), + }), + z.object({ + ...sdkPartBaseShape, + type: z.literal('compaction'), + auto: z.boolean(), + overflow: z.boolean().optional(), + tail_start_id: messageIdSchema.optional(), + }), +]); +export type KiloSdkPart = z.infer; + +export const kiloSdkStoredMessageSchema = z.object({ + info: kiloSdkMessageSchema, + parts: z.array(kiloSdkPartSchema), +}); +export type KiloSdkStoredMessage = z.infer; + +export const kiloSdkMessageHistoryPageSchema = z.object({ + messages: z.array(kiloSdkStoredMessageSchema), + nextCursor: z.string().nullable(), + omittedItemCount: z.number().int().nonnegative().default(0), +}); +export type KiloSdkMessageHistoryPage = z.infer; + +export const kiloSdkHistoryTooLargeSchema = z.object({ + kind: z.literal('too_large'), + maximumBytes: z.number().int().positive(), + phase: z.enum(['message_scan', 'page_parts']), +}); +export type KiloSdkHistoryTooLarge = z.infer; + +export const kiloSdkHistoryRetryableFailureSchema = z.object({ + kind: z.literal('retryable_failure'), + phase: z.enum(['message_scan', 'page_parts']), +}); +export type KiloSdkHistoryRetryableFailure = z.infer; + +export const kiloSdkInvalidDataSchema = z.object({ kind: z.literal('invalid_data') }); +export type KiloSdkInvalidData = z.infer; + +export const kiloSdkMessageHistorySchema = z.union([ + kiloSdkMessageHistoryPageSchema, + kiloSdkHistoryTooLargeSchema, + kiloSdkHistoryRetryableFailureSchema, + kiloSdkInvalidDataSchema, +]); + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function normalizePersistedSnapshotFileDiffs(value: unknown): unknown { + const parsed = persistedSnapshotFileDiffsSchema.safeParse(value); + return parsed.success ? parsed.data : value; +} + +function normalizePersistedKiloSdkStoredMessage(value: unknown): unknown { + if (!isRecord(value) || !isRecord(value.info) || value.info.role !== 'user') { + return value; + } + const summary = value.info.summary; + if (!isRecord(summary) || !('diffs' in summary)) { + return value; + } + return { + ...value, + info: { + ...value.info, + summary: { + ...summary, + diffs: normalizePersistedSnapshotFileDiffs(summary.diffs), + }, + }, + }; +} + +function normalizePersistedKiloSdkMessageHistory(value: unknown): unknown { + if (!isRecord(value) || !Array.isArray(value.messages)) { + return value; + } + return { ...value, messages: value.messages.map(normalizePersistedKiloSdkStoredMessage) }; +} + +export const persistedKiloSdkMessageHistorySchema = z.preprocess( + normalizePersistedKiloSdkMessageHistory, + kiloSdkMessageHistorySchema +); +export type KiloSdkMessageHistory = z.infer; + +export const kiloSdkMessagesLegacyCursorSchema = z + .object({ + id: messageIdSchema, + time: z.number().nonnegative(), + }) + .strict(); +export type KiloSdkMessagesLegacyCursor = z.infer; + +export const kiloSdkMessagesCursorSchema = kiloSdkMessagesLegacyCursorSchema; +export type KiloSdkMessagesCursor = KiloSdkMessagesLegacyCursor; + +export function encodeKiloSdkMessagesCursor(cursor: KiloSdkMessagesCursor): string { + const parsed = kiloSdkMessagesCursorSchema.parse(cursor); + const bytes = new TextEncoder().encode(JSON.stringify(parsed)); + let binary = ''; + for (const byte of bytes) { + binary += String.fromCharCode(byte); + } + return btoa(binary).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/g, ''); +} + +export function decodeKiloSdkMessagesCursor(cursor: string): KiloSdkMessagesCursor { + const base64 = cursor.replace(/-/g, '+').replace(/_/g, '/'); + const paddingLength = (4 - (base64.length % 4)) % 4; + const binary = atob(base64 + '='.repeat(paddingLength)); + const bytes = Uint8Array.from(binary, character => character.charCodeAt(0)); + return kiloSdkMessagesCursorSchema.parse(JSON.parse(new TextDecoder().decode(bytes))); +} + +export function validateKiloSdkMessagesCursor(cursor: string): boolean { + try { + decodeKiloSdkMessagesCursor(cursor); + return true; + } catch { + return false; + } +} + +export const getCloudAgentRootSessionMessagesSchema = resolveCloudAgentRootSessionSchema + .extend({ + limit: z.number().int().nonnegative().max(MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE).optional(), + before: z.string().min(1).optional(), + }) + .superRefine((params, ctx) => { + if (params.before !== undefined && (params.limit === undefined || params.limit === 0)) { + ctx.addIssue({ + code: 'custom', + path: ['before'], + message: 'before requires a positive limit', + }); + return; + } + if (params.before !== undefined && !validateKiloSdkMessagesCursor(params.before)) { + ctx.addIssue({ + code: 'custom', + path: ['before'], + message: 'before is not a valid message cursor', + }); + } + }); +export type GetCloudAgentRootSessionMessagesParams = z.input< + typeof getCloudAgentRootSessionMessagesSchema +>; +export type CloudAgentRootSessionMessages = { + kiloSessionId: string; + cloudAgentSessionId: string; + history: KiloSdkMessageHistory | null; +}; +export type GetCloudAgentRootSessionMessagesResult = CloudAgentRootSessionMessages | null; + +export type SessionIngestRpcMethods = { + createSessionForCloudAgent: (params: CreateSessionForCloudAgentParams) => Promise; + deleteSessionForCloudAgent: (params: DeleteSessionForCloudAgentParams) => Promise; + resolveCloudAgentRootSessionForKiloSession: ( + params: ResolveCloudAgentRootSessionForKiloSessionParams + ) => Promise; + getCloudAgentRootSessionSnapshot: ( + params: GetCloudAgentRootSessionSnapshotParams + ) => Promise; + listCloudAgentRootSessions: ( + params: ListCloudAgentRootSessionsParams + ) => Promise; + getCloudAgentRootSessionMessages: ( + params: GetCloudAgentRootSessionMessagesParams + ) => Promise; +}; diff --git a/packages/session-ingest-contracts/tsconfig.json b/packages/session-ingest-contracts/tsconfig.json new file mode 100644 index 0000000000..76473b226e --- /dev/null +++ b/packages/session-ingest-contracts/tsconfig.json @@ -0,0 +1,17 @@ +{ + "compilerOptions": { + "target": "ESNext", + "module": "ESNext", + "moduleResolution": "bundler", + "lib": ["ESNext", "WebWorker"], + "strict": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "noEmit": true, + "isolatedModules": true, + "noImplicitReturns": true, + "noFallthroughCasesInSwitch": true + }, + "include": ["src/**/*"], + "exclude": ["node_modules"] +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bf48d30cfc..6ef386fbdf 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1178,6 +1178,19 @@ importers: specifier: 'catalog:' version: 4.1.6(@opentelemetry/api@1.9.1)(@types/node@25.5.2)(@vitest/coverage-v8@4.1.6)(@vitest/ui@4.1.6)(esbuild@0.27.4)(jiti@2.7.0)(terser@5.46.0)(tsx@4.21.0)(yaml@2.8.4) + packages/session-ingest-contracts: + dependencies: + zod: + specifier: 'catalog:' + version: 4.4.3 + devDependencies: + '@typescript/native-preview': + specifier: 'catalog:' + version: 7.0.0-dev.20260514.1 + typescript: + specifier: 'catalog:' + version: 5.9.3 + packages/trpc: dependencies: '@kilocode/kiloclaw-secret-catalog': @@ -1509,6 +1522,9 @@ importers: '@kilocode/notifications': specifier: workspace:* version: link:../../packages/notifications + '@kilocode/session-ingest-contracts': + specifier: workspace:* + version: link:../../packages/session-ingest-contracts '@kilocode/worker-utils': specifier: workspace:* version: link:../../packages/worker-utils @@ -1543,6 +1559,9 @@ importers: '@cloudflare/vitest-pool-workers': specifier: 'catalog:' version: 0.16.4(@cloudflare/workers-types@4.20260511.1)(@vitest/runner@4.1.6)(@vitest/snapshot@4.1.6)(bufferutil@4.1.0)(utf-8-validate@6.0.6)(vitest@4.1.6) + '@kilocode/sdk': + specifier: 7.2.52 + version: 7.2.52 '@types/jsonwebtoken': specifier: 'catalog:' version: 9.0.10 @@ -2502,6 +2521,9 @@ importers: '@kilocode/db': specifier: workspace:* version: link:../../packages/db + '@kilocode/session-ingest-contracts': + specifier: workspace:* + version: link:../../packages/session-ingest-contracts '@kilocode/worker-utils': specifier: workspace:* version: link:../../packages/worker-utils diff --git a/services/cloud-agent-next/package.json b/services/cloud-agent-next/package.json index 5a63c9f78d..c126a9d078 100644 --- a/services/cloud-agent-next/package.json +++ b/services/cloud-agent-next/package.json @@ -36,6 +36,7 @@ "@kilocode/db": "workspace:*", "@kilocode/encryption": "workspace:*", "@kilocode/notifications": "workspace:*", + "@kilocode/session-ingest-contracts": "workspace:*", "@kilocode/worker-utils": "workspace:*", "@trpc/server": "catalog:", "drizzle-orm": "catalog:", @@ -49,6 +50,7 @@ "devDependencies": { "@cloudflare/containers": "0.3.5", "@cloudflare/vitest-pool-workers": "catalog:", + "@kilocode/sdk": "7.2.52", "@types/jsonwebtoken": "catalog:", "@types/node": "catalog:", "@types/ws": "^8.5.12", diff --git a/services/cloud-agent-next/src/index.ts b/services/cloud-agent-next/src/index.ts index d15d592963..ccc57f09f3 100644 --- a/services/cloud-agent-next/src/index.ts +++ b/services/cloud-agent-next/src/index.ts @@ -1,3 +1,4 @@ export { default } from './server.js'; export { Sandbox, Sandbox as SandboxSmall, Sandbox as SandboxDIND } from '@cloudflare/sandbox'; export { CloudAgentSession } from './persistence/CloudAgentSession.js'; +export { UserKiloFacade } from './kilo-facade/user-kilo-facade.js'; diff --git a/services/cloud-agent-next/src/kilo-facade/basic-prompt.ts b/services/cloud-agent-next/src/kilo-facade/basic-prompt.ts new file mode 100644 index 0000000000..c06c2ed3dc --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/basic-prompt.ts @@ -0,0 +1,46 @@ +import * as z from 'zod'; +import { Limits } from '../schema.js'; +import { MessageIdSchema } from '../router/schemas.js'; + +const basicTextPartSchema = z + .object({ + type: z.literal('text'), + text: z.string(), + }) + .strict(); + +const basicPromptBodySchema = z + .object({ + messageID: MessageIdSchema.optional(), + parts: z.array(basicTextPartSchema).min(1), + }) + .strict(); + +export type BasicKiloPrompt = { + messageId?: string; + prompt: string; +}; + +export type BasicKiloPromptParseResult = + | { success: true; prompt: BasicKiloPrompt } + | { success: false }; + +export function parseBasicKiloPrompt(value: unknown): BasicKiloPromptParseResult { + const result = basicPromptBodySchema.safeParse(value); + if (!result.success) { + return { success: false }; + } + + const prompt = result.data.parts.map(part => part.text).join(''); + if (prompt.length === 0 || prompt.length > Limits.MAX_PROMPT_LENGTH) { + return { success: false }; + } + + return { + success: true, + prompt: { + messageId: result.data.messageID, + prompt, + }, + }; +} diff --git a/services/cloud-agent-next/src/kilo-facade/cloud-agent-extension-events.test.ts b/services/cloud-agent-next/src/kilo-facade/cloud-agent-extension-events.test.ts new file mode 100644 index 0000000000..f48ae67008 --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/cloud-agent-extension-events.test.ts @@ -0,0 +1,131 @@ +import { describe, expect, it } from 'vitest'; +import { projectPublicCloudAgentExtensionEvent } from './cloud-agent-extension-events'; + +const kiloSessionId = 'ses_12345678901234567890123456'; + +function source(streamEventType: string, payload: unknown) { + return { stream_event_type: streamEventType, payload: JSON.stringify(payload) }; +} + +describe('projectPublicCloudAgentExtensionEvent', () => { + it('projects durable message lifecycle events without prompt or diagnostic data', () => { + expect( + projectPublicCloudAgentExtensionEvent( + source('cloud.message.queued', { + messageId: 'msg_queued', + content: 'secret prompt', + delivery: 'queued', + }), + kiloSessionId + ) + ).toEqual({ + type: 'cloud.message.queued', + properties: { sessionID: kiloSessionId, messageId: 'msg_queued', delivery: 'queued' }, + }); + expect( + projectPublicCloudAgentExtensionEvent( + source('cloud.message.sent', { messageId: 'msg_sent', delivery: 'sent' }), + kiloSessionId + ) + ).toEqual({ + type: 'cloud.message.sent', + properties: { sessionID: kiloSessionId, messageId: 'msg_sent', delivery: 'sent' }, + }); + expect( + projectPublicCloudAgentExtensionEvent( + source('cloud.message.completed', { + messageId: 'msg_completed', + status: 'completed', + delivery: 'sent', + accepted: true, + assistantMessageId: 'msg_private_assistant', + }), + kiloSessionId + ) + ).toEqual({ + type: 'cloud.message.completed', + properties: { + sessionID: kiloSessionId, + messageId: 'msg_completed', + status: 'completed', + delivery: 'sent', + accepted: true, + }, + }); + expect( + projectPublicCloudAgentExtensionEvent( + source('cloud.message.failed', { + messageId: 'msg_failed', + status: 'failed', + delivery: 'queued', + accepted: false, + error: '/workspace/private/token failure', + attempts: 3, + }), + kiloSessionId + ) + ).toEqual({ + type: 'cloud.message.failed', + properties: { + sessionID: kiloSessionId, + messageId: 'msg_failed', + status: 'failed', + delivery: 'queued', + accepted: false, + }, + }); + }); + + it('projects only public cloud status vocabulary and drops arbitrary progress messages', () => { + expect( + projectPublicCloudAgentExtensionEvent( + source('cloud.status', { + cloudStatus: { + type: 'preparing', + step: 'kilo_session', + message: 'Restoring /workspace/private/session', + }, + }), + kiloSessionId + ) + ).toEqual({ + type: 'cloud.status', + properties: { + sessionID: kiloSessionId, + cloudStatus: { type: 'preparing', step: 'kilo_session' }, + }, + }); + expect( + projectPublicCloudAgentExtensionEvent( + source('cloud.status', { + cloudStatus: { type: 'preparing', step: 'private_step', message: 'secret' }, + }), + kiloSessionId + ) + ).toEqual({ + type: 'cloud.status', + properties: { sessionID: kiloSessionId, cloudStatus: { type: 'preparing' } }, + }); + }); + + it('drops unselected or malformed source events', () => { + expect( + projectPublicCloudAgentExtensionEvent( + source('preparing', { step: 'kilo_session' }), + kiloSessionId + ) + ).toBeNull(); + expect( + projectPublicCloudAgentExtensionEvent( + source('cloud.message.queued', { content: 'missing id' }), + kiloSessionId + ) + ).toBeNull(); + expect( + projectPublicCloudAgentExtensionEvent( + { stream_event_type: 'cloud.status', payload: '{malformed' }, + kiloSessionId + ) + ).toBeNull(); + }); +}); diff --git a/services/cloud-agent-next/src/kilo-facade/cloud-agent-extension-events.ts b/services/cloud-agent-next/src/kilo-facade/cloud-agent-extension-events.ts new file mode 100644 index 0000000000..8ae9c46aa1 --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/cloud-agent-extension-events.ts @@ -0,0 +1,179 @@ +import type { StoredEvent } from '../websocket/types.js'; + +export type PublicCloudAgentStatusStep = + | 'disk_check' + | 'workspace_setup' + | 'cloning' + | 'branch' + | 'devcontainer_setup' + | 'setup_commands' + | 'kilo_server' + | 'kilo_session' + | 'ready' + | 'failed'; + +export type PublicCloudAgentExtensionEvent = + | { + type: 'cloud.message.queued'; + properties: { sessionID: string; messageId: string; delivery: 'queued' }; + } + | { + type: 'cloud.message.sent'; + properties: { sessionID: string; messageId: string; delivery: 'sent' }; + } + | { + type: 'cloud.message.completed'; + properties: { + sessionID: string; + messageId: string; + status: 'completed'; + delivery: 'sent'; + accepted: true; + }; + } + | { + type: 'cloud.message.failed'; + properties: { + sessionID: string; + messageId: string; + status: 'failed' | 'interrupted'; + delivery: 'queued' | 'sent'; + accepted: boolean; + }; + } + | { + type: 'cloud.status'; + properties: { + sessionID: string; + cloudStatus: { + type: 'preparing' | 'ready' | 'finalizing' | 'error'; + step?: PublicCloudAgentStatusStep; + }; + }; + }; + +const FORWARDED_SOURCE_EVENT_TYPES = new Set([ + 'cloud.message.queued', + 'cloud.message.sent', + 'cloud.message.completed', + 'cloud.message.failed', + 'cloud.status', +]); + +export function isPublicCloudAgentExtensionSourceType(streamEventType: string): boolean { + return FORWARDED_SOURCE_EVENT_TYPES.has(streamEventType); +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function parsePayload(payload: string): Record | undefined { + try { + const parsed: unknown = JSON.parse(payload); + return isRecord(parsed) ? parsed : undefined; + } catch { + return undefined; + } +} + +function messageIdFromPayload(payload: Record): string | undefined { + return typeof payload.messageId === 'string' && payload.messageId.length > 0 + ? payload.messageId + : undefined; +} + +function isCloudStatusType( + value: unknown +): value is 'preparing' | 'ready' | 'finalizing' | 'error' { + return value === 'preparing' || value === 'ready' || value === 'finalizing' || value === 'error'; +} + +function isCloudStatusStep(value: unknown): value is PublicCloudAgentStatusStep { + return ( + value === 'disk_check' || + value === 'workspace_setup' || + value === 'cloning' || + value === 'branch' || + value === 'devcontainer_setup' || + value === 'setup_commands' || + value === 'kilo_server' || + value === 'kilo_session' || + value === 'ready' || + value === 'failed' + ); +} + +export function projectPublicCloudAgentExtensionEvent( + event: Pick, + kiloSessionId: string +): PublicCloudAgentExtensionEvent | null { + const payload = parsePayload(event.payload); + if (!payload) return null; + + if (event.stream_event_type === 'cloud.status') { + const cloudStatus = payload.cloudStatus; + if (!isRecord(cloudStatus) || !isCloudStatusType(cloudStatus.type)) return null; + return { + type: 'cloud.status', + properties: { + sessionID: kiloSessionId, + cloudStatus: { + type: cloudStatus.type, + ...(isCloudStatusStep(cloudStatus.step) ? { step: cloudStatus.step } : {}), + }, + }, + }; + } + + const messageId = messageIdFromPayload(payload); + if (!messageId) return null; + + switch (event.stream_event_type) { + case 'cloud.message.queued': + return { + type: 'cloud.message.queued', + properties: { sessionID: kiloSessionId, messageId, delivery: 'queued' }, + }; + case 'cloud.message.sent': + return { + type: 'cloud.message.sent', + properties: { sessionID: kiloSessionId, messageId, delivery: 'sent' }, + }; + case 'cloud.message.completed': + if (payload.status !== 'completed') return null; + return { + type: 'cloud.message.completed', + properties: { + sessionID: kiloSessionId, + messageId, + status: 'completed', + delivery: 'sent', + accepted: true, + }, + }; + case 'cloud.message.failed': { + const status = payload.status; + const delivery = payload.delivery; + if ( + (status !== 'failed' && status !== 'interrupted') || + (delivery !== 'queued' && delivery !== 'sent') || + typeof payload.accepted !== 'boolean' + ) { + return null; + } + return { + type: 'cloud.message.failed', + properties: { + sessionID: kiloSessionId, + messageId, + status, + delivery, + accepted: payload.accepted, + }, + }; + } + default: + return null; + } +} diff --git a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts new file mode 100644 index 0000000000..9b446479e3 --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts @@ -0,0 +1,451 @@ +import { describe, expect, it } from 'vitest'; + +import { + hasUnprojectedPrivateStructuredPath, + projectPublicEvent, + projectPublicStoredMessage, +} from './public-sdk-projection'; + +const kiloSessionId = 'ses_12345678901234567890123456'; +const foreignKiloSessionId = 'ses_22222222222222222222222222'; + +describe('projectPublicEvent', () => { + it('suppresses message.updated when embedded message identity conflicts with the exposed root', () => { + expect( + projectPublicEvent( + { + type: 'message.updated', + properties: { + sessionID: kiloSessionId, + info: { + id: 'msg_foreign', + sessionID: foreignKiloSessionId, + role: 'assistant', + }, + }, + }, + kiloSessionId + ) + ).toBeNull(); + }); + + it('suppresses message.part.updated when embedded part identity conflicts with the exposed root', () => { + expect( + projectPublicEvent( + { + type: 'message.part.updated', + properties: { + sessionID: kiloSessionId, + part: { + id: 'prt_foreign', + sessionID: foreignKiloSessionId, + messageID: 'msg_foreign', + type: 'text', + text: 'private foreign part', + }, + }, + }, + kiloSessionId + ) + ).toBeNull(); + }); + + it('suppresses message.part.updated when a nested tool attachment identity conflicts with the exposed root', () => { + expect( + projectPublicEvent( + { + type: 'message.part.updated', + properties: { + sessionID: kiloSessionId, + part: { + id: 'prt_tool', + sessionID: kiloSessionId, + messageID: 'msg_tool', + type: 'tool', + state: { + status: 'completed', + attachments: [ + { + id: 'prt_foreign_attachment', + sessionID: foreignKiloSessionId, + messageID: 'msg_tool', + type: 'file', + mime: 'text/plain', + url: 'data:text/plain,safe', + }, + ], + }, + }, + }, + }, + kiloSessionId + ) + ).toBeNull(); + }); + + it('suppresses session.updated when a top-level session identity conflicts with the exposed root', () => { + expect( + projectPublicEvent( + { + type: 'session.updated', + properties: { + sessionID: foreignKiloSessionId, + info: { id: kiloSessionId }, + }, + }, + kiloSessionId + ) + ).toBeNull(); + }); + + it('suppresses unreviewed event variants when their top-level identity conflicts with the exposed root', () => { + expect( + projectPublicEvent( + { + type: 'todo.updated', + properties: { sessionID: foreignKiloSessionId, todos: [] }, + }, + kiloSessionId + ) + ).toBeNull(); + }); + + it('preserves root message and part events with matching embedded identities', () => { + expect( + projectPublicEvent( + { + type: 'message.updated', + properties: { + sessionID: kiloSessionId, + info: { id: 'msg_root', sessionID: kiloSessionId, role: 'assistant' }, + }, + }, + kiloSessionId + ) + ).toEqual({ + type: 'message.updated', + properties: { + sessionID: kiloSessionId, + info: { id: 'msg_root', sessionID: kiloSessionId, role: 'assistant' }, + }, + }); + expect( + projectPublicEvent( + { + type: 'message.part.updated', + properties: { + sessionID: kiloSessionId, + part: { + id: 'prt_root', + sessionID: kiloSessionId, + messageID: 'msg_root', + type: 'text', + text: 'public root part', + }, + }, + }, + kiloSessionId + ) + ).toEqual({ + type: 'message.part.updated', + properties: { + sessionID: kiloSessionId, + part: { + id: 'prt_root', + sessionID: kiloSessionId, + messageID: 'msg_root', + type: 'text', + text: 'public root part', + }, + }, + }); + }); + + it.each(['file.edited', 'pty.exited', 'indexing.status', 'session.error'])( + 'suppresses identityless %s variants fail-closed', + type => { + expect(projectPublicEvent({ type, properties: {} }, kiloSessionId)).toBeNull(); + } + ); + + it('suppresses identityless message.updated variants fail-closed', () => { + expect( + projectPublicEvent( + { type: 'message.updated', properties: { info: { id: 'msg_identityless' } } }, + kiloSessionId + ) + ).toBeNull(); + }); + + it('does not trust arbitrary nested matching session identities for attribution', () => { + expect( + projectPublicEvent( + { + type: 'todo.updated', + properties: { metadata: { sessionID: kiloSessionId }, todos: [] }, + }, + kiloSessionId + ) + ).toBeNull(); + }); + + it('does not treat arbitrary nested foreign session identities as authoritative', () => { + expect( + projectPublicEvent( + { + type: 'todo.updated', + properties: { + sessionID: kiloSessionId, + metadata: { sessionID: foreignKiloSessionId }, + todos: [], + }, + }, + kiloSessionId + ) + ).toEqual({ + type: 'todo.updated', + properties: { + sessionID: kiloSessionId, + metadata: { sessionID: foreignKiloSessionId }, + todos: [], + }, + }); + }); + + it('does not treat tool input or structured output session identities as authoritative', () => { + const event = { + type: 'message.part.updated', + properties: { + sessionID: kiloSessionId, + part: { + id: 'prt_tool_data', + sessionID: kiloSessionId, + messageID: 'msg_tool_data', + type: 'tool', + state: { + status: 'completed', + input: { sessionID: foreignKiloSessionId }, + output: { sessionID: foreignKiloSessionId }, + metadata: { sessionID: foreignKiloSessionId }, + }, + }, + }, + }; + expect(projectPublicEvent(event, kiloSessionId)).toEqual(event); + }); + + it('suppresses explicit child session.idle variants', () => { + expect( + projectPublicEvent( + { type: 'session.idle', properties: { sessionID: foreignKiloSessionId } }, + kiloSessionId + ) + ).toBeNull(); + }); +}); + +describe('projectPublicEvent file URLs', () => { + it('redacts wrapper-local file URLs from loose event file parts and completed tool attachments', () => { + expect( + projectPublicEvent( + { + type: 'message.part.updated', + properties: { + sessionID: kiloSessionId, + part: { + id: 'prt_event_file', + sessionID: kiloSessionId, + messageID: 'msg_event_file', + type: 'file', + mime: 'text/plain', + url: 'file:///workspace/private/event.txt', + }, + }, + }, + kiloSessionId + ) + ).toMatchObject({ properties: { part: { url: '' } } }); + expect( + projectPublicEvent( + { + type: 'message.part.updated', + properties: { + sessionID: kiloSessionId, + part: { + id: 'prt_event_tool', + sessionID: kiloSessionId, + messageID: 'msg_event_tool', + type: 'tool', + state: { + status: 'completed', + attachments: [ + { + id: 'prt_event_attachment', + sessionID: kiloSessionId, + messageID: 'msg_event_tool', + type: 'file', + mime: 'text/plain', + url: 'file:///workspace/private/attachment.txt', + }, + { + id: 'prt_event_data_attachment', + sessionID: kiloSessionId, + messageID: 'msg_event_tool', + type: 'file', + mime: 'text/plain', + url: 'data:text/plain,safe', + }, + ], + }, + }, + }, + }, + kiloSessionId + ) + ).toMatchObject({ + properties: { + part: { state: { attachments: [{ url: '' }, { url: 'data:text/plain,safe' }] } }, + }, + }); + }); + + it('fails closed when private local file URIs remain outside reviewed redaction paths', () => { + expect( + projectPublicEvent( + { + type: 'message.part.updated', + properties: { + sessionID: kiloSessionId, + part: { + id: 'prt_resource', + sessionID: kiloSessionId, + messageID: 'msg_resource', + type: 'file', + mime: 'text/plain', + url: 'data:text/plain,safe', + source: { + type: 'resource', + uri: 'file:///workspace/private/resource.txt', + }, + }, + }, + }, + kiloSessionId + ) + ).toBeNull(); + expect( + projectPublicEvent( + { + type: 'todo.updated', + properties: { + sessionID: kiloSessionId, + artifact: { href: 'FILE:///workspace/private/unreviewed.txt' }, + }, + }, + kiloSessionId + ) + ).toBeNull(); + }); +}); + +describe('projectPublicStoredMessage', () => { + it('redacts wrapper-local file URLs from typed stored file parts while preserving data URLs', () => { + const projected = projectPublicStoredMessage( + { + info: { + id: 'msg_files', + sessionID: kiloSessionId, + role: 'user', + time: { created: 100 }, + agent: 'build', + model: { providerID: 'test', modelID: 'fake' }, + }, + parts: [ + { + id: 'prt_local_file', + sessionID: kiloSessionId, + messageID: 'msg_files', + type: 'file', + mime: 'text/plain', + url: 'file:///workspace/private/secret.txt', + }, + { + id: 'prt_data_file', + sessionID: kiloSessionId, + messageID: 'msg_files', + type: 'file', + mime: 'text/plain', + url: 'data:text/plain,safe', + }, + { + id: 'prt_tool', + sessionID: kiloSessionId, + messageID: 'msg_files', + type: 'tool', + callID: 'call_files', + tool: 'read', + state: { + status: 'completed', + input: {}, + output: 'safe', + title: 'read', + metadata: {}, + time: { start: 100, end: 101 }, + attachments: [ + { + id: 'prt_local_attachment', + sessionID: kiloSessionId, + messageID: 'msg_files', + type: 'file', + mime: 'text/plain', + url: 'file:///workspace/private/attachment.txt', + }, + ], + }, + }, + ], + }, + kiloSessionId + ); + + expect(projected.parts).toMatchObject([ + { url: '' }, + { url: 'data:text/plain,safe' }, + { state: { attachments: [{ url: '' }] } }, + ]); + }); + + it('exposes typed resource file URIs to the final fail-closed boundary check', () => { + const projected = projectPublicStoredMessage( + { + info: { + id: 'msg_resource', + sessionID: kiloSessionId, + role: 'user', + time: { created: 100 }, + agent: 'build', + model: { providerID: 'test', modelID: 'fake' }, + }, + parts: [ + { + id: 'prt_resource', + sessionID: kiloSessionId, + messageID: 'msg_resource', + type: 'file', + mime: 'text/plain', + url: 'data:text/plain,safe', + source: { + type: 'resource', + text: { value: 'private', start: 0, end: 7 }, + clientName: 'wrapper', + uri: 'file:///workspace/private/resource.txt', + }, + }, + ], + }, + kiloSessionId + ); + + expect(hasUnprojectedPrivateStructuredPath(projected)).toBe(true); + }); +}); diff --git a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts new file mode 100644 index 0000000000..69ac922235 --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts @@ -0,0 +1,449 @@ +import type { + CloudAgentRootSessionSummary, + KiloSdkAssistantMessage, + KiloSdkFilePart, + KiloSdkMessageInfo, + KiloSdkPart, + KiloSdkSessionInfo, + KiloSdkStoredMessage, + KiloSdkToolState, + KiloSdkUserMessage, +} from '../session-ingest-binding.js'; + +export function publicCloudAgentDirectory(kiloSessionId: string): string { + return `/cloud-agent/sessions/${kiloSessionId}`; +} + +export function projectPublicListedSession( + summary: CloudAgentRootSessionSummary +): KiloSdkSessionInfo { + return { + id: summary.kiloSessionId, + slug: summary.kiloSessionId, + projectID: 'cloud-agent', + directory: publicCloudAgentDirectory(summary.kiloSessionId), + title: summary.title ?? '', + version: 'cloud-agent', + time: { created: summary.created, updated: summary.updated }, + }; +} + +function isAbsoluteStructuredPath(value: string): boolean { + return /^(?:\/|[A-Za-z]:[\\/]|\\\\)/.test(value); +} + +function publicPath(path: string, kiloSessionId: string): string { + return isAbsoluteStructuredPath(path) ? publicCloudAgentDirectory(kiloSessionId) : path; +} + +function isPrivateStructuredPath(path: string): boolean { + return ( + isAbsoluteStructuredPath(path) && + path !== '/cloud-agent' && + !path.startsWith('/cloud-agent/sessions/') + ); +} + +export function projectPublicSession( + session: KiloSdkSessionInfo, + kiloSessionId: string +): KiloSdkSessionInfo { + const projected = { ...session }; + delete projected.path; + return { + ...projected, + directory: publicCloudAgentDirectory(kiloSessionId), + ...(projected.summary?.diffs + ? { + summary: { + ...projected.summary, + diffs: projected.summary.diffs.map(diff => ({ + ...diff, + file: publicPath(diff.file, kiloSessionId), + })), + }, + } + : {}), + ...(projected.permission + ? { + permission: projected.permission.map(rule => ({ + ...rule, + pattern: publicPath(rule.pattern, kiloSessionId), + })), + } + : {}), + }; +} + +function projectPublicUserMessage( + message: KiloSdkUserMessage, + kiloSessionId: string +): KiloSdkUserMessage { + return { + ...message, + ...(message.summary + ? { + summary: { + ...message.summary, + diffs: message.summary.diffs.map(diff => ({ + ...diff, + file: publicPath(diff.file, kiloSessionId), + })), + }, + } + : {}), + ...(message.editorContext + ? { + editorContext: { + ...message.editorContext, + ...(message.editorContext.visibleFiles + ? { + visibleFiles: message.editorContext.visibleFiles.map(file => + publicPath(file, kiloSessionId) + ), + } + : {}), + ...(message.editorContext.openTabs + ? { + openTabs: message.editorContext.openTabs.map(file => + publicPath(file, kiloSessionId) + ), + } + : {}), + ...(message.editorContext.activeFile + ? { activeFile: publicPath(message.editorContext.activeFile, kiloSessionId) } + : {}), + }, + } + : {}), + }; +} + +function projectPublicAssistantMessage( + message: KiloSdkAssistantMessage, + kiloSessionId: string +): KiloSdkAssistantMessage { + const directory = publicCloudAgentDirectory(kiloSessionId); + return { + ...message, + path: { cwd: directory, root: directory }, + }; +} + +export function projectPublicMessageInfo( + message: KiloSdkMessageInfo, + kiloSessionId: string +): KiloSdkMessageInfo { + return message.role === 'assistant' + ? projectPublicAssistantMessage(message, kiloSessionId) + : projectPublicUserMessage(message, kiloSessionId); +} + +function isLocalFileUri(value: string): boolean { + return /^file:/i.test(value); +} + +function publicFilePartUrl(url: string): string { + return isLocalFileUri(url) ? '' : url; +} + +function projectPublicFilePart(part: KiloSdkFilePart, kiloSessionId: string): KiloSdkFilePart { + const projected = { ...part, url: publicFilePartUrl(part.url) }; + if (!part.source || part.source.type === 'resource') return projected; + return { + ...projected, + source: { ...part.source, path: publicPath(part.source.path, kiloSessionId) }, + }; +} + +function projectPublicToolState(state: KiloSdkToolState, kiloSessionId: string): KiloSdkToolState { + if (state.status !== 'completed' || !state.attachments) return state; + return { + ...state, + attachments: state.attachments.map(part => projectPublicFilePart(part, kiloSessionId)), + }; +} + +export function projectPublicPart(part: KiloSdkPart, kiloSessionId: string): KiloSdkPart { + if (part.type === 'file') return projectPublicFilePart(part, kiloSessionId); + if (part.type === 'tool') { + return { ...part, state: projectPublicToolState(part.state, kiloSessionId) }; + } + if (part.type === 'patch') { + return { ...part, files: part.files.map(file => publicPath(file, kiloSessionId)) }; + } + return part; +} + +export function projectPublicStoredMessage( + message: KiloSdkStoredMessage, + kiloSessionId: string +): KiloSdkStoredMessage { + return { + info: projectPublicMessageInfo(message.info, kiloSessionId), + parts: message.parts.map(part => projectPublicPart(part, kiloSessionId)), + }; +} + +export function projectPublicStoredMessages( + messages: KiloSdkStoredMessage[], + kiloSessionId: string +): KiloSdkStoredMessage[] { + return messages.map(message => projectPublicStoredMessage(message, kiloSessionId)); +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function projectLooseDiffs(value: unknown, kiloSessionId: string): unknown { + if (!Array.isArray(value)) return value; + return value.map((diff: unknown) => + isRecord(diff) && typeof diff.file === 'string' + ? { ...diff, file: publicPath(diff.file, kiloSessionId) } + : diff + ); +} + +function projectLooseSessionInfo( + info: Record, + kiloSessionId: string +): Record { + const projected = { ...info }; + delete projected.path; + return { + ...projected, + directory: publicCloudAgentDirectory(kiloSessionId), + ...(isRecord(projected.summary) + ? { + summary: { + ...projected.summary, + diffs: projectLooseDiffs(projected.summary.diffs, kiloSessionId), + }, + } + : {}), + ...(Array.isArray(projected.permission) + ? { + permission: projected.permission.map((rule: unknown) => + isRecord(rule) && typeof rule.pattern === 'string' + ? { ...rule, pattern: publicPath(rule.pattern, kiloSessionId) } + : rule + ), + } + : {}), + }; +} + +function projectLooseMessageInfo>( + info: T, + kiloSessionId: string +): T { + if (info.role === 'assistant' && isRecord(info.path)) { + const directory = publicCloudAgentDirectory(kiloSessionId); + return { ...info, path: { ...info.path, cwd: directory, root: directory } }; + } + if (info.role !== 'user') return info; + return { + ...info, + ...(isRecord(info.summary) + ? { + summary: { + ...info.summary, + diffs: projectLooseDiffs(info.summary.diffs, kiloSessionId), + }, + } + : {}), + ...(isRecord(info.editorContext) + ? { + editorContext: { + ...info.editorContext, + ...(Array.isArray(info.editorContext.visibleFiles) + ? { + visibleFiles: info.editorContext.visibleFiles.map((file: unknown) => + typeof file === 'string' ? publicPath(file, kiloSessionId) : file + ), + } + : {}), + ...(Array.isArray(info.editorContext.openTabs) + ? { + openTabs: info.editorContext.openTabs.map((file: unknown) => + typeof file === 'string' ? publicPath(file, kiloSessionId) : file + ), + } + : {}), + ...(typeof info.editorContext.activeFile === 'string' + ? { activeFile: publicPath(info.editorContext.activeFile, kiloSessionId) } + : {}), + }, + } + : {}), + }; +} + +function projectLooseFilePart>( + part: T, + kiloSessionId: string +): T { + if (part.type !== 'file') return part; + const projected = + typeof part.url === 'string' ? { ...part, url: publicFilePartUrl(part.url) } : part; + if (!isRecord(part.source)) return projected; + if ( + (part.source.type === 'file' || part.source.type === 'symbol') && + typeof part.source.path === 'string' + ) { + return { + ...projected, + source: { ...part.source, path: publicPath(part.source.path, kiloSessionId) }, + }; + } + return projected; +} + +function projectLoosePart>(part: T, kiloSessionId: string): T { + const filePart = projectLooseFilePart(part, kiloSessionId); + if (filePart.type === 'patch' && Array.isArray(filePart.files)) { + return { + ...filePart, + files: filePart.files.map((file: unknown) => + typeof file === 'string' ? publicPath(file, kiloSessionId) : file + ), + }; + } + if (filePart.type !== 'tool' || !isRecord(filePart.state)) return filePart; + if (filePart.state.status !== 'completed' || !Array.isArray(filePart.state.attachments)) { + return filePart; + } + return { + ...filePart, + state: { + ...filePart.state, + attachments: filePart.state.attachments.map((attachment: unknown) => + isRecord(attachment) ? projectLooseFilePart(attachment, kiloSessionId) : attachment + ), + }, + }; +} + +export type PublicProjectableEvent = { + id?: string; + type: string; + properties: Record; +}; + +export function hasUnprojectedPrivateStructuredPath(value: unknown, key?: string): boolean { + if (typeof value === 'string') { + return ( + isLocalFileUri(value) || + (key !== undefined && + ['path', 'directory', 'cwd', 'root', 'file', 'pattern', 'url'].includes(key) && + isPrivateStructuredPath(value)) + ); + } + if (Array.isArray(value)) { + return value.some(item => + hasUnprojectedPrivateStructuredPath(item, key === 'files' ? 'file' : key) + ); + } + if (!isRecord(value)) return false; + return Object.entries(value).some(([childKey, child]) => + hasUnprojectedPrivateStructuredPath(child, childKey) + ); +} + +function projectSessionEventProperties( + properties: Record, + kiloSessionId: string +): Record | null { + const info = properties.info; + if (!isRecord(info)) return hasUnprojectedPrivateStructuredPath(properties) ? null : properties; + if (info.id !== kiloSessionId) return null; + const projected = { ...properties, info: projectLooseSessionInfo(info, kiloSessionId) }; + return hasUnprojectedPrivateStructuredPath(projected) ? null : projected; +} + +function hasConflictingEntitySessionIdentity( + entity: Record, + kiloSessionId: string +): boolean { + return ( + Object.prototype.hasOwnProperty.call(entity, 'sessionID') && entity.sessionID !== kiloSessionId + ); +} + +function hasConflictingMessageEventIdentity( + properties: Record, + kiloSessionId: string +): boolean { + return ( + isRecord(properties.info) && hasConflictingEntitySessionIdentity(properties.info, kiloSessionId) + ); +} + +function hasConflictingPartEventIdentity( + properties: Record, + kiloSessionId: string +): boolean { + const part = properties.part; + if (!isRecord(part)) return false; + if (hasConflictingEntitySessionIdentity(part, kiloSessionId)) return true; + if (part.type !== 'tool' || !isRecord(part.state) || !Array.isArray(part.state.attachments)) { + return false; + } + return part.state.attachments.some( + attachment => + isRecord(attachment) && hasConflictingEntitySessionIdentity(attachment, kiloSessionId) + ); +} + +function projectMessageEventProperties( + properties: Record, + kiloSessionId: string +): Record | null { + const info = properties.info; + if (!isRecord(info)) return hasUnprojectedPrivateStructuredPath(properties) ? null : properties; + if (info.role !== 'user' && info.role !== 'assistant') { + return hasUnprojectedPrivateStructuredPath(properties) ? null : properties; + } + const projected = { ...properties, info: projectLooseMessageInfo(info, kiloSessionId) }; + return hasUnprojectedPrivateStructuredPath(projected) ? null : projected; +} + +function projectPartEventProperties( + properties: Record, + kiloSessionId: string +): Record | null { + const part = properties.part; + if (!isRecord(part) || typeof part.type !== 'string') { + return hasUnprojectedPrivateStructuredPath(properties) ? null : properties; + } + const projected = { ...properties, part: projectLoosePart(part, kiloSessionId) }; + return hasUnprojectedPrivateStructuredPath(projected) ? null : projected; +} + +export function projectPublicEvent( + event: PublicProjectableEvent, + kiloSessionId: string +): PublicProjectableEvent | null { + if (event.properties.sessionID !== kiloSessionId) return null; + let properties: Record | null; + switch (event.type) { + case 'session.updated': + properties = projectSessionEventProperties(event.properties, kiloSessionId); + break; + case 'message.updated': + properties = hasConflictingMessageEventIdentity(event.properties, kiloSessionId) + ? null + : projectMessageEventProperties(event.properties, kiloSessionId); + break; + case 'message.part.updated': + properties = hasConflictingPartEventIdentity(event.properties, kiloSessionId) + ? null + : projectPartEventProperties(event.properties, kiloSessionId); + break; + default: + properties = hasUnprojectedPrivateStructuredPath(event.properties) ? null : event.properties; + break; + } + return properties ? { ...event, properties } : null; +} diff --git a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts new file mode 100644 index 0000000000..bcf731f402 --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts @@ -0,0 +1,85 @@ +import { getSandbox } from '@cloudflare/sandbox'; +import { findWrapperForSession } from '../kilo/wrapper-manager.js'; +import { generateSandboxId, getSandboxNamespace } from '../sandbox-id.js'; +import { fetchSessionMetadata } from '../session-service.js'; +import type { Env, SandboxInstance, SandboxId, SessionId } from '../types.js'; + +export type SessionKiloFacadeDecision = + | { kind: 'proxy-live-wrapper' } + | { kind: 'reject'; status: number; code: string; message: string }; + +export type SessionKiloFacadePolicyInput = { + method: string; + kiloRelativePath: string; + search: string; + userId: string; + kiloSessionId: string; + cloudAgentSessionId: string; +}; + +export type LiveWrapperTarget = { + sandbox: SandboxInstance; + port: number; +}; + +export function decideSessionKiloFacadeRoute( + input: SessionKiloFacadePolicyInput +): SessionKiloFacadeDecision { + const suffix = input.kiloRelativePath.slice( + `/session/${encodeURIComponent(input.kiloSessionId)}`.length + ); + const supported = + (input.method === 'GET' && (suffix === '' || suffix === '/message')) || + (input.method === 'POST' && (suffix === '/prompt_async' || suffix === '/abort')); + if (!supported) { + return { + kind: 'reject', + status: 501, + code: 'KILO_ROUTE_UNSUPPORTED', + message: 'Kilo facade route is not supported', + }; + } + return { kind: 'proxy-live-wrapper' }; +} + +export function buildWrapperKiloProxyUrl(params: { + wrapperPort: number; + kiloRelativePath: string; + search: string; +}): string { + const url = new URL(`http://localhost:${params.wrapperPort}/kilo-proxy`); + url.pathname = `/kilo-proxy${params.kiloRelativePath}`; + url.search = params.search; + return url.toString(); +} + +export async function resolveLiveWrapperTarget(params: { + env: Env; + userId: string; + cloudAgentSessionId: string; +}): Promise { + const { env, userId, cloudAgentSessionId } = params; + const metadata = await fetchSessionMetadata(env, userId, cloudAgentSessionId); + if (!metadata) { + return null; + } + + const sessionId = cloudAgentSessionId as SessionId; + const sandboxId: SandboxId = + metadata.workspace?.sandboxId ?? + (await generateSandboxId( + env.PER_SESSION_SANDBOX_ORG_IDS, + metadata.identity.orgId, + userId, + metadata.identity.sessionId, + metadata.identity.botId + )); + + const sandbox = getSandbox(getSandboxNamespace(env, sandboxId), sandboxId); + const wrapperInfo = await findWrapperForSession(sandbox, sessionId); + if (!wrapperInfo) { + return null; + } + + return { sandbox, port: wrapperInfo.port }; +} diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts new file mode 100644 index 0000000000..09ca173728 --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts @@ -0,0 +1,3072 @@ +import { + encodeKiloSdkMessagesCursor, + MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE, +} from '@kilocode/session-ingest-contracts'; +import { TRPCError } from '@trpc/server'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const { preflightExistingPromptModelMock } = vi.hoisted(() => ({ + preflightExistingPromptModelMock: vi.fn(), +})); + +vi.mock('../session/model-preflight.js', () => ({ + preflightExistingPromptModel: preflightExistingPromptModelMock, +})); + +vi.mock('cloudflare:workers', () => ({ + DurableObject: class DurableObject { + ctx: unknown; + env: unknown; + + constructor(ctx: unknown, env: unknown) { + this.ctx = ctx; + this.env = env; + } + }, +})); + +vi.mock('@cloudflare/sandbox', () => ({ + getSandbox: vi.fn(), + Sandbox: class Sandbox {}, +})); + +import type { KiloSdkStoredMessage } from '../session-ingest-binding'; +import type { Env } from '../types'; +import { + handleKiloFacadeRequest, + isSyntheticGlobalEvent, + publicCloudAgentDirectory, + rewriteGlobalEventDirectory, + UserKiloFacade, +} from './user-kilo-facade'; +import type { LiveWrapperTarget, SessionKiloFacadeDecision } from './session-proxy'; + +const kiloSessionId = 'ses_12345678901234567890123456'; +type ContainerFetch = (request: Request, port: number) => Promise; + +function envStub(): Env { + return { + SESSION_INGEST: { + resolveCloudAgentRootSessionForKiloSession: vi.fn(), + getCloudAgentRootSessionSnapshot: vi.fn(), + listCloudAgentRootSessions: vi.fn(), + getCloudAgentRootSessionMessages: vi.fn(), + }, + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ + hasMessageAdmission: vi.fn().mockResolvedValue(false), + admitSubmittedMessage: vi.fn(), + interruptExecution: vi.fn(), + })), + }, + } as unknown as Env; +} + +function sdkSessionInfo(id: string, directory = '/workspace/root') { + return { + id, + slug: 'rainy-fox', + projectID: 'global', + directory, + path: '/workspace/private/session', + summary: { + additions: 1, + deletions: 0, + files: 1, + diffs: [{ file: '/workspace/private/changed.ts', additions: 1, deletions: 0 }], + }, + permission: [{ permission: 'read', pattern: '/workspace/private', action: 'allow' as const }], + title: 'SDK session fixture', + version: '1.0.0', + time: { created: 100, updated: 200 }, + }; +} + +function projectedSdkSessionInfo(id: string) { + const { path: _privatePath, ...info } = sdkSessionInfo(id); + return { + ...info, + directory: publicCloudAgentDirectory(id), + summary: { + ...info.summary, + diffs: [ + { + ...info.summary.diffs[0], + file: publicCloudAgentDirectory(id), + }, + ], + }, + permission: [{ permission: 'read', pattern: publicCloudAgentDirectory(id), action: 'allow' }], + }; +} + +function listedSdkSessionInfo(id: string) { + return { + id, + slug: id, + projectID: 'cloud-agent', + directory: publicCloudAgentDirectory(id), + title: 'Mapped session', + version: 'cloud-agent', + time: { created: 100, updated: 200 }, + }; +} + +function sdkMessageHistory(): KiloSdkStoredMessage[] { + return [ + { + info: { + id: 'msg_user_1', + sessionID: kiloSessionId, + role: 'user' as const, + time: { created: 110 }, + summary: { + diffs: [ + { + file: '/workspace/private/user-change.ts', + patch: 'unchanged text payload', + additions: 1, + deletions: 0, + }, + ], + }, + editorContext: { + visibleFiles: ['/workspace/private/visible.ts', 'src/public.ts'], + openTabs: ['/workspace/private/open.ts'], + activeFile: '/workspace/private/active.ts', + }, + agent: 'build', + model: { providerID: 'test', modelID: 'fake' }, + }, + parts: [ + { + id: 'prt_user_1', + sessionID: kiloSessionId, + messageID: 'msg_user_1', + type: 'text' as const, + text: 'hello', + }, + ], + }, + { + info: { + id: 'msg_assistant_1', + sessionID: kiloSessionId, + role: 'assistant' as const, + time: { created: 120, completed: 130 }, + parentID: 'msg_user_1', + modelID: 'fake', + providerID: 'test', + mode: 'build', + agent: 'build', + path: { cwd: '/workspace/private/session', root: '/workspace/private' }, + cost: 0, + tokens: { input: 1, output: 1, reasoning: 0, cache: { read: 0, write: 0 } }, + }, + parts: [ + { + id: 'prt_file_1', + sessionID: kiloSessionId, + messageID: 'msg_assistant_1', + type: 'file' as const, + mime: 'text/plain', + url: 'data:text/plain,private', + source: { + type: 'file' as const, + text: { value: 'private', start: 0, end: 7 }, + path: '/workspace/private/source.ts', + }, + }, + { + id: 'prt_tool_1', + sessionID: kiloSessionId, + messageID: 'msg_assistant_1', + type: 'tool' as const, + callID: 'call_1', + tool: 'read', + state: { + status: 'completed' as const, + input: {}, + output: 'unchanged model-authored /workspace/private text', + title: 'read', + metadata: {}, + time: { start: 120, end: 121 }, + attachments: [ + { + id: 'prt_attachment_1', + sessionID: kiloSessionId, + messageID: 'msg_assistant_1', + type: 'file' as const, + mime: 'text/plain', + url: 'data:text/plain,symbol', + source: { + type: 'symbol' as const, + text: { value: 'symbol', start: 0, end: 6 }, + path: '/workspace/private/symbol.ts', + range: { + start: { line: 1, character: 0 }, + end: { line: 1, character: 6 }, + }, + name: 'symbol', + kind: 1, + }, + }, + ], + }, + }, + ], + }, + ]; +} + +function projectedSdkMessageHistory(): KiloSdkStoredMessage[] { + const history = sdkMessageHistory(); + const userMessage = history[0]; + const assistantMessage = history[1]; + if (!userMessage || userMessage.info.role !== 'user' || !userMessage.info.summary) { + throw new Error('Expected user message projection fixture'); + } + if (!assistantMessage || assistantMessage.info.role !== 'assistant') { + throw new Error('Expected assistant message projection fixture'); + } + const filePart = assistantMessage.parts[0]; + const toolPart = assistantMessage.parts[1]; + if ( + !filePart || + filePart.type !== 'file' || + !filePart.source || + filePart.source.type === 'resource' + ) { + throw new Error('Expected file source fixture'); + } + if ( + !toolPart || + toolPart.type !== 'tool' || + toolPart.state.status !== 'completed' || + !toolPart.state.attachments?.[0] + ) { + throw new Error('Expected tool attachment fixture'); + } + const attachment = toolPart.state.attachments[0]; + if (!attachment.source || attachment.source.type === 'resource') { + throw new Error('Expected attachment source fixture'); + } + const directory = publicCloudAgentDirectory(kiloSessionId); + return [ + { + ...userMessage, + info: { + ...userMessage.info, + summary: { + diffs: [{ ...userMessage.info.summary.diffs[0], file: directory }], + }, + editorContext: { + visibleFiles: [directory, 'src/public.ts'], + openTabs: [directory], + activeFile: directory, + }, + }, + }, + { + ...assistantMessage, + info: { + ...assistantMessage.info, + path: { cwd: directory, root: directory }, + }, + parts: [ + { + ...filePart, + source: { ...filePart.source, path: directory }, + }, + { + ...toolPart, + state: { + ...toolPart.state, + attachments: [ + { + ...attachment, + source: { + ...attachment.source, + path: directory, + }, + }, + ], + }, + }, + ], + }, + ]; +} + +function liveWrapperTarget(containerFetch: ContainerFetch): LiveWrapperTarget { + return { + port: 5123, + sandbox: { containerFetch } as unknown as LiveWrapperTarget['sandbox'], + }; +} + +async function readSseTextFrame(reader: ReadableStreamDefaultReader): Promise { + const frame = await reader.read(); + if (frame.done) { + throw new Error('Expected SSE frame'); + } + return new TextDecoder().decode(frame.value); +} + +async function readSseFrame(reader: ReadableStreamDefaultReader): Promise { + return JSON.parse((await readSseTextFrame(reader)).slice('data: '.length)); +} + +describe('handleKiloFacadeRequest', () => { + beforeEach(() => { + vi.resetAllMocks(); + preflightExistingPromptModelMock.mockResolvedValue(undefined); + }); + + it('returns 501 for intentionally unsupported Kilo route families', async () => { + for (const request of [ + new Request('http://worker.test/kilo/session', { method: 'POST' }), + new Request('http://worker.test/kilo/session/status'), + new Request('http://worker.test/kilo/session/viewed', { method: 'POST' }), + new Request('http://worker.test/kilo/sync/history', { method: 'POST' }), + new Request('http://worker.test/kilo/config'), + new Request('http://worker.test/kilo/provider'), + new Request('http://worker.test/kilo/project/current'), + new Request('http://worker.test/kilo/global/health'), + ]) { + const response = await handleKiloFacadeRequest({ + request, + env: envStub(), + userId: 'usr_1', + deps: { + globalEvents: { + openPublicGlobalEventStream: () => new Response('unused'), + openPublicSessionEventStream: () => new Response('unused'), + }, + }, + }); + + expect(response.status).toBe(501); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_ROUTE_UNSUPPORTED' }); + } + }); + + it('lists mapped Cloud Agent roots without requiring private SDK snapshots', async () => { + const env = envStub(); + const listCloudAgentRootSessions = vi.mocked(env.SESSION_INGEST.listCloudAgentRootSessions); + listCloudAgentRootSessions.mockResolvedValue([ + { + kiloSessionId, + cloudAgentSessionId: 'agent_live', + title: 'Mapped session', + created: 100, + updated: 200, + }, + ]); + + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/session?limit=15&start=1234'), + env, + userId: 'usr_1', + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual([listedSdkSessionInfo(kiloSessionId)]); + expect(listCloudAgentRootSessions).toHaveBeenCalledWith({ + kiloUserId: 'usr_1', + limit: 15, + start: 1234, + }); + }); + + it('rejects selectors on the global event route before opening the stream', async () => { + const openPublicGlobalEventStream = vi.fn(); + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/global/event?directory=%2Fworkspace%2Fprivate'), + env: envStub(), + userId: 'usr_1', + deps: { + globalEvents: { openPublicGlobalEventStream, openPublicSessionEventStream: vi.fn() }, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_EVENT_SELECTOR_UNSUPPORTED', + }); + expect(openPublicGlobalEventStream).not.toHaveBeenCalled(); + }); + + it('opens the public global event stream through the facade coordinator', async () => { + const openPublicGlobalEventStream = vi.fn( + () => new Response('stream', { headers: { 'Content-Type': 'text/event-stream' } }) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/global/event'), + env: envStub(), + userId: 'usr_1', + deps: { + globalEvents: { openPublicGlobalEventStream, openPublicSessionEventStream: vi.fn() }, + }, + }); + + expect(response.status).toBe(200); + expect(response.headers.get('content-type')).toContain('text/event-stream'); + expect(openPublicGlobalEventStream).toHaveBeenCalledOnce(); + }); + + it('requires a directory selector for a scoped event subscription', async () => { + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/event'), + env: envStub(), + userId: 'usr_1', + deps: { + globalEvents: { + openPublicGlobalEventStream: vi.fn(), + openPublicSessionEventStream: vi.fn(), + }, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_EVENT_DIRECTORY_REQUIRED', + }); + }); + + it('rejects unsupported scoped event selector semantics before ownership lookup', async () => { + const resolveRootSessionForKiloSession = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/event?workspace=private'), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession, + globalEvents: { + openPublicGlobalEventStream: vi.fn(), + openPublicSessionEventStream: vi.fn(), + }, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_EVENT_SELECTOR_UNSUPPORTED', + }); + expect(resolveRootSessionForKiloSession).not.toHaveBeenCalled(); + }); + + it('rejects a non-Cloud-Agent scoped event directory selector', async () => { + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/event?directory=%2Fworkspace%2Fprivate'), + env: envStub(), + userId: 'usr_1', + deps: { + globalEvents: { + openPublicGlobalEventStream: vi.fn(), + openPublicSessionEventStream: vi.fn(), + }, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_EVENT_SELECTOR_UNSUPPORTED', + }); + }); + + it('hides a malformed Cloud Agent session event selector as not found', async () => { + const resolveRootSessionForKiloSession = vi.fn(); + const response = await handleKiloFacadeRequest({ + request: new Request( + 'http://worker.test/kilo/event?directory=%2Fcloud-agent%2Fsessions%2Fbad-id' + ), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession, + globalEvents: { + openPublicGlobalEventStream: vi.fn(), + openPublicSessionEventStream: vi.fn(), + }, + }, + }); + + expect(response.status).toBe(404); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_SESSION_NOT_FOUND' }); + expect(resolveRootSessionForKiloSession).not.toHaveBeenCalled(); + }); + + it('opens a scoped event stream for an owned root without resolving a live wrapper', async () => { + const openPublicSessionEventStream = vi.fn( + () => new Response('stream', { headers: { 'Content-Type': 'text/event-stream' } }) + ); + const resolveLiveWrapper = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/event?directory=${encodeURIComponent(publicCloudAgentDirectory(kiloSessionId))}` + ), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper, + globalEvents: { openPublicGlobalEventStream: vi.fn(), openPublicSessionEventStream }, + }, + }); + + expect(response.status).toBe(200); + expect(openPublicSessionEventStream).toHaveBeenCalledWith(kiloSessionId); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + }); + + it('hides a foreign scoped event selector as not found', async () => { + const openPublicSessionEventStream = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/event?directory=${encodeURIComponent(publicCloudAgentDirectory(kiloSessionId))}` + ), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => null), + globalEvents: { openPublicGlobalEventStream: vi.fn(), openPublicSessionEventStream }, + }, + }); + + expect(response.status).toBe(404); + expect(openPublicSessionEventStream).not.toHaveBeenCalled(); + }); + + it('returns 404 for unresolved root Kilo sessions before live wrapper lookup', async () => { + const resolveLiveWrapper = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => null), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(404); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + }); + + it('returns 404 for malformed Kilo session IDs before resolving ownership', async () => { + const resolveRootSessionForKiloSession = vi.fn(async () => { + throw new Error('Malformed session IDs should not reach session ingest'); + }); + + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/session/not-a-session/message'), + env: envStub(), + userId: 'usr_1', + deps: { resolveRootSessionForKiloSession }, + }); + + expect(response.status).toBe(404); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_SESSION_NOT_FOUND', + }); + expect(resolveRootSessionForKiloSession).not.toHaveBeenCalled(); + }); + + it('runs the session route policy seam before live wrapper lookup', async () => { + const resolveLiveWrapper = vi.fn(); + const decision: SessionKiloFacadeDecision = { + kind: 'reject', + status: 418, + code: 'POLICY_REJECTED', + message: 'rejected by policy', + }; + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + decideSessionRoute: vi.fn(() => decision), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(418); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + }); + + it('returns a persisted session snapshot when an authorized detail read is cold', async () => { + const env = envStub(); + const getCloudAgentRootSessionSnapshot = vi.mocked( + env.SESSION_INGEST.getCloudAgentRootSessionSnapshot + ); + getCloudAgentRootSessionSnapshot.mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + snapshot: { kind: 'value', info: sdkSessionInfo(kiloSessionId), byteLength: 123 }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual(projectedSdkSessionInfo(kiloSessionId)); + expect(getCloudAgentRootSessionSnapshot).toHaveBeenCalledWith({ + kiloUserId: 'usr_1', + kiloSessionId, + }); + }); + + it('falls back to the persisted detail snapshot when the wrapper has no private Kilo runtime', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_live', + snapshot: { kind: 'value', info: sdkSessionInfo(kiloSessionId), byteLength: 123 }, + }); + const containerFetch = vi.fn(async () => + Response.json( + { error: 'KILO_RUNTIME_UNAVAILABLE', message: 'Kilo runtime is not bootstrapped' }, + { status: 503 } + ) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + directory: publicCloudAgentDirectory(kiloSessionId), + }); + }); + + it('falls back to the persisted detail snapshot when live wrapper discovery is unavailable', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_live', + snapshot: { kind: 'value', info: sdkSessionInfo(kiloSessionId), byteLength: 123 }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => { + throw new Error('wrapper discovery failed'); + }), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + directory: publicCloudAgentDirectory(kiloSessionId), + }); + }); + + it('falls back to the persisted detail snapshot when a live proxy request cannot be reached', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_live', + snapshot: { kind: 'value', info: sdkSessionInfo(kiloSessionId), byteLength: 123 }, + }); + const containerFetch = vi.fn(async () => { + throw new Error('wrapper request failed'); + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + directory: publicCloudAgentDirectory(kiloSessionId), + }); + }); + + it('fails closed when a cold session snapshot contains an unreviewed structured path', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + snapshot: { + kind: 'value', + byteLength: 123, + info: { + ...sdkSessionInfo(kiloSessionId), + share: { url: '/workspace/private/share' }, + }, + }, + }); + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(502); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + }); + }); + + it('returns a stable pending response for a cold detail read without a materialized snapshot', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + snapshot: { kind: 'pending' }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(503); + expect(response.headers.get('retry-after')).toBe('1'); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_SESSION_SNAPSHOT_PENDING', + }); + }); + + it('maps an oversized cold session snapshot to stable 413', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + snapshot: { kind: 'too_large', maximumBytes: 8 * 1024 * 1024 }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(413); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_SESSION_SNAPSHOT_TOO_LARGE', + }); + }); + + it('maps a retryable cold session snapshot failure to stable retryable 503', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + snapshot: { kind: 'retryable_failure' }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(503); + expect(response.headers.get('retry-after')).toBe('1'); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_SESSION_READ_RETRYABLE' }); + }); + + it('maps invalid persisted snapshot data to the warm-path upstream invalid response', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + snapshot: { kind: 'invalid_data' }, + } as never); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(502); + await expect(response.json()).resolves.toEqual({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + message: 'Kilo session response is not valid', + }); + }); + + it('falls back to the persisted detail snapshot when the live private Kilo request fails', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_live', + snapshot: { kind: 'value', info: sdkSessionInfo(kiloSessionId), byteLength: 123 }, + }); + const containerFetch = vi.fn(async () => + Response.json({ error: 'KILO_PROXY_ERROR', message: 'fetch failed' }, { status: 502 }) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + directory: publicCloudAgentDirectory(kiloSessionId), + }); + }); + + it('rewrites a fresh live session detail response without exposing its private directory', async () => { + const containerFetch = vi.fn(async () => + Response.json(sdkSessionInfo(kiloSessionId), { + headers: { 'X-Upstream': 'kept' }, + }) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + expect(response.headers.get('x-upstream')).toBe('kept'); + await expect(response.json()).resolves.toEqual(projectedSdkSessionInfo(kiloSessionId)); + }); + + it('bounds successful live detail JSON before projection', async () => { + const containerFetch = vi.fn( + async () => + new Response(JSON.stringify(sdkSessionInfo(kiloSessionId)), { + status: 200, + headers: { 'Content-Length': String(8 * 1024 * 1024 + 1) }, + }) + ); + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(502); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + }); + }); + + it('bounds live detail without consuming an oversized streamed body', async () => { + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => + liveWrapperTarget( + async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(8 * 1024 * 1024 + 1)); + controller.close(); + }, + }) + ) + ) + ), + }, + }); + + expect(response.status).toBe(502); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + }); + }); + + it('rejects malformed successful live detail JSON instead of returning an SDK-invalid session', async () => { + const containerFetch = vi.fn(async () => + Response.json({ directory: '/workspace/root' }) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(502); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + }); + }); + + it('rejects unsupported SDK session detail selectors instead of ignoring them', async () => { + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}?directory=%2Fworkspace%2Fprivate` + ), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(), + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_SESSION_SELECTOR_UNSUPPORTED', + }); + }); + + it('rejects unsupported session list selectors instead of broadening the result', async () => { + const env = envStub(); + + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/session?directory=%2Fworkspace%2Fprivate'), + env, + userId: 'usr_1', + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_SESSION_LIST_SELECTOR_UNSUPPORTED', + }); + expect(env.SESSION_INGEST.listCloudAgentRootSessions).not.toHaveBeenCalled(); + }); + + it('returns persisted SDK messages and public pagination headers when a read is cold', async () => { + const env = envStub(); + const nextCursor = 'eyJpZCI6Im1zZ191c2VyXzEiLCJ0aW1lIjoxMTB9'; + const getCloudAgentRootSessionMessages = vi.mocked( + env.SESSION_INGEST.getCloudAgentRootSessionMessages + ); + getCloudAgentRootSessionMessages.mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + history: { messages: sdkMessageHistory(), nextCursor, omittedItemCount: 0 }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message?limit=1`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); + expect(response.headers.get('x-kilo-omitted-item-count')).toBe('0'); + expect(response.headers.get('x-next-cursor')).toBe(nextCursor); + expect(response.headers.get('access-control-expose-headers')).toBe( + 'X-Kilo-Omitted-Item-Count, Link, X-Next-Cursor' + ); + expect(response.headers.get('link')).toBe( + `; rel="next"` + ); + expect(getCloudAgentRootSessionMessages).toHaveBeenCalledWith({ + kiloUserId: 'usr_1', + kiloSessionId, + limit: 1, + }); + }); + + it('returns persisted omission metadata without changing the projected message array body', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + history: { messages: sdkMessageHistory(), nextCursor: null, omittedItemCount: 3 }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message?limit=1`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); + expect(response.headers.get('x-kilo-omitted-item-count')).toBe('3'); + expect(response.headers.get('access-control-expose-headers')).toBe('X-Kilo-Omitted-Item-Count'); + expect(response.headers.get('link')).toBeNull(); + expect(response.headers.get('x-next-cursor')).toBeNull(); + }); + + it('keeps legacy message cursors on the live-first path', async () => { + const env = envStub(); + const cursor = encodeKiloSdkMessagesCursor({ id: 'msg_user_1', time: 110 }); + const containerFetch = vi.fn(async () => Response.json(sdkMessageHistory())); + const resolveLiveWrapper = vi.fn(async () => liveWrapperTarget(containerFetch)); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/message?limit=1&before=${cursor}` + ), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); + expect(resolveLiveWrapper).toHaveBeenCalledOnce(); + expect(containerFetch).toHaveBeenCalledOnce(); + expect(env.SESSION_INGEST.getCloudAgentRootSessionMessages).not.toHaveBeenCalled(); + expect(response.headers.get('x-kilo-omitted-item-count')).toBeNull(); + }); + + it('rejects invalid message pagination before durable lookup', async () => { + const env = envStub(); + const resolveLiveWrapper = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/message?limit=1&before=not-a-cursor` + ), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_QUERY_INVALID' }); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + expect(env.SESSION_INGEST.getCloudAgentRootSessionMessages).not.toHaveBeenCalled(); + }); + + it('rejects a decodable non-message cursor before durable lookup', async () => { + const env = envStub(); + const resolveLiveWrapper = vi.fn(); + const cursor = btoa(JSON.stringify({ id: 'other_01', time: 1 })).replace(/=+$/g, ''); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/message?limit=1&before=${cursor}` + ), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_QUERY_INVALID' }); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + expect(env.SESSION_INGEST.getCloudAgentRootSessionMessages).not.toHaveBeenCalled(); + }); + + it('rejects message page limits above the shared maximum before live or persisted reads', async () => { + const env = envStub(); + const resolveLiveWrapper = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/message?limit=${MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE + 1}` + ), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_QUERY_INVALID' }); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + expect(env.SESSION_INGEST.getCloudAgentRootSessionMessages).not.toHaveBeenCalled(); + }); + + it('rejects repeated message pagination parameters before live or persisted reads', async () => { + const env = envStub(); + const resolveLiveWrapper = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/message?limit=1&limit=${MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE + 1}` + ), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_QUERY_INVALID' }); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + expect(env.SESSION_INGEST.getCloudAgentRootSessionMessages).not.toHaveBeenCalled(); + }); + + it('falls back to persisted messages when the wrapper has no private Kilo runtime', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_live', + history: { messages: sdkMessageHistory(), nextCursor: null, omittedItemCount: 0 }, + }); + const containerFetch = vi.fn(async () => + Response.json( + { error: 'KILO_RUNTIME_UNAVAILABLE', message: 'Kilo runtime is not bootstrapped' }, + { status: 503 } + ) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); + }); + + it('falls back to persisted messages when the live private Kilo request fails', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_live', + history: { messages: sdkMessageHistory(), nextCursor: null, omittedItemCount: 0 }, + }); + const containerFetch = vi.fn(async () => + Response.json({ error: 'KILO_PROXY_ERROR', message: 'fetch failed' }, { status: 502 }) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); + }); + + it('falls back to persisted messages when a live proxy request cannot be reached', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_live', + history: { messages: sdkMessageHistory(), nextCursor: null, omittedItemCount: 0 }, + }); + const containerFetch = vi.fn(async () => { + throw new Error('wrapper request failed'); + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); + }); + + it('continues forwarding live message reads for freshest available history without private selectors', async () => { + const containerFetch = vi.fn(async () => Response.json(sdkMessageHistory())); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/message?directory=${encodeURIComponent(publicCloudAgentDirectory(kiloSessionId))}&limit=1` + ), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); + const [forwardedRequest] = containerFetch.mock.calls[0]; + expect(new URL(forwardedRequest.url).search).toBe('?limit=1'); + }); + + it('bounds and validates successful live message JSON before public projection', async () => { + const oversizedResponse = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => + liveWrapperTarget( + async () => + new Response(JSON.stringify(sdkMessageHistory()), { + status: 200, + headers: { 'Content-Length': String(8 * 1024 * 1024 + 1) }, + }) + ) + ), + }, + }); + expect(oversizedResponse.status).toBe(502); + await expect(oversizedResponse.json()).resolves.toMatchObject({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + }); + + const malformedResponse = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => + liveWrapperTarget(async () => Response.json([{ info: {} }])) + ), + }, + }); + expect(malformedResponse.status).toBe(502); + await expect(malformedResponse.json()).resolves.toMatchObject({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + }); + }); + + it('bounds live messages without consuming an oversized streamed body', async () => { + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => + liveWrapperTarget( + async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(8 * 1024 * 1024 + 1)); + controller.close(); + }, + }) + ) + ) + ), + }, + }); + + expect(response.status).toBe(502); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + }); + }); + + it('fails closed when cold messages contain an unreviewed structured path', async () => { + const env = envStub(); + const history = sdkMessageHistory(); + const assistant = history[1]; + if (!assistant || assistant.info.role !== 'assistant') { + throw new Error('Expected assistant message fixture'); + } + history[1] = { + ...assistant, + parts: [ + ...assistant.parts, + { + id: 'prt_metadata_private', + sessionID: kiloSessionId, + messageID: assistant.info.id, + type: 'text', + text: 'safe freeform text', + metadata: { path: '/workspace/private/unreviewed.ts' }, + }, + ], + }; + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + history: { messages: history, nextCursor: null, omittedItemCount: 0 }, + }); + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(502); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + }); + }); + + it('preserves non-enumerating not-found behavior for a cold transcript miss', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue(null); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(404); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_SESSION_NOT_FOUND' }); + }); + + it('returns snapshot-pending for a listed root without persisted transcript state', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_pending', + history: null, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_pending', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(503); + expect(response.headers.get('retry-after')).toBe('1'); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_SESSION_SNAPSHOT_PENDING', + }); + }); + + it('maps an oversized cold transcript materialization result to stable 413', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + history: { kind: 'too_large', maximumBytes: 8 * 1024 * 1024, phase: 'message_scan' }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message?limit=1`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(413); + await expect(response.json()).resolves.toEqual({ + error: 'KILO_TRANSCRIPT_TOO_LARGE', + message: + 'Persisted Kilo transcript exceeds the safe cold-read budget; use smaller bounded history when possible, or retry while a live runtime is available', + }); + }); + + it('maps retryable cold transcript failure to stable retryable 503', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + history: { kind: 'retryable_failure', phase: 'page_parts' }, + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(503); + expect(response.headers.get('retry-after')).toBe('1'); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_SESSION_READ_RETRYABLE' }); + }); + + it('maps invalid persisted transcript data to the warm-path upstream invalid response', async () => { + const env = envStub(); + vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ + kiloSessionId, + cloudAgentSessionId: 'agent_cold', + history: { kind: 'invalid_data' }, + } as never); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper: vi.fn(async () => null), + }, + }); + + expect(response.status).toBe(502); + await expect(response.json()).resolves.toEqual({ + error: 'KILO_UPSTREAM_RESPONSE_INVALID', + message: 'Kilo messages response is not valid', + }); + }); + + it('accepts the exact public directory selector on prompt_async while retaining prompt admission', async () => { + const admitPrompt = vi.fn().mockResolvedValue({ + success: true, + outcome: 'queued', + messageId: 'msg_match_directory', + compatibilityDelivery: 'queued', + }); + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/prompt_async?directory=${encodeURIComponent(publicCloudAgentDirectory(kiloSessionId))}`, + { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ parts: [{ type: 'text', text: 'matching selector' }] }), + } + ), + env: envStub(), + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance: vi.fn(async () => ({ success: true as const })), + admitPrompt, + }, + }); + + expect(response.status).toBe(204); + expect(admitPrompt).toHaveBeenCalledOnce(); + }); + + it('rejects workspace or mismatching prompt selectors before balance or admission side effects', async () => { + const validatePromptBalance = vi.fn(); + const admitPrompt = vi.fn(); + for (const search of [ + '?workspace=private', + `?directory=${encodeURIComponent(publicCloudAgentDirectory('ses_22222222222222222222222222'))}`, + ]) { + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/prompt_async${search}`, + { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ parts: [{ type: 'text', text: 'blocked' }] }), + } + ), + env: envStub(), + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance, + admitPrompt, + }, + }); + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_SESSION_SELECTOR_UNSUPPORTED', + }); + } + expect(validatePromptBalance).not.toHaveBeenCalled(); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + + it('rejects abort selectors before interruption side effects', async () => { + const interruptPrompt = vi.fn(); + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/abort?workspace=private`, + { + method: 'POST', + } + ), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + interruptPrompt, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_SESSION_SELECTOR_UNSUPPORTED', + }); + expect(interruptPrompt).not.toHaveBeenCalled(); + }); + + it('admits text-only prompt_async through the session DO while its wrapper is cold', async () => { + const env = envStub(); + const admitSubmittedMessage = vi.fn().mockResolvedValue({ + success: true, + outcome: 'queued', + messageId: 'msg_018f1e2d3c4bAsynPrmtAbCdEf', + compatibilityDelivery: 'queued', + }); + vi.spyOn(env.CLOUD_AGENT_SESSION, 'get').mockReturnValue({ + hasMessageAdmission: vi.fn().mockResolvedValue(false), + admitSubmittedMessage, + } as never); + const idFromName = vi.spyOn(env.CLOUD_AGENT_SESSION, 'idFromName'); + const resolveLiveWrapper = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + messageID: 'msg_018f1e2d3c4bAsynPrmtAbCdEf', + parts: [ + { type: 'text', text: 'hello' }, + { type: 'text', text: ' world' }, + ], + }), + }), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + resolveLiveWrapper, + validatePromptBalance: vi.fn(async () => ({ success: true as const })), + }, + authToken: 'validated-token', + }); + + expect(response.status).toBe(204); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + expect(idFromName).toHaveBeenCalledWith('usr_1:agent_cold'); + expect(admitSubmittedMessage).toHaveBeenCalledWith({ + userId: 'usr_1', + turn: { + type: 'prompt', + id: 'msg_018f1e2d3c4bAsynPrmtAbCdEf', + prompt: 'hello world', + }, + }); + }); + + it('preflights the stored session model before admitting a new prompt_async message', async () => { + const env = envStub(); + const admitSubmittedMessage = vi.fn().mockResolvedValue({ + success: true, + outcome: 'queued', + messageId: 'msg_018f1e2d3c4bPreflightAbCdEf', + compatibilityDelivery: 'queued', + }); + vi.spyOn(env.CLOUD_AGENT_SESSION, 'get').mockReturnValue({ + admitSubmittedMessage, + } as never); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ parts: [{ type: 'text', text: 'preflight me' }] }), + }), + env, + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance: vi.fn(async () => ({ success: true as const })), + }, + }); + + expect(response.status).toBe(204); + expect(preflightExistingPromptModelMock).toHaveBeenCalledWith({ + env, + userId: 'usr_1', + cloudAgentSessionId: 'agent_cold', + procedure: 'kilo.prompt_async', + }); + expect(preflightExistingPromptModelMock.mock.invocationCallOrder[0]).toBeLessThan( + admitSubmittedMessage.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY + ); + }); + + it('replays an admitted matching prompt_async message before model preflight', async () => { + const env = envStub(); + const admitSubmittedMessage = vi.fn().mockResolvedValue({ + success: true, + outcome: 'queued', + messageId: 'msg_018f1e2d3c4bReplayAbCdEfGh', + compatibilityDelivery: 'queued', + }); + const hasMessageAdmission = vi.fn().mockResolvedValue(true); + vi.spyOn(env.CLOUD_AGENT_SESSION, 'get').mockReturnValue({ + admitSubmittedMessage, + hasMessageAdmission, + } as never); + preflightExistingPromptModelMock.mockRejectedValue( + new Error('Selected model is no longer available') + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + messageID: 'msg_018f1e2d3c4bReplayAbCdEfGh', + parts: [{ type: 'text', text: 'retry me' }], + }), + }), + env, + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance: vi.fn(async () => ({ success: true as const })), + }, + }); + + expect(response.status).toBe(204); + expect(hasMessageAdmission).toHaveBeenCalledWith('msg_018f1e2d3c4bReplayAbCdEfGh'); + expect(preflightExistingPromptModelMock).not.toHaveBeenCalled(); + expect(admitSubmittedMessage).toHaveBeenCalledOnce(); + }); + + it('maps unavailable stored-session models to an SDK error before prompt_async admission', async () => { + const env = envStub(); + const admitPrompt = vi.fn(); + preflightExistingPromptModelMock.mockRejectedValue( + new TRPCError({ + code: 'BAD_REQUEST', + message: 'Selected model is not available for this cloud agent session', + }) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ parts: [{ type: 'text', text: 'reject me' }] }), + }), + env, + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance: vi.fn(async () => ({ success: true as const })), + admitPrompt, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual({ + error: 'KILO_PROMPT_ADMISSION_REJECTED', + message: 'Selected model is not available for this cloud agent session', + }); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + + it('maps unavailable stored-session model validation to retryable SDK error before admission', async () => { + const env = envStub(); + const admitPrompt = vi.fn(); + preflightExistingPromptModelMock.mockRejectedValue( + new TRPCError({ + code: 'SERVICE_UNAVAILABLE', + message: 'Model availability could not be verified', + }) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ parts: [{ type: 'text', text: 'retry later' }] }), + }), + env, + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance: vi.fn(async () => ({ success: true as const })), + admitPrompt, + }, + }); + + expect(response.status).toBe(503); + await expect(response.json()).resolves.toEqual({ + error: 'MODEL_VALIDATION_UNAVAILABLE', + message: 'Model availability could not be verified', + }); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + + it('rejects unsupported prompt_async semantic fields without admitting work', async () => { + const env = envStub(); + const admitSubmittedMessage = vi.fn(); + vi.spyOn(env.CLOUD_AGENT_SESSION, 'get').mockReturnValue({ + admitSubmittedMessage, + } as never); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + parts: [{ type: 'text', text: 'hello' }], + noReply: true, + }), + }), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_BASIC_PROMPT_UNSUPPORTED', + }); + expect(admitSubmittedMessage).not.toHaveBeenCalled(); + }); + + it('rejects prompt_async attempts to bypass public balance validation after owner resolution', async () => { + const env = envStub(); + const validatePromptBalance = vi.fn(); + const admitPrompt = vi.fn(); + const resolveRootSessionForKiloSession = vi.fn(async () => ({ + cloudAgentSessionId: 'agent_owned', + })); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-skip-balance-check': 'true', + }, + body: JSON.stringify({ parts: [{ type: 'text', text: 'local acceptance' }] }), + }), + env, + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession, + validatePromptBalance, + admitPrompt, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual({ + error: 'KILO_BALANCE_BYPASS_UNSUPPORTED', + message: 'Balance bypass is not supported for public Kilo prompt mutations', + }); + expect(resolveRootSessionForKiloSession).toHaveBeenCalledWith({ + env, + userId: 'usr_1', + kiloSessionId, + }); + expect(validatePromptBalance).not.toHaveBeenCalled(); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + + it('does not support synchronous prompt mutation or admit its body', async () => { + const validatePromptBalance = vi.fn(); + const admitPrompt = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + messageID: 'msg_018f1e2d3c4bSyncGoneAbCdEf', + parts: [{ type: 'text', text: 'do not admit this' }], + }), + }), + env: envStub(), + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_owned', + })), + validatePromptBalance, + admitPrompt, + }, + }); + + expect(response.status).toBe(501); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_ROUTE_UNSUPPORTED' }); + expect(validatePromptBalance).not.toHaveBeenCalled(); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + + it('applies public balance validation before durable prompt admission', async () => { + const env = envStub(); + const admitPrompt = vi.fn(); + const validatePromptBalance = vi.fn().mockResolvedValue({ + success: false, + status: 402, + message: 'Insufficient credits: $1 minimum required', + }); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ parts: [{ type: 'text', text: 'hello' }] }), + }), + env, + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance, + admitPrompt, + }, + }); + + expect(response.status).toBe(402); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_BALANCE_VALIDATION_FAILED', + }); + expect(validatePromptBalance).toHaveBeenCalledWith({ + env, + authToken: 'validated-token', + userId: 'usr_1', + cloudAgentSessionId: 'agent_cold', + }); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + + it('applies abort idempotently and returns the native boolean success response', async () => { + const interruptPrompt = vi.fn().mockResolvedValue({ + success: false, + message: 'No accepted wrapper messages or pending queued messages', + }); + const resolveLiveWrapper = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/abort`, { + method: 'POST', + }), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + interruptPrompt, + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toBe(true); + expect(interruptPrompt).toHaveBeenCalledWith({ + env: expect.anything(), + userId: 'usr_1', + cloudAgentSessionId: 'agent_cold', + }); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + }); + + it('rejects unsupported root-scoped session routes without contacting a warm wrapper', async () => { + const containerFetch = vi.fn(async () => new Response('unexpected')); + const resolveLiveWrapper = vi.fn(async () => liveWrapperTarget(containerFetch)); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/share?cursor=abc`, { + method: 'POST', + body: JSON.stringify({ prompt: 'hello' }), + }), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(501); + await expect(response.json()).resolves.toEqual({ + error: 'KILO_ROUTE_UNSUPPORTED', + message: 'Kilo facade route is not supported', + }); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + expect(containerFetch).not.toHaveBeenCalled(); + }); + + it('preserves the exact no-suffix session route shape', async () => { + const containerFetch = vi.fn(async () => new Response('ok')); + + await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + const [forwardedRequest] = containerFetch.mock.calls[0]; + expect(new URL((forwardedRequest as Request).url).pathname).toBe( + `/kilo-proxy/session/${kiloSessionId}` + ); + }); +}); + +describe('UserKiloFacade producer messages', () => { + it('does not let an older routed producer replace a newer connected feed', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const close = vi.fn(); + const currentSocket = { + close, + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + const ctx = { + getWebSockets: vi.fn(() => [currentSocket]), + } as unknown as DurableObjectState; + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade(ctx, env); + const response = await facade.fetch( + new Request( + 'http://worker.test/internal/kilo/global-feed?userId=usr_1&cloudAgentSessionId=agent_live&kiloSessionId=ses_12345678901234567890123456&wrapperRunId=wr_old&wrapperGeneration=1&wrapperConnectionId=conn_old', + { headers: { Upgrade: 'websocket' } } + ) + ); + + expect(response.status).toBe(409); + expect(close).not.toHaveBeenCalled(); + }); + + it('emits a native-shaped global connection event in the public virtual-server directory', async () => { + const env = { + CLOUD_AGENT_SESSION: { idFromName: vi.fn(), get: vi.fn() }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicGlobalEventStream(); + const body = response.body; + if (!body) throw new Error('Expected public global event response body'); + const reader = body.getReader(); + + try { + const event = await readSseFrame<{ + directory?: string; + payload: { id?: string; type: string }; + }>(reader); + expect(event.directory).toBe('/cloud-agent'); + expect(event.payload.type).toBe('server.connected'); + expect(event.payload.id).toMatch(/^evt_/); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('drops malformed substantive global producer envelopes without a usable directory or payload type', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicGlobalEventStream(); + const body = response.body; + if (!body) throw new Error('Expected public global event response body'); + const reader = body.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + + try { + await readSseFrame(reader); + await facade.webSocketMessage( + ws, + JSON.stringify({ + payload: { type: 'message.updated', properties: { sessionID: kiloSessionId } }, + }) + ); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/root', + payload: { properties: { sessionID: kiloSessionId } }, + }) + ); + expect((ws as unknown as { send: ReturnType }).send).toHaveBeenCalledTimes(2); + expect((ws as unknown as { send: ReturnType }).send).toHaveBeenLastCalledWith( + JSON.stringify({ error: 'Invalid global event envelope' }) + ); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('accepts the deployed native substantive message envelope without a payload id', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicGlobalEventStream(); + const body = response.body; + if (!body) throw new Error('Expected public global event response body'); + const reader = body.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + const payload = { + type: 'message.updated', + properties: { sessionID: kiloSessionId, info: { id: 'msg_native' } }, + }; + + try { + await readSseFrame(reader); + await facade.webSocketMessage( + ws, + JSON.stringify({ directory: '/workspace/native', payload }) + ); + await expect(readSseFrame(reader)).resolves.toEqual({ + directory: publicCloudAgentDirectory(kiloSessionId), + payload, + }); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('rewrites a valid substantive global producer event at the public stream boundary', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicGlobalEventStream(); + const body = response.body; + if (!body) throw new Error('Expected public global event response body'); + const reader = body.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + + try { + await readSseFrame(reader); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + project: 'project-1', + payload: { + id: 'evt_public', + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'msg_public' }, + }, + }) + ); + await expect(readSseFrame(reader)).resolves.toEqual({ + directory: publicCloudAgentDirectory(kiloSessionId), + project: 'project-1', + payload: { + id: 'evt_public', + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'msg_public' }, + }, + }); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('suppresses rewritten global envelopes with unsafe sibling paths before safe sentinel fanout', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicGlobalEventStream(); + const body = response.body; + if (!body) throw new Error('Expected public global event response body'); + const reader = body.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + + try { + await readSseFrame(reader); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + artifact: { path: '/workspace/private/envelope-secret.ts' }, + payload: { + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'unsafe-envelope' }, + }, + }) + ); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'safe-after-envelope-suppressed' }, + }, + }) + ); + const next = await readSseFrame<{ payload: { properties: { id: string } } }>(reader); + expect(next.payload.properties.id).toBe('safe-after-envelope-suppressed'); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('uses comment frames for public keepalives rather than undeclared heartbeat events', async () => { + vi.useFakeTimers(); + const env = { + CLOUD_AGENT_SESSION: { idFromName: vi.fn(), get: vi.fn() }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const body = facade.openPublicSessionEventStream(kiloSessionId).body; + if (!body) throw new Error('Expected public scoped event response body'); + const reader = body.getReader(); + try { + expect((await readSseFrame<{ type: string }>(reader)).type).toBe('server.connected'); + await vi.advanceTimersByTimeAsync(10_000); + const keepalive = await readSseTextFrame(reader); + expect(keepalive).toBe(': heartbeat\n\n'); + expect(keepalive).not.toContain('server.heartbeat'); + } finally { + await reader.cancel().catch(() => undefined); + vi.useRealTimers(); + } + }); + + it('emits a native-shaped scoped connection event while no producer is live', async () => { + const env = { + CLOUD_AGENT_SESSION: { idFromName: vi.fn(), get: vi.fn() }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicSessionEventStream(kiloSessionId); + const body = response.body; + if (!body) throw new Error('Expected public scoped event response body'); + const reader = body.getReader(); + + try { + const event = await readSseFrame<{ + id?: string; + type: string; + properties: Record; + }>(reader); + expect(event).toEqual({ + id: expect.stringMatching(/^evt_/), + type: 'server.connected', + properties: {}, + }); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('fans Cloud Agent extension events out before a native wrapper producer is live', async () => { + const env = { + CLOUD_AGENT_SESSION: { idFromName: vi.fn(), get: vi.fn() }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const globalBody = facade.openPublicGlobalEventStream().body; + const scopedBody = facade.openPublicSessionEventStream(kiloSessionId).body; + if (!globalBody || !scopedBody) throw new Error('Expected public event response bodies'); + const globalReader = globalBody.getReader(); + const scopedReader = scopedBody.getReader(); + const event = { + type: 'cloud.message.queued' as const, + properties: { + sessionID: kiloSessionId, + messageId: 'msg_extension', + delivery: 'queued' as const, + }, + }; + + try { + await readSseFrame(globalReader); + await readSseFrame(scopedReader); + await facade.publishCloudAgentExtensionEvent({ + kiloUserId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + event, + }); + await expect(readSseFrame(scopedReader)).resolves.toEqual(event); + await expect(readSseFrame(globalReader)).resolves.toEqual({ + directory: publicCloudAgentDirectory(kiloSessionId), + payload: event, + }); + } finally { + await globalReader.cancel().catch(() => undefined); + await scopedReader.cancel().catch(() => undefined); + } + }); + + it('suppresses organization lifecycle extensions after current root access is lost', async () => { + const resolveCloudAgentRootSessionForKiloSession = vi.fn(async () => null); + const env = { + SESSION_INGEST: { resolveCloudAgentRootSessionForKiloSession }, + CLOUD_AGENT_SESSION: { idFromName: vi.fn(), get: vi.fn() }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const body = facade.openPublicGlobalEventStream().body; + if (!body) throw new Error('Expected public global event response body'); + const reader = body.getReader(); + const deniedEvent = { + type: 'cloud.message.queued' as const, + properties: { + sessionID: kiloSessionId, + messageId: 'msg_denied', + delivery: 'queued' as const, + }, + }; + const personalEvent = { + type: 'cloud.message.sent' as const, + properties: { + sessionID: kiloSessionId, + messageId: 'msg_personal', + delivery: 'sent' as const, + }, + }; + + try { + await readSseFrame(reader); + await facade.publishCloudAgentExtensionEvent({ + kiloUserId: 'usr_1', + cloudAgentSessionId: 'agent_org', + kiloSessionId, + organizationId: 'org_removed', + event: deniedEvent, + }); + await facade.publishCloudAgentExtensionEvent({ + kiloUserId: 'usr_1', + cloudAgentSessionId: 'agent_personal', + kiloSessionId, + event: personalEvent, + }); + await expect(readSseFrame(reader)).resolves.toEqual({ + directory: publicCloudAgentDirectory(kiloSessionId), + payload: personalEvent, + }); + expect(resolveCloudAgentRootSessionForKiloSession).toHaveBeenCalledWith({ + kiloUserId: 'usr_1', + kiloSessionId, + }); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('opens a scoped stream while cold and emits only matching bare native events', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicSessionEventStream(kiloSessionId); + const body = response.body; + if (!body) throw new Error('Expected public scoped event response body'); + const reader = body.getReader(); + const matchingWs = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + const otherWs = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_other', + kiloSessionId: 'ses_22222222222222222222222222', + wrapperRunId: 'wr_other', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_other', + }), + } as unknown as WebSocket; + + try { + const connected = await readSseFrame<{ id?: string; type: string }>(reader); + expect(connected.type).toBe('server.connected'); + expect(connected.id).toMatch(/^evt_/); + await facade.webSocketMessage( + otherWs, + JSON.stringify({ + directory: '/workspace/other', + payload: { + id: 'evt_other', + type: 'message.updated', + properties: { sessionID: 'ses_22222222222222222222222222', id: 'msg_other' }, + }, + }) + ); + await facade.webSocketMessage( + matchingWs, + JSON.stringify({ + payload: { id: 'evt_synthetic', type: 'server.connected', properties: {} }, + }) + ); + await facade.webSocketMessage( + matchingWs, + JSON.stringify({ + directory: '/workspace/root', + payload: { + type: 'message.updated', + properties: { sessionID: kiloSessionId, info: { id: 'msg_match' } }, + }, + }) + ); + const event = await readSseFrame<{ + type: string; + properties: { sessionID: string; info: { id: string } }; + }>(reader); + expect(event).toEqual({ + type: 'message.updated', + properties: { sessionID: kiloSessionId, info: { id: 'msg_match' } }, + }); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('projects embedded session and message entities before scoped and global fanout', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const globalBody = facade.openPublicGlobalEventStream().body; + const scopedBody = facade.openPublicSessionEventStream(kiloSessionId).body; + if (!globalBody || !scopedBody) throw new Error('Expected public event response bodies'); + const globalReader = globalBody.getReader(); + const scopedReader = scopedBody.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + + try { + await readSseFrame(globalReader); + await readSseFrame(scopedReader); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'session.updated', + properties: { sessionID: kiloSessionId, info: sdkSessionInfo(kiloSessionId) }, + }, + }) + ); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'message.updated', + properties: { sessionID: kiloSessionId, info: sdkMessageHistory()[1].info }, + }, + }) + ); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'message.updated', + properties: { sessionID: kiloSessionId, info: sdkMessageHistory()[0].info }, + }, + }) + ); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'message.part.updated', + properties: { sessionID: kiloSessionId, part: sdkMessageHistory()[1].parts[1] }, + }, + }) + ); + + const globalSession = await readSseFrame<{ payload: { properties: { info: unknown } } }>( + globalReader + ); + const scopedSession = await readSseFrame<{ properties: { info: unknown } }>(scopedReader); + expect(globalSession.payload.properties.info).toEqual(projectedSdkSessionInfo(kiloSessionId)); + expect(scopedSession.properties.info).toEqual(projectedSdkSessionInfo(kiloSessionId)); + const globalMessage = await readSseFrame<{ + payload: { properties: { info: { path: unknown } } }; + }>(globalReader); + const scopedMessage = await readSseFrame<{ properties: { info: { path: unknown } } }>( + scopedReader + ); + const directory = publicCloudAgentDirectory(kiloSessionId); + expect(globalMessage.payload.properties.info.path).toEqual({ + cwd: directory, + root: directory, + }); + expect(scopedMessage.properties.info.path).toEqual({ cwd: directory, root: directory }); + const globalUserMessage = await readSseFrame<{ + payload: { properties: { info: { editorContext: { activeFile: string } } } }; + }>(globalReader); + const scopedUserMessage = await readSseFrame<{ + properties: { info: { editorContext: { activeFile: string } } }; + }>(scopedReader); + expect(globalUserMessage.payload.properties.info.editorContext.activeFile).toBe(directory); + expect(scopedUserMessage.properties.info.editorContext.activeFile).toBe(directory); + const globalPart = await readSseFrame<{ + payload: { + properties: { part: { state: { attachments: Array<{ source: { path: string } }> } } }; + }; + }>(globalReader); + const scopedPart = await readSseFrame<{ + properties: { part: { state: { attachments: Array<{ source: { path: string } }> } } }; + }>(scopedReader); + expect(globalPart.payload.properties.part.state.attachments[0].source.path).toBe(directory); + expect(scopedPart.properties.part.state.attachments[0].source.path).toBe(directory); + } finally { + await globalReader.cancel().catch(() => undefined); + await scopedReader.cancel().catch(() => undefined); + } + }); + + it('does not accept Cloud Agent extension namespace frames from wrapper producers', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const body = facade.openPublicGlobalEventStream().body; + if (!body) throw new Error('Expected public global event response body'); + const reader = body.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + + try { + await readSseFrame(reader); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'cloud.status', + properties: { internalDetail: 'secret backend lifecycle detail' }, + }, + }) + ); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'safe-native-after-extension' }, + }, + }) + ); + const next = await readSseFrame<{ payload: { properties: { id: string } } }>(reader); + expect(next.payload.properties.id).toBe('safe-native-after-extension'); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('suppresses identityless and unreviewed wrapper native variants rather than leaking them', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicGlobalEventStream(); + const body = response.body; + if (!body) throw new Error('Expected public global event response body'); + const reader = body.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + + try { + await readSseFrame(reader); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { type: 'file.edited', properties: { path: '/workspace/private/secret.ts' } }, + }) + ); + for (const type of ['pty.exited', 'indexing.status', 'session.error']) { + await facade.webSocketMessage( + ws, + JSON.stringify({ directory: '/workspace/private', payload: { type, properties: {} } }) + ); + } + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'session.idle', + properties: { sessionID: 'ses_22222222222222222222222222' }, + }, + }) + ); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'message.updated', + properties: { + sessionID: kiloSessionId, + info: sdkMessageHistory()[1].info, + artifact: { path: '/workspace/private/unreviewed.ts' }, + }, + }, + }) + ); + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/private', + payload: { + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'safe-after-suppressed' }, + }, + }) + ); + const next = await readSseFrame<{ payload: { properties: { id: string } } }>(reader); + expect(next.payload.properties.id).toBe('safe-after-suppressed'); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('closes producer messages once their wrapper fence is no longer current', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ + success: false as const, + status: 409, + message: 'Stale wrapper connection', + })); + const idFromName = vi.fn(() => 'session-do-id'); + const env = { + CLOUD_AGENT_SESSION: { + idFromName, + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const close = vi.fn(); + const ws = { + close, + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_stale', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_stale', + }), + } as unknown as WebSocket; + + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/root', + payload: { + id: 'evt_msg_1', + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'msg_1' }, + }, + }) + ); + + expect(idFromName).toHaveBeenCalledWith('usr_1:agent_live'); + expect(validateKiloGlobalFeedProducer).toHaveBeenCalledWith({ + kiloSessionId, + wrapperRunId: 'wr_stale', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_stale', + }); + expect(close).toHaveBeenCalledWith(4401, 'Stale wrapper connection'); + }); + + it('terminates a public subscriber that falls behind the live event stream', async () => { + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicGlobalEventStream(); + const body = response.body; + if (!body) { + throw new Error('Expected public global event response body'); + } + const reader = body.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + + for (let index = 0; index < 256; index++) { + await facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/root', + payload: { + id: `evt_${index}`, + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: `msg_${index}` }, + }, + }) + ); + } + + try { + await expect(reader.read()).rejects.toThrow('Global event subscriber fell behind'); + } finally { + await reader.cancel().catch(() => undefined); + } + }); + + it('preserves producer event order while fence validation is asynchronous', async () => { + let resolveFirstValidation: ((result: { success: true }) => void) | undefined; + const validateKiloGlobalFeedProducer = vi + .fn() + .mockImplementationOnce( + () => + new Promise<{ success: true }>(resolve => { + resolveFirstValidation = resolve; + }) + ) + .mockResolvedValue({ success: true as const }); + const env = { + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(() => 'session-do-id'), + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + const response = facade.openPublicGlobalEventStream(); + const body = response.body; + if (!body) { + throw new Error('Expected public global event response body'); + } + const reader = body.getReader(); + const ws = { + close: vi.fn(), + send: vi.fn(), + deserializeAttachment: () => ({ + userId: 'usr_1', + cloudAgentSessionId: 'agent_live', + kiloSessionId, + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }), + } as unknown as WebSocket; + + const first = facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/root', + payload: { + id: 'evt_first', + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'first' }, + }, + }) + ); + const second = facade.webSocketMessage( + ws, + JSON.stringify({ + directory: '/workspace/root', + payload: { + id: 'evt_second', + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'second' }, + }, + }) + ); + await vi.waitFor(() => { + expect(validateKiloGlobalFeedProducer).toHaveBeenCalledTimes(1); + }); + if (!resolveFirstValidation) { + throw new Error('Expected the first producer validation to be pending'); + } + resolveFirstValidation({ success: true }); + await Promise.all([first, second]); + + const readEvent = async (): Promise<{ payload: { properties?: { id?: string } } }> => { + const frame = await reader.read(); + if (frame.done) { + throw new Error('Expected global event frame'); + } + return JSON.parse(new TextDecoder().decode(frame.value).slice('data: '.length)); + }; + + try { + await readEvent(); + expect((await readEvent()).payload.properties?.id).toBe('first'); + expect((await readEvent()).payload.properties?.id).toBe('second'); + } finally { + await reader.cancel().catch(() => undefined); + } + }); +}); + +describe('global event helpers', () => { + it('rewrites only directory to the public Cloud Agent session namespace', () => { + const payload = { + id: 'evt_msg_1', + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'msg_1' }, + }; + const event = { + directory: '/workspace/project/sessions/agent_live', + project: 'project-1', + workspace: 'workspace-1', + payload, + }; + + expect(rewriteGlobalEventDirectory(event, kiloSessionId)).toEqual({ + directory: publicCloudAgentDirectory(kiloSessionId), + project: 'project-1', + workspace: 'workspace-1', + payload, + }); + }); + + it('identifies wrapper synthetic global frames for suppression', () => { + expect( + isSyntheticGlobalEvent({ + payload: { id: 'evt_connected', type: 'server.connected', properties: {} }, + }) + ).toBe(true); + expect( + isSyntheticGlobalEvent({ + payload: { id: 'evt_heartbeat', type: 'server.heartbeat', properties: {} }, + }) + ).toBe(true); + expect( + isSyntheticGlobalEvent({ + payload: { id: 'evt_message', type: 'message.updated', properties: {} }, + }) + ).toBe(false); + }); +}); diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts new file mode 100644 index 0000000000..bd6c5d4e68 --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -0,0 +1,1634 @@ +import { + MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE, + validateKiloSdkMessagesCursor, +} from '@kilocode/session-ingest-contracts'; +import { DurableObject } from 'cloudflare:workers'; +import { TRPCError } from '@trpc/server'; +import type { + GetCloudAgentRootSessionMessagesParams, + KiloSdkSessionInfo, + KiloSdkStoredMessage, + ListCloudAgentRootSessionsParams, +} from '../session-ingest-binding.js'; +import { + fetchOrgIdForSession, + validateBalanceOnly, + type BalanceOnlyResult, +} from '../balance-validation.js'; +import type { + SubmittedSessionMessageRequest, + SessionMessageAdmissionResult, +} from '../execution/types.js'; +import type { CloudAgentSession } from '../persistence/CloudAgentSession.js'; +import type { UserId } from '../types/ids.js'; +import type { Env } from '../types.js'; +import { withDORetry } from '../utils/do-retry.js'; +import { preflightAndAdmitPromptMessage } from '../session/queue-message.js'; +import { parseBasicKiloPrompt } from './basic-prompt.js'; +import { + isPublicCloudAgentExtensionSourceType, + type PublicCloudAgentExtensionEvent, +} from './cloud-agent-extension-events.js'; +import { createProxyRequest } from '../shared/http-proxy.js'; +import { + hasUnprojectedPrivateStructuredPath, + projectPublicEvent, + projectPublicListedSession, + projectPublicSession, + projectPublicStoredMessages, + publicCloudAgentDirectory, +} from './public-sdk-projection.js'; +import { + buildWrapperKiloProxyUrl, + decideSessionKiloFacadeRoute, + resolveLiveWrapperTarget, + type LiveWrapperTarget, + type SessionKiloFacadeDecision, + type SessionKiloFacadePolicyInput, +} from './session-proxy.js'; + +export const KILO_FACADE_USER_ID_HEADER = 'x-kilo-facade-user-id'; +export const KILO_FACADE_AUTH_TOKEN_HEADER = 'x-kilo-facade-auth-token'; + +export const KILO_FACADE_GLOBAL_FEED_PATH = '/internal/kilo/global-feed'; +const HEARTBEAT_INTERVAL_MS = 10_000; +const MAX_PUBLIC_GLOBAL_EVENT_QUEUE_SIZE = 64; +const SLOW_SUBSCRIBER_ERROR = 'Global event subscriber fell behind'; +const SUPPORTED_SESSION_LIST_QUERY_PARAMS = new Set(['limit', 'start']); +const SUPPORTED_SESSION_MESSAGES_QUERY_PARAMS = new Set(['directory', 'limit', 'before']); +const KNOWN_UNSUPPORTED_ROUTES = new Set([ + 'GET /session/status', + 'POST /session/viewed', + 'POST /sync/history', + 'GET /config', + 'GET /provider', + 'GET /project/current', + 'GET /global/health', +]); +const MAX_KILO_SESSION_JSON_BYTES = 8 * 1024 * 1024; +const MAX_KILO_ERROR_JSON_BYTES = 64 * 1024; +const MAX_KILO_PROMPT_JSON_BYTES = 256 * 1024; +const MAX_TIMESTAMP_MILLISECONDS = 8_640_000_000_000_000; +const PUBLIC_VIRTUAL_SERVER_DIRECTORY = '/cloud-agent'; + +type KiloEventPayload = Record & { + id?: string; + type: string; + properties: Record; +}; + +type KiloGlobalEventEnvelope = Record & { + directory?: string; + project?: unknown; + workspace?: unknown; + payload?: KiloEventPayload; +}; + +type GlobalFeedSource = { + userId: string; + cloudAgentSessionId: string; + kiloSessionId: string; + wrapperRunId: string; + wrapperGeneration: number; + wrapperConnectionId: string; +}; + +type PublicSubscriber = { + controller: ReadableStreamDefaultController; + heartbeat: ReturnType; + scope: { kind: 'global' } | { kind: 'session'; kiloSessionId: string }; +}; + +export type KiloFacadeGlobalEvents = { + openPublicGlobalEventStream(): Response; + openPublicSessionEventStream(kiloSessionId: string): Response; +}; + +export type KiloFacadeRequestDeps = { + resolveRootSessionForKiloSession?: (params: { + env: Env; + userId: string; + kiloSessionId: string; + }) => Promise<{ cloudAgentSessionId: string } | null>; + decideSessionRoute?: (input: SessionKiloFacadePolicyInput) => SessionKiloFacadeDecision; + resolveLiveWrapper?: (params: { + env: Env; + userId: string; + cloudAgentSessionId: string; + }) => Promise; + admitPrompt?: (params: { + env: Env; + userId: string; + cloudAgentSessionId: string; + request: SubmittedSessionMessageRequest; + }) => Promise; + validatePromptBalance?: (params: { + env: Env; + authToken: string; + userId: string; + cloudAgentSessionId: string; + }) => Promise; + interruptPrompt?: (params: { + env: Env; + userId: string; + cloudAgentSessionId: string; + }) => Promise>>; + globalEvents?: KiloFacadeGlobalEvents; +}; + +const encoder = new TextEncoder(); + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function facadeError(status: number, code: string, message: string): Response { + return Response.json({ error: code, message }, { status }); +} + +function kiloRelativePath(pathname: string): string { + if (pathname === '/kilo') { + return '/'; + } + if (pathname.startsWith('/kilo/')) { + return pathname.slice('/kilo'.length); + } + return pathname; +} + +function parseRootSessionRoute(kiloPath: string): { + encodedKiloSessionId: string; + kiloSessionId: string; +} | null { + const match = /^\/session\/([^/]+)(?:\/.*)?$/.exec(kiloPath); + if (!match?.[1]) { + return null; + } + try { + return { + encodedKiloSessionId: match[1], + kiloSessionId: decodeURIComponent(match[1]), + }; + } catch { + return null; + } +} + +function isSessionIngestKiloSessionId(kiloSessionId: string): boolean { + return kiloSessionId.startsWith('ses_') && kiloSessionId.length === 30; +} + +function missingRootKiloSessionResponse(): Response { + return facadeError(404, 'KILO_SESSION_NOT_FOUND', 'Cloud Agent root Kilo session was not found'); +} + +function pendingSessionSnapshotResponse(): Response { + const response = facadeError( + 503, + 'KILO_SESSION_SNAPSHOT_PENDING', + 'Cloud Agent Kilo session snapshot is not available yet' + ); + response.headers.set('Retry-After', '1'); + return response; +} + +function unsupportedStructuredPathsResponse(entity: 'session' | 'messages'): Response { + return facadeError( + 502, + 'KILO_UPSTREAM_RESPONSE_INVALID', + `Kilo ${entity} response contains unsupported structured paths` + ); +} + +function projectSafePublicSession( + info: KiloSdkSessionInfo, + kiloSessionId: string +): KiloSdkSessionInfo | Response { + const publicInfo = projectPublicSession(info, kiloSessionId); + return hasUnprojectedPrivateStructuredPath(publicInfo) + ? unsupportedStructuredPathsResponse('session') + : publicInfo; +} + +function projectSafePublicMessages( + messages: KiloSdkStoredMessage[], + kiloSessionId: string +): KiloSdkStoredMessage[] | Response { + const publicMessages = projectPublicStoredMessages(messages, kiloSessionId); + return hasUnprojectedPrivateStructuredPath(publicMessages) + ? unsupportedStructuredPathsResponse('messages') + : publicMessages; +} + +function sessionSnapshotTooLargeResponse(): Response { + return facadeError( + 413, + 'KILO_SESSION_SNAPSHOT_TOO_LARGE', + 'Persisted Kilo session snapshot exceeds the safe cold-read budget' + ); +} + +function retryableSessionReadResponse(): Response { + const response = facadeError( + 503, + 'KILO_SESSION_READ_RETRYABLE', + 'Persisted Kilo session data is temporarily unavailable; retry the request' + ); + response.headers.set('Retry-After', '1'); + return response; +} + +function invalidPersistedSessionDataResponse(entity: 'session' | 'messages'): Response { + return facadeError(502, 'KILO_UPSTREAM_RESPONSE_INVALID', `Kilo ${entity} response is not valid`); +} + +function transcriptTooLargeResponse(): Response { + return facadeError( + 413, + 'KILO_TRANSCRIPT_TOO_LARGE', + 'Persisted Kilo transcript exceeds the safe cold-read budget; use smaller bounded history when possible, or retry while a live runtime is available' + ); +} + +function isKiloEventPayload(value: unknown): value is KiloEventPayload { + return isRecord(value) && typeof value.type === 'string' && isRecord(value.properties); +} + +function isKiloGlobalEventEnvelope(value: unknown): value is KiloGlobalEventEnvelope { + return isRecord(value) && isKiloEventPayload(value.payload); +} + +function isSubstantiveKiloGlobalEventEnvelope( + value: unknown +): value is KiloGlobalEventEnvelope & { directory: string; payload: KiloEventPayload } { + return ( + isKiloGlobalEventEnvelope(value) && + typeof value.directory === 'string' && + value.directory.length > 0 + ); +} + +function createPublicEventPayload(type: 'server.connected'): KiloEventPayload { + return { + id: `evt_${crypto.randomUUID()}`, + type, + properties: {}, + }; +} + +export { publicCloudAgentDirectory }; + +function isKiloSdkSessionInfo(value: unknown, kiloSessionId: string): value is KiloSdkSessionInfo { + return ( + isRecord(value) && + value.id === kiloSessionId && + typeof value.slug === 'string' && + typeof value.projectID === 'string' && + typeof value.directory === 'string' && + typeof value.title === 'string' && + typeof value.version === 'string' && + isRecord(value.time) && + typeof value.time.created === 'number' && + typeof value.time.updated === 'number' + ); +} + +async function readBoundedBody( + response: Response, + maximumBytes: number +): Promise { + if (!response.body) { + return new Uint8Array(); + } + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let byteLength = 0; + for (;;) { + const chunk = await reader.read(); + if (chunk.done) { + break; + } + const value: unknown = chunk.value; + if (!(value instanceof Uint8Array)) { + await reader.cancel().catch(() => undefined); + return null; + } + if (byteLength + value.byteLength > maximumBytes) { + await reader.cancel().catch(() => undefined); + return null; + } + chunks.push(value); + byteLength += value.byteLength; + } + const bytes = new Uint8Array(byteLength); + let offset = 0; + for (const chunk of chunks) { + bytes.set(chunk, offset); + offset += chunk.byteLength; + } + return bytes; +} + +async function isUnavailableKiloRuntimeResponse(response: Response): Promise { + if (response.status !== 502 && response.status !== 503) { + return false; + } + const contentType = response.headers.get('content-type'); + if (!contentType?.toLowerCase().includes('application/json')) { + return false; + } + const bytes = await readBoundedBody(response.clone(), MAX_KILO_ERROR_JSON_BYTES); + if (!bytes) { + return false; + } + try { + const parsed: unknown = JSON.parse(new TextDecoder().decode(bytes)); + return ( + isRecord(parsed) && + (parsed.error === 'KILO_RUNTIME_UNAVAILABLE' || parsed.error === 'KILO_PROXY_ERROR') + ); + } catch { + return false; + } +} + +async function parseLiveSessionDetail( + response: Response, + kiloSessionId: string +): Promise { + const declaredLength = response.headers.get('content-length'); + if (declaredLength !== null) { + const bodyBytes = Number(declaredLength); + if (!Number.isSafeInteger(bodyBytes) || bodyBytes > MAX_KILO_SESSION_JSON_BYTES) { + return facadeError( + 502, + 'KILO_UPSTREAM_RESPONSE_INVALID', + 'Kilo session response exceeds supported size' + ); + } + } + + const bytes = await readBoundedBody(response, MAX_KILO_SESSION_JSON_BYTES); + if (!bytes) { + return facadeError( + 502, + 'KILO_UPSTREAM_RESPONSE_INVALID', + 'Kilo session response exceeds supported size' + ); + } + let parsed: unknown; + try { + parsed = JSON.parse(new TextDecoder().decode(bytes)); + } catch { + return facadeError( + 502, + 'KILO_UPSTREAM_RESPONSE_INVALID', + 'Kilo session response is not valid JSON' + ); + } + if (!isKiloSdkSessionInfo(parsed, kiloSessionId)) { + return facadeError(502, 'KILO_UPSTREAM_RESPONSE_INVALID', 'Kilo session response is not valid'); + } + return parsed; +} + +async function rewriteLiveSessionDetailResponse( + response: Response, + kiloSessionId: string +): Promise { + if (!response.ok) { + return response; + } + const info = await parseLiveSessionDetail(response, kiloSessionId); + if (info instanceof Response) { + return info; + } + const publicInfo = projectSafePublicSession(info, kiloSessionId); + if (publicInfo instanceof Response) return publicInfo; + const headers = new Headers(response.headers); + headers.delete('content-length'); + headers.delete('content-encoding'); + headers.set('content-type', 'application/json'); + return new Response(JSON.stringify(publicInfo), { + status: response.status, + statusText: response.statusText, + headers, + }); +} + +function isKiloSdkMessageInfo(value: unknown, kiloSessionId: string): boolean { + if (!isRecord(value) || value.sessionID !== kiloSessionId || typeof value.id !== 'string') { + return false; + } + if (value.role === 'user') return true; + return ( + value.role === 'assistant' && + isRecord(value.path) && + typeof value.path.cwd === 'string' && + typeof value.path.root === 'string' + ); +} + +function isKiloSdkPart(value: unknown, kiloSessionId: string): boolean { + if ( + !isRecord(value) || + value.sessionID !== kiloSessionId || + typeof value.messageID !== 'string' || + typeof value.id !== 'string' || + typeof value.type !== 'string' + ) { + return false; + } + if (value.type === 'file' && value.source !== undefined) { + return ( + isRecord(value.source) && + (value.source.type === 'resource' || typeof value.source.path === 'string') + ); + } + if (value.type === 'tool' && isRecord(value.state) && Array.isArray(value.state.attachments)) { + return value.state.attachments.every(part => isKiloSdkPart(part, kiloSessionId)); + } + return true; +} + +function isKiloSdkStoredMessage( + value: unknown, + kiloSessionId: string +): value is KiloSdkStoredMessage { + return ( + isRecord(value) && + isKiloSdkMessageInfo(value.info, kiloSessionId) && + Array.isArray(value.parts) && + value.parts.every(part => isKiloSdkPart(part, kiloSessionId)) + ); +} + +async function rewriteLiveMessagesResponse( + response: Response, + kiloSessionId: string +): Promise { + if (!response.ok) return response; + const declaredLength = response.headers.get('content-length'); + if (declaredLength !== null) { + const bodyBytes = Number(declaredLength); + if (!Number.isSafeInteger(bodyBytes) || bodyBytes > MAX_KILO_SESSION_JSON_BYTES) { + return facadeError( + 502, + 'KILO_UPSTREAM_RESPONSE_INVALID', + 'Kilo messages response exceeds supported size' + ); + } + } + const bytes = await readBoundedBody(response, MAX_KILO_SESSION_JSON_BYTES); + if (!bytes) { + return facadeError( + 502, + 'KILO_UPSTREAM_RESPONSE_INVALID', + 'Kilo messages response exceeds supported size' + ); + } + let parsed: unknown; + try { + parsed = JSON.parse(new TextDecoder().decode(bytes)); + } catch { + return facadeError( + 502, + 'KILO_UPSTREAM_RESPONSE_INVALID', + 'Kilo messages response is not valid JSON' + ); + } + if ( + !Array.isArray(parsed) || + !parsed.every(message => isKiloSdkStoredMessage(message, kiloSessionId)) + ) { + return facadeError( + 502, + 'KILO_UPSTREAM_RESPONSE_INVALID', + 'Kilo messages response is not valid' + ); + } + const publicMessages = projectSafePublicMessages(parsed, kiloSessionId); + if (publicMessages instanceof Response) return publicMessages; + const headers = new Headers(response.headers); + headers.delete('content-length'); + headers.delete('content-encoding'); + headers.set('content-type', 'application/json'); + return new Response(JSON.stringify(publicMessages), { + status: response.status, + statusText: response.statusText, + headers, + }); +} + +function parseSessionListQuery( + url: URL +): Omit | Response { + for (const key of url.searchParams.keys()) { + if (!SUPPORTED_SESSION_LIST_QUERY_PARAMS.has(key)) { + return facadeError( + 400, + 'KILO_SESSION_LIST_SELECTOR_UNSUPPORTED', + `Session list query parameter is not supported: ${key}` + ); + } + } + + const params: Omit = {}; + const limitParam = url.searchParams.get('limit'); + if (limitParam !== null) { + const limit = Number(limitParam); + if (!Number.isInteger(limit) || limit < 1 || limit > 100) { + return facadeError( + 400, + 'KILO_QUERY_INVALID', + 'Session list limit must be an integer from 1 to 100' + ); + } + params.limit = limit; + } + const startParam = url.searchParams.get('start'); + if (startParam !== null) { + const start = Number(startParam); + if (!Number.isSafeInteger(start) || start < 0 || start > MAX_TIMESTAMP_MILLISECONDS) { + return facadeError( + 400, + 'KILO_QUERY_INVALID', + 'Session list start must be a non-negative integer' + ); + } + params.start = start; + } + return params; +} + +function unsupportedSessionSelectorResponse(): Response { + return facadeError( + 400, + 'KILO_SESSION_SELECTOR_UNSUPPORTED', + 'Only the matching public Cloud Agent session directory selector is supported' + ); +} + +function validateIdScopedSelectors( + url: URL, + kiloSessionId: string, + paginationKeys: Set +): Response | null { + for (const key of url.searchParams.keys()) { + if (paginationKeys.has(key)) continue; + if (key !== 'directory') return unsupportedSessionSelectorResponse(); + } + const directory = url.searchParams.get('directory'); + if (directory !== null && directory !== publicCloudAgentDirectory(kiloSessionId)) { + return unsupportedSessionSelectorResponse(); + } + return null; +} + +function parseSessionMessagesQuery( + url: URL +): Pick | Response { + const seenParams = new Set(); + for (const key of url.searchParams.keys()) { + if (!SUPPORTED_SESSION_MESSAGES_QUERY_PARAMS.has(key)) { + return unsupportedSessionSelectorResponse(); + } + if (seenParams.has(key)) { + return facadeError( + 400, + 'KILO_QUERY_INVALID', + 'Session messages query parameters must be unique' + ); + } + seenParams.add(key); + } + + const params: Pick = {}; + const limitParam = url.searchParams.get('limit'); + if (limitParam !== null) { + const limit = Number(limitParam); + if (!Number.isInteger(limit) || limit < 0 || limit > MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE) { + return facadeError( + 400, + 'KILO_QUERY_INVALID', + `Session messages limit must be an integer from 0 to ${MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE}` + ); + } + params.limit = limit; + } + const before = url.searchParams.get('before'); + if (before !== null) { + if (before.length === 0 || params.limit === undefined || params.limit === 0) { + return facadeError( + 400, + 'KILO_QUERY_INVALID', + 'Session messages before requires a positive limit' + ); + } + if (!validateKiloSdkMessagesCursor(before)) { + return facadeError( + 400, + 'KILO_QUERY_INVALID', + 'Session messages before is not a valid cursor' + ); + } + params.before = before; + } + return params; +} + +function isExactSessionDetailRead(method: string, kiloPath: string, routePath: string): boolean { + return method === 'GET' && kiloPath === `/session/${routePath}`; +} + +function isExactSessionMessagesRead(method: string, kiloPath: string, routePath: string): boolean { + return method === 'GET' && kiloPath === `/session/${routePath}/message`; +} + +function isExactSessionPromptAsync(method: string, kiloPath: string, routePath: string): boolean { + return method === 'POST' && kiloPath === `/session/${routePath}/prompt_async`; +} + +function isExactSessionAbort(method: string, kiloPath: string, routePath: string): boolean { + return method === 'POST' && kiloPath === `/session/${routePath}/abort`; +} + +function promptAdmissionError( + result: Extract +): Response { + switch (result.code) { + case 'BAD_REQUEST': + return facadeError(400, 'KILO_PROMPT_ADMISSION_REJECTED', result.error); + case 'NOT_FOUND': + return missingRootKiloSessionResponse(); + case 'PENDING_QUEUE_FULL': + return facadeError(429, 'KILO_PROMPT_QUEUE_FULL', result.error); + case 'SANDBOX_CONNECT_FAILED': + case 'WORKSPACE_SETUP_FAILED': + case 'KILO_SERVER_FAILED': + case 'WRAPPER_START_FAILED': + return facadeError(503, result.code, result.error); + case 'INTERNAL': + return facadeError(500, 'KILO_PROMPT_ADMISSION_FAILED', result.error); + } +} + +function promptPreflightError(error: unknown): Response { + if (!(error instanceof TRPCError)) throw error; + switch (error.code) { + case 'BAD_REQUEST': + return facadeError(400, 'KILO_PROMPT_ADMISSION_REJECTED', error.message); + case 'NOT_FOUND': + return missingRootKiloSessionResponse(); + case 'FORBIDDEN': + return facadeError(403, 'KILO_PROMPT_ADMISSION_REJECTED', error.message); + case 'SERVICE_UNAVAILABLE': + return facadeError(503, 'MODEL_VALIDATION_UNAVAILABLE', error.message); + default: + throw error; + } +} + +async function readRequestJson(request: Request): Promise { + const declaredLength = request.headers.get('content-length'); + if (declaredLength !== null) { + const bodyBytes = Number(declaredLength); + if (!Number.isSafeInteger(bodyBytes) || bodyBytes > MAX_KILO_PROMPT_JSON_BYTES) { + return facadeError( + 400, + 'KILO_BASIC_PROMPT_UNSUPPORTED', + 'Basic Kilo prompt body is not supported' + ); + } + } + const response = new Response(request.body); + const bytes = await readBoundedBody(response, MAX_KILO_PROMPT_JSON_BYTES); + if (!bytes) { + return facadeError( + 400, + 'KILO_BASIC_PROMPT_UNSUPPORTED', + 'Basic Kilo prompt body is not supported' + ); + } + try { + return JSON.parse(new TextDecoder().decode(bytes)); + } catch { + return facadeError( + 400, + 'KILO_BASIC_PROMPT_UNSUPPORTED', + 'Basic Kilo prompt body is not supported' + ); + } +} + +async function defaultValidatePromptBalance(params: { + env: Env; + authToken: string; + userId: string; + cloudAgentSessionId: string; +}): Promise { + const orgId = await fetchOrgIdForSession(params.env, params.userId, params.cloudAgentSessionId); + return validateBalanceOnly(params.authToken, orgId, params.env); +} + +async function defaultAdmitPrompt(params: { + env: Env; + userId: string; + cloudAgentSessionId: string; + request: SubmittedSessionMessageRequest; +}): Promise { + const id = params.env.CLOUD_AGENT_SESSION.idFromName( + `${params.userId}:${params.cloudAgentSessionId}` + ); + return withDORetry, SessionMessageAdmissionResult>( + () => params.env.CLOUD_AGENT_SESSION.get(id), + stub => stub.admitSubmittedMessage(params.request), + 'admitSubmittedMessage' + ); +} + +async function admitBasicPrompt(params: { + request: Request; + env: Env; + userId: string; + authToken?: string; + cloudAgentSessionId: string; + deps?: KiloFacadeRequestDeps; +}): Promise { + const value = await readRequestJson(params.request); + if (value instanceof Response) { + return value; + } + const parsed = parseBasicKiloPrompt(value); + if (!parsed.success) { + return facadeError( + 400, + 'KILO_BASIC_PROMPT_UNSUPPORTED', + 'Basic Kilo prompt body is not supported' + ); + } + if (!params.authToken) { + return facadeError(500, 'KILO_FACADE_UNAVAILABLE', 'Durable prompt admission is unavailable'); + } + if (params.request.headers.get('x-skip-balance-check') !== null) { + return facadeError( + 400, + 'KILO_BALANCE_BYPASS_UNSUPPORTED', + 'Balance bypass is not supported for public Kilo prompt mutations' + ); + } + const balance = await (params.deps?.validatePromptBalance ?? defaultValidatePromptBalance)({ + env: params.env, + authToken: params.authToken, + userId: params.userId, + cloudAgentSessionId: params.cloudAgentSessionId, + }); + if (!balance.success) { + return facadeError(balance.status, 'KILO_BALANCE_VALIDATION_FAILED', balance.message); + } + const request: SubmittedSessionMessageRequest = { + userId: params.userId as UserId, + turn: { + type: 'prompt', + id: parsed.prompt.messageId, + prompt: parsed.prompt.prompt, + }, + }; + const admitPrompt = params.deps?.admitPrompt ?? defaultAdmitPrompt; + try { + return await preflightAndAdmitPromptMessage( + { cloudAgentSessionId: params.cloudAgentSessionId, turn: request.turn }, + { env: params.env, userId: params.userId }, + 'kilo.prompt_async', + () => + admitPrompt({ + env: params.env, + userId: params.userId, + cloudAgentSessionId: params.cloudAgentSessionId, + request, + }) + ); + } catch (error) { + return promptPreflightError(error); + } +} + +async function handlePromptAsyncMutation(params: { + request: Request; + env: Env; + userId: string; + authToken?: string; + cloudAgentSessionId: string; + deps?: KiloFacadeRequestDeps; +}): Promise { + const admission = await admitBasicPrompt(params); + if (admission instanceof Response) { + return admission; + } + if (!admission.success) { + return promptAdmissionError(admission); + } + return new Response(null, { status: 204 }); +} + +async function defaultInterruptPrompt(params: { + env: Env; + userId: string; + cloudAgentSessionId: string; +}): Promise>> { + const id = params.env.CLOUD_AGENT_SESSION.idFromName( + `${params.userId}:${params.cloudAgentSessionId}` + ); + return withDORetry< + DurableObjectStub, + Awaited> + >( + () => params.env.CLOUD_AGENT_SESSION.get(id), + stub => stub.interruptExecution(), + 'interruptExecution' + ); +} + +async function handleAbortMutation(params: { + env: Env; + userId: string; + cloudAgentSessionId: string; + deps?: KiloFacadeRequestDeps; +}): Promise { + await (params.deps?.interruptPrompt ?? defaultInterruptPrompt)({ + env: params.env, + userId: params.userId, + cloudAgentSessionId: params.cloudAgentSessionId, + }); + return Response.json(true); +} + +function messagesPageResponse( + requestUrl: URL, + messages: unknown, + nextCursor: string | null, + omittedItemCount: number +): Response { + const exposedHeaders = ['X-Kilo-Omitted-Item-Count']; + const headers = new Headers({ + 'content-type': 'application/json', + 'X-Kilo-Omitted-Item-Count': String(omittedItemCount), + }); + if (nextCursor !== null) { + const nextUrl = new URL(requestUrl); + nextUrl.searchParams.set('before', nextCursor); + exposedHeaders.push('Link', 'X-Next-Cursor'); + headers.set('Link', `<${nextUrl.toString()}>; rel="next"`); + headers.set('X-Next-Cursor', nextCursor); + } + headers.set('Access-Control-Expose-Headers', exposedHeaders.join(', ')); + return new Response(JSON.stringify(messages), { headers }); +} + +async function persistedSessionDetailResponse(params: { + env: Env; + userId: string; + kiloSessionId: string; +}): Promise { + const snapshot = await params.env.SESSION_INGEST.getCloudAgentRootSessionSnapshot({ + kiloUserId: params.userId, + kiloSessionId: params.kiloSessionId, + }); + if (!snapshot) { + return missingRootKiloSessionResponse(); + } + switch (snapshot.snapshot.kind) { + case 'pending': + return pendingSessionSnapshotResponse(); + case 'too_large': + return sessionSnapshotTooLargeResponse(); + case 'retryable_failure': + return retryableSessionReadResponse(); + case 'invalid_data': + return invalidPersistedSessionDataResponse('session'); + case 'value': { + const publicInfo = projectSafePublicSession(snapshot.snapshot.info, snapshot.kiloSessionId); + return publicInfo instanceof Response ? publicInfo : Response.json(publicInfo); + } + } +} + +async function persistedSessionMessagesResponse(params: { + env: Env; + userId: string; + kiloSessionId: string; + url: URL; + query: Pick; +}): Promise { + const result = await params.env.SESSION_INGEST.getCloudAgentRootSessionMessages({ + kiloUserId: params.userId, + kiloSessionId: params.kiloSessionId, + ...params.query, + }); + if (!result) { + return missingRootKiloSessionResponse(); + } + if (result.history === null) { + return pendingSessionSnapshotResponse(); + } + if ('kind' in result.history) { + switch (result.history.kind) { + case 'too_large': + return transcriptTooLargeResponse(); + case 'retryable_failure': + return retryableSessionReadResponse(); + case 'invalid_data': + return invalidPersistedSessionDataResponse('messages'); + } + } + const publicMessages = projectSafePublicMessages(result.history.messages, params.kiloSessionId); + if (publicMessages instanceof Response) return publicMessages; + return messagesPageResponse( + params.url, + publicMessages, + result.history.nextCursor, + result.history.omittedItemCount ?? 0 + ); +} + +type PersistedSessionRead = + | { kind: 'detail' } + | { + kind: 'messages'; + query: Pick; + }; + +function persistedSessionReadResponse(params: { + env: Env; + userId: string; + kiloSessionId: string; + url: URL; + read: PersistedSessionRead; +}): Promise { + switch (params.read.kind) { + case 'detail': + return persistedSessionDetailResponse(params); + case 'messages': + return persistedSessionMessagesResponse({ ...params, query: params.read.query }); + } +} + +export function isSyntheticGlobalEvent(event: KiloGlobalEventEnvelope): boolean { + const type = event.payload?.type; + return type === 'server.connected' || type === 'server.heartbeat'; +} + +function parsePublicSessionDirectory(directory: string): string | null { + const match = /^\/cloud-agent\/sessions\/([^/]+)$/.exec(directory); + if (!match?.[1]) return null; + try { + return decodeURIComponent(match[1]); + } catch { + return null; + } +} + +export function rewriteGlobalEventDirectory( + event: KiloGlobalEventEnvelope, + kiloSessionId: string +): KiloGlobalEventEnvelope { + return { + ...event, + directory: publicCloudAgentDirectory(kiloSessionId), + }; +} + +export function encodeSseData(data: unknown): Uint8Array { + return encoder.encode(`data: ${JSON.stringify(data)}\n\n`); +} + +function encodeSseComment(comment: string): Uint8Array { + return encoder.encode(`: ${comment}\n\n`); +} + +export async function defaultResolveRootSessionForKiloSession(params: { + env: Env; + userId: string; + kiloSessionId: string; +}): Promise<{ cloudAgentSessionId: string } | null> { + return params.env.SESSION_INGEST.resolveCloudAgentRootSessionForKiloSession({ + kiloUserId: params.userId, + kiloSessionId: params.kiloSessionId, + }); +} + +async function eventStreamResponse(params: { + request: Request; + env: Env; + userId: string; + deps?: KiloFacadeRequestDeps; + url: URL; + kiloPath: string; +}): Promise { + const { request, env, userId, deps, url, kiloPath } = params; + if (kiloPath === '/global/event') { + if (request.method !== 'GET') { + return facadeError(501, 'KILO_ROUTE_UNSUPPORTED', 'Kilo facade route is not supported'); + } + if (url.search.length > 0) { + return facadeError( + 400, + 'KILO_EVENT_SELECTOR_UNSUPPORTED', + 'Global event selectors are not supported' + ); + } + const globalEvents = deps?.globalEvents; + if (!globalEvents) { + return facadeError(500, 'KILO_FACADE_UNAVAILABLE', 'Kilo facade stream is unavailable'); + } + return globalEvents.openPublicGlobalEventStream(); + } + + if (kiloPath !== '/event') { + return null; + } + if (request.method !== 'GET') { + return facadeError(501, 'KILO_ROUTE_UNSUPPORTED', 'Kilo facade route is not supported'); + } + if ([...url.searchParams.keys()].some(key => key !== 'directory')) { + return facadeError( + 400, + 'KILO_EVENT_SELECTOR_UNSUPPORTED', + 'Only the Cloud Agent session directory selector is supported' + ); + } + const directory = url.searchParams.get('directory'); + if (!directory) { + return facadeError( + 400, + 'KILO_EVENT_DIRECTORY_REQUIRED', + 'A Cloud Agent session directory selector is required' + ); + } + const kiloSessionId = parsePublicSessionDirectory(directory); + if (!kiloSessionId) { + return facadeError( + 400, + 'KILO_EVENT_SELECTOR_UNSUPPORTED', + 'Only a Cloud Agent root session directory selector is supported' + ); + } + if (!isSessionIngestKiloSessionId(kiloSessionId)) { + return missingRootKiloSessionResponse(); + } + const resolveRoot = + deps?.resolveRootSessionForKiloSession ?? defaultResolveRootSessionForKiloSession; + const root = await resolveRoot({ env, userId, kiloSessionId }); + if (!root) { + return missingRootKiloSessionResponse(); + } + const globalEvents = deps?.globalEvents; + if (!globalEvents) { + return facadeError(500, 'KILO_FACADE_UNAVAILABLE', 'Kilo facade stream is unavailable'); + } + return globalEvents.openPublicSessionEventStream(kiloSessionId); +} + +async function proxyOwnedKiloSessionRequest(params: { + request: Request; + env: Env; + userId: string; + deps?: KiloFacadeRequestDeps; + url: URL; + kiloPath: string; + kiloSessionId: string; + cloudAgentSessionId: string; + persistedRead: PersistedSessionRead | null; +}): Promise { + const { request, env, userId, deps, url, kiloPath, kiloSessionId, cloudAgentSessionId } = params; + const persistedFallback = () => + params.persistedRead + ? persistedSessionReadResponse({ + env, + userId, + kiloSessionId, + url, + read: params.persistedRead, + }) + : null; + + let liveWrapper: LiveWrapperTarget | null; + try { + liveWrapper = await (deps?.resolveLiveWrapper ?? resolveLiveWrapperTarget)({ + env, + userId, + cloudAgentSessionId, + }); + } catch (error) { + const fallback = persistedFallback(); + if (fallback) return fallback; + throw error; + } + if (!liveWrapper) { + const fallback = persistedFallback(); + if (fallback) return fallback; + return facadeError( + 503, + 'KILO_LIVE_RUNTIME_UNAVAILABLE', + 'Cloud Agent Kilo runtime is not live' + ); + } + + const upstreamSearchParams = new URLSearchParams(url.searchParams); + upstreamSearchParams.delete('directory'); + const upstreamSearch = upstreamSearchParams.size > 0 ? `?${upstreamSearchParams.toString()}` : ''; + const targetUrl = buildWrapperKiloProxyUrl({ + wrapperPort: liveWrapper.port, + kiloRelativePath: kiloPath, + search: upstreamSearch, + }); + const proxyRequest = createProxyRequest(request, targetUrl); + let response: Response; + try { + response = await liveWrapper.sandbox.containerFetch(proxyRequest, liveWrapper.port); + } catch (error) { + const fallback = persistedFallback(); + if (fallback) return fallback; + throw error; + } + + if (!params.persistedRead) { + return response; + } + if (await isUnavailableKiloRuntimeResponse(response)) { + return persistedSessionReadResponse({ + env, + userId, + kiloSessionId, + url, + read: params.persistedRead, + }); + } + switch (params.persistedRead.kind) { + case 'detail': + return rewriteLiveSessionDetailResponse(response, kiloSessionId); + case 'messages': + return rewriteLiveMessagesResponse(response, kiloSessionId); + } +} + +export async function handleKiloFacadeRequest(params: { + request: Request; + env: Env; + userId: string; + authToken?: string; + deps?: KiloFacadeRequestDeps; +}): Promise { + const { request, env, userId, authToken, deps } = params; + const url = new URL(request.url); + const kiloPath = kiloRelativePath(url.pathname); + + const streamResponse = await eventStreamResponse({ request, env, userId, deps, url, kiloPath }); + if (streamResponse) { + return streamResponse; + } + + if (kiloPath === '/session' && request.method === 'GET') { + const query = parseSessionListQuery(url); + if (query instanceof Response) { + return query; + } + const sessions = await env.SESSION_INGEST.listCloudAgentRootSessions({ + ...query, + kiloUserId: userId, + }); + return Response.json(sessions.map(projectPublicListedSession)); + } + + if (kiloPath === '/session' || KNOWN_UNSUPPORTED_ROUTES.has(`${request.method} ${kiloPath}`)) { + return facadeError(501, 'KILO_ROUTE_UNSUPPORTED', 'Kilo facade route is not supported'); + } + + const route = parseRootSessionRoute(kiloPath); + if (!route) { + return facadeError(501, 'KILO_ROUTE_UNSUPPORTED', 'Kilo facade route is not supported'); + } + if (!isSessionIngestKiloSessionId(route.kiloSessionId)) { + return missingRootKiloSessionResponse(); + } + + const resolveRoot = + deps?.resolveRootSessionForKiloSession ?? defaultResolveRootSessionForKiloSession; + const root = await resolveRoot({ env, userId, kiloSessionId: route.kiloSessionId }); + if (!root) { + return missingRootKiloSessionResponse(); + } + + const policyInput: SessionKiloFacadePolicyInput = { + method: request.method, + kiloRelativePath: kiloPath, + search: url.search, + userId, + kiloSessionId: route.kiloSessionId, + cloudAgentSessionId: root.cloudAgentSessionId, + }; + const decision = (deps?.decideSessionRoute ?? decideSessionKiloFacadeRoute)(policyInput); + if (decision.kind === 'reject') { + return facadeError(decision.status, decision.code, decision.message); + } + + const detailRead = isExactSessionDetailRead(request.method, kiloPath, route.encodedKiloSessionId); + const messagesRoute = isExactSessionMessagesRead( + request.method, + kiloPath, + route.encodedKiloSessionId + ); + const mutationRoute = + isExactSessionPromptAsync(request.method, kiloPath, route.encodedKiloSessionId) || + isExactSessionAbort(request.method, kiloPath, route.encodedKiloSessionId); + if (detailRead || messagesRoute || mutationRoute) { + const paginationKeys = + request.method === 'GET' && messagesRoute ? new Set(['limit', 'before']) : new Set(); + const selectorResponse = validateIdScopedSelectors(url, route.kiloSessionId, paginationKeys); + if (selectorResponse) return selectorResponse; + } + + if (isExactSessionPromptAsync(request.method, kiloPath, route.encodedKiloSessionId)) { + return handlePromptAsyncMutation({ + request, + env, + userId, + authToken, + cloudAgentSessionId: root.cloudAgentSessionId, + deps, + }); + } + if (isExactSessionAbort(request.method, kiloPath, route.encodedKiloSessionId)) { + return handleAbortMutation({ + env, + userId, + cloudAgentSessionId: root.cloudAgentSessionId, + deps, + }); + } + + const sessionDetailRead = isExactSessionDetailRead( + request.method, + kiloPath, + route.encodedKiloSessionId + ); + const sessionMessagesRead = isExactSessionMessagesRead( + request.method, + kiloPath, + route.encodedKiloSessionId + ); + const messageQuery = sessionMessagesRead ? parseSessionMessagesQuery(url) : null; + if (messageQuery instanceof Response) { + return messageQuery; + } + const persistedRead: PersistedSessionRead | null = sessionDetailRead + ? { kind: 'detail' } + : sessionMessagesRead && messageQuery !== null + ? { kind: 'messages', query: messageQuery } + : null; + return proxyOwnedKiloSessionRequest({ + request, + env, + userId, + deps, + url, + kiloPath, + kiloSessionId: route.kiloSessionId, + cloudAgentSessionId: root.cloudAgentSessionId, + persistedRead, + }); +} + +function parseGlobalFeedSource(request: Request): GlobalFeedSource | Response { + const url = new URL(request.url); + const userId = url.searchParams.get('userId'); + const cloudAgentSessionId = url.searchParams.get('cloudAgentSessionId'); + const kiloSessionId = url.searchParams.get('kiloSessionId'); + const wrapperRunId = url.searchParams.get('wrapperRunId'); + const wrapperGenerationParam = url.searchParams.get('wrapperGeneration'); + const wrapperConnectionId = url.searchParams.get('wrapperConnectionId'); + const wrapperGeneration = wrapperGenerationParam ? Number(wrapperGenerationParam) : NaN; + + if ( + !userId || + !cloudAgentSessionId || + !kiloSessionId || + !wrapperRunId || + !Number.isInteger(wrapperGeneration) || + wrapperGeneration < 0 || + !wrapperConnectionId + ) { + return facadeError(400, 'INVALID_GLOBAL_FEED_SOURCE', 'Invalid global feed source'); + } + + return { + userId, + cloudAgentSessionId, + kiloSessionId, + wrapperRunId, + wrapperGeneration, + wrapperConnectionId, + }; +} + +function producerTag(source: Pick): string { + return `kilo-global:${source.cloudAgentSessionId}`; +} + +function isSameGlobalFeedProducer(left: GlobalFeedSource, right: GlobalFeedSource): boolean { + return ( + left.wrapperRunId === right.wrapperRunId && + left.wrapperGeneration === right.wrapperGeneration && + left.wrapperConnectionId === right.wrapperConnectionId + ); +} + +function mayReplaceGlobalFeedProducer( + existing: GlobalFeedSource, + candidate: GlobalFeedSource +): boolean { + if (candidate.wrapperGeneration > existing.wrapperGeneration) return true; + if (candidate.wrapperGeneration < existing.wrapperGeneration) return false; + return isSameGlobalFeedProducer(existing, candidate); +} + +export class UserKiloFacade extends DurableObject implements KiloFacadeGlobalEvents { + private subscribers = new Map(); + private producerMessageTails = new WeakMap>(); + + async fetch(request: Request): Promise { + const url = new URL(request.url); + + if (url.pathname === KILO_FACADE_GLOBAL_FEED_PATH) { + return this.handleGlobalFeedRequest(request); + } + + const userId = request.headers.get(KILO_FACADE_USER_ID_HEADER); + if (!userId) { + return facadeError(401, 'UNAUTHENTICATED', 'Missing authenticated user context'); + } + const authToken = request.headers.get(KILO_FACADE_AUTH_TOKEN_HEADER) ?? undefined; + + return handleKiloFacadeRequest({ + request, + env: this.env, + userId, + authToken, + deps: { globalEvents: this }, + }); + } + + openPublicGlobalEventStream(): Response { + return this.openPublicEventStream({ kind: 'global' }); + } + + openPublicSessionEventStream(kiloSessionId: string): Response { + return this.openPublicEventStream({ kind: 'session', kiloSessionId }); + } + + async publishCloudAgentExtensionEvent(input: { + kiloUserId: string; + cloudAgentSessionId: string; + kiloSessionId: string; + organizationId?: string; + event: PublicCloudAgentExtensionEvent; + }): Promise { + if (input.organizationId) { + const sessionIngest = this.env.SESSION_INGEST; + if (!sessionIngest) return; + const resolved = await sessionIngest.resolveCloudAgentRootSessionForKiloSession({ + kiloUserId: input.kiloUserId, + kiloSessionId: input.kiloSessionId, + }); + if (resolved?.cloudAgentSessionId !== input.cloudAgentSessionId) return; + } + this.broadcastGlobalEvent( + { directory: publicCloudAgentDirectory(input.kiloSessionId), payload: input.event }, + input.kiloSessionId + ); + } + + private openPublicEventStream(scope: PublicSubscriber['scope']): Response { + const subscriberId = crypto.randomUUID(); + const publicConnectionEvent = () => + scope.kind === 'global' + ? { + directory: PUBLIC_VIRTUAL_SERVER_DIRECTORY, + payload: createPublicEventPayload('server.connected'), + } + : createPublicEventPayload('server.connected'); + + const stream = new ReadableStream( + { + start: controller => { + const subscriber: PublicSubscriber = { + controller, + scope, + heartbeat: setInterval(() => { + this.sendCommentToSubscriber(subscriberId, 'heartbeat'); + }, HEARTBEAT_INTERVAL_MS), + }; + this.subscribers.set(subscriberId, subscriber); + this.sendToSubscriber(subscriberId, publicConnectionEvent()); + }, + cancel: () => { + this.removeSubscriber(subscriberId); + }, + }, + { highWaterMark: MAX_PUBLIC_GLOBAL_EVENT_QUEUE_SIZE, size: () => 1 } + ); + + return new Response(stream, { + headers: { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache, no-transform', + 'X-Accel-Buffering': 'no', + 'X-Content-Type-Options': 'nosniff', + }, + }); + } + + async webSocketMessage(ws: WebSocket, message: string | ArrayBuffer): Promise { + const previous = this.producerMessageTails.get(ws) ?? Promise.resolve(); + const next = previous + .catch(() => undefined) + .then(() => this.handleGlobalFeedMessage(ws, message)); + this.producerMessageTails.set(ws, next); + try { + await next; + } finally { + if (this.producerMessageTails.get(ws) === next) { + this.producerMessageTails.delete(ws); + } + } + } + + private validateGlobalFeedProducer(source: GlobalFeedSource) { + const sessionDoId = this.env.CLOUD_AGENT_SESSION.idFromName( + `${source.userId}:${source.cloudAgentSessionId}` + ); + const sessionStub = this.env.CLOUD_AGENT_SESSION.get(sessionDoId); + return sessionStub.validateKiloGlobalFeedProducer({ + kiloSessionId: source.kiloSessionId, + wrapperRunId: source.wrapperRunId, + wrapperGeneration: source.wrapperGeneration, + wrapperConnectionId: source.wrapperConnectionId, + }); + } + + private async handleGlobalFeedMessage( + ws: WebSocket, + message: string | ArrayBuffer + ): Promise { + if (typeof message !== 'string') { + ws.send(JSON.stringify({ error: 'Binary global feed messages are not supported' })); + return; + } + + const source = ws.deserializeAttachment() as GlobalFeedSource | null; + if (!source) { + ws.close(1011, 'Missing global feed source'); + return; + } + + const validation = await this.validateGlobalFeedProducer(source); + if (!validation.success) { + ws.close(4401, validation.message); + return; + } + + let parsed: unknown; + try { + parsed = JSON.parse(message); + } catch { + ws.send(JSON.stringify({ error: 'Invalid global feed JSON' })); + return; + } + + if (!isKiloGlobalEventEnvelope(parsed)) { + ws.send(JSON.stringify({ error: 'Invalid global event envelope' })); + return; + } + + if (isSyntheticGlobalEvent(parsed)) { + return; + } + if (!isSubstantiveKiloGlobalEventEnvelope(parsed)) { + ws.send(JSON.stringify({ error: 'Invalid global event envelope' })); + return; + } + if (isPublicCloudAgentExtensionSourceType(parsed.payload.type)) { + return; + } + + const projectedPayload = projectPublicEvent(parsed.payload, source.kiloSessionId); + if (!projectedPayload) { + return; + } + const publicEnvelope = rewriteGlobalEventDirectory( + { ...parsed, payload: projectedPayload }, + source.kiloSessionId + ); + if (hasUnprojectedPrivateStructuredPath(publicEnvelope)) { + return; + } + this.broadcastGlobalEvent(publicEnvelope, source.kiloSessionId); + } + + async webSocketClose(ws: WebSocket): Promise { + const source = ws.deserializeAttachment() as GlobalFeedSource | null; + if (!source) return; + const tag = producerTag(source); + const remaining = this.ctx.getWebSockets(tag).filter(existing => existing !== ws); + if (remaining.length === 0) { + await this.ctx.storage.delete(`producer:${source.cloudAgentSessionId}`); + } + } + + private async handleGlobalFeedRequest(request: Request): Promise { + if (request.headers.get('Upgrade')?.toLowerCase() !== 'websocket') { + return facadeError(426, 'WEBSOCKET_REQUIRED', 'Expected WebSocket upgrade'); + } + + const source = parseGlobalFeedSource(request); + if (source instanceof Response) { + return source; + } + + const validation = await this.validateGlobalFeedProducer(source); + if (!validation.success) { + return new Response(validation.message, { status: validation.status }); + } + + const tag = producerTag(source); + const existingSockets = this.ctx.getWebSockets(tag); + for (const existing of existingSockets) { + const existingSource = existing.deserializeAttachment() as GlobalFeedSource | null; + if (existingSource && !mayReplaceGlobalFeedProducer(existingSource, source)) { + return new Response('A newer global feed producer is already connected', { status: 409 }); + } + } + for (const existing of existingSockets) { + try { + existing.close(1000, 'Replaced by newer global feed'); + } catch { + // Ignore already-closed producer sockets. + } + } + + const pair = new WebSocketPair(); + const client = pair[0]; + const server = pair[1]; + this.ctx.acceptWebSocket(server, [tag]); + server.serializeAttachment(source); + await this.ctx.storage.put(`producer:${source.cloudAgentSessionId}`, { + wrapperRunId: source.wrapperRunId, + wrapperGeneration: source.wrapperGeneration, + wrapperConnectionId: source.wrapperConnectionId, + }); + + return new Response(null, { status: 101, webSocket: client }); + } + + private sendFrameToSubscriber(subscriberId: string, frame: Uint8Array): void { + const subscriber = this.subscribers.get(subscriberId); + if (!subscriber) return; + if (subscriber.controller.desiredSize !== null && subscriber.controller.desiredSize <= 0) { + this.removeSubscriber(subscriberId); + subscriber.controller.error(new Error(SLOW_SUBSCRIBER_ERROR)); + return; + } + try { + subscriber.controller.enqueue(frame); + } catch { + this.removeSubscriber(subscriberId); + } + } + + private sendToSubscriber(subscriberId: string, event: unknown): void { + this.sendFrameToSubscriber(subscriberId, encodeSseData(event)); + } + + private sendCommentToSubscriber(subscriberId: string, comment: string): void { + this.sendFrameToSubscriber(subscriberId, encodeSseComment(comment)); + } + + private broadcastGlobalEvent(event: KiloGlobalEventEnvelope, kiloSessionId: string): void { + for (const [subscriberId, subscriber] of this.subscribers.entries()) { + if (subscriber.scope.kind === 'global') { + this.sendToSubscriber(subscriberId, event); + continue; + } + if (subscriber.scope.kiloSessionId === kiloSessionId && event.payload) { + this.sendToSubscriber(subscriberId, event.payload); + } + } + } + + private removeSubscriber(subscriberId: string): void { + const subscriber = this.subscribers.get(subscriberId); + if (!subscriber) return; + clearInterval(subscriber.heartbeat); + this.subscribers.delete(subscriberId); + } +} diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index e3f6ae6318..218196a72b 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -58,7 +58,12 @@ import { } from '../websocket/ingest.js'; import type { StoredEvent } from '../websocket/types.js'; import type { WrapperCommand, CloudStatusData } from '../shared/protocol.js'; +import { + isPublicCloudAgentExtensionSourceType, + projectPublicCloudAgentExtensionEvent, +} from '../kilo-facade/cloud-agent-extension-events.js'; import { commandsOrDefault, type SlashCommandInfo } from '../shared/slash-commands.js'; +import { withDORetry } from '../utils/do-retry.js'; import type { AcceptedExecutionTurn, AgentSelection, @@ -95,6 +100,11 @@ import { getWrapperRuntimeState, nextWrapperLeaseDeadline, } from '../session/wrapper-runtime-state.js'; +import { + kiloGlobalFeedValidationSchema, + validateKiloGlobalFeedProducerIdentity, + type KiloGlobalFeedValidationResult, +} from '../session/wrapper-global-feed-validation.js'; import { getSessionMessageState, listNonTerminalAcceptedMessages, @@ -276,6 +286,7 @@ export class CloudAgentSession extends DurableObject { private messageSettlementOutbox?: MessageSettlementOutbox; private sessionMessageQueue?: SessionMessageQueue; private wrapperSupervisor?: WrapperSupervisor; + private publicExtensionPublicationTail: Promise = Promise.resolve(); private isTerminalStatus( status: ExecutionStatus ): status is 'completed' | 'failed' | 'interrupted' { @@ -806,6 +817,7 @@ export class CloudAgentSession extends DurableObject { * @param event - The stored event to broadcast */ broadcastEvent(event: StoredEvent): void { + this.publishPublicCloudAgentExtensionEvent(event); if (this.streamHandler) { this.streamHandler.broadcastEvent(event); return; @@ -824,6 +836,39 @@ export class CloudAgentSession extends DurableObject { }); } + private publishPublicCloudAgentExtensionEvent(event: StoredEvent): void { + if (!isPublicCloudAgentExtensionSourceType(event.stream_event_type)) return; + const publication = this.publicExtensionPublicationTail + .catch(() => undefined) + .then(async () => { + const metadata = await this.getMetadata(); + const kiloSessionId = metadata?.auth.kiloSessionId; + if (!metadata || !kiloSessionId) return; + const projected = projectPublicCloudAgentExtensionEvent(event, kiloSessionId); + if (!projected) return; + const facadeId = this.env.USER_KILO_FACADE.idFromName(metadata.identity.userId); + await withDORetry( + () => this.env.USER_KILO_FACADE.get(facadeId), + facade => + facade.publishCloudAgentExtensionEvent({ + kiloUserId: metadata.identity.userId, + cloudAgentSessionId: metadata.identity.sessionId, + kiloSessionId, + organizationId: metadata.identity.orgId, + event: projected, + }), + 'publishCloudAgentExtensionEvent' + ); + }) + .catch(error => { + logger + .withFields({ error: error instanceof Error ? error.message : String(error) }) + .warn('Failed to publish public Cloud Agent extension event'); + }); + this.publicExtensionPublicationTail = publication; + this.ctx.waitUntil(publication); + } + private insertAndBroadcastEvent(params: { executionId: EventSourceId; sessionId: string; @@ -972,6 +1017,26 @@ export class CloudAgentSession extends DurableObject { return metadata ? parseSessionMetadata(metadata) : null; } + async validateKiloGlobalFeedProducer(params: { + kiloSessionId: string; + wrapperRunId: string; + wrapperGeneration: number; + wrapperConnectionId: string; + }): Promise { + const parsed = kiloGlobalFeedValidationSchema.safeParse(params); + if (!parsed.success) { + return { success: false, status: 400, message: 'Invalid global feed producer identity' }; + } + + const metadata = await this.getMetadata(); + const runtimeState = await getWrapperRuntimeState(this.ctx.storage); + return validateKiloGlobalFeedProducerIdentity({ + metadata, + runtimeState, + producer: parsed.data, + }); + } + async getLatestAssistantMessage(): Promise { const sessionId = await this.requireSessionId(); const metadata = await this.getMetadata(); diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index 163ace1f3e..4c0101d784 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/services/cloud-agent-next/src/router.test.ts @@ -334,6 +334,7 @@ describe('router sessionId validation', () => { getCurrentRuntimeExecution: vi.fn().mockResolvedValue(null), })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], + USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], @@ -721,6 +722,7 @@ describe('router sessionId validation', () => { idFromName: vi.fn((id: string) => ({ id })), get: vi.fn(() => mockSessionStub), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], + USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], @@ -829,6 +831,7 @@ describe('router sessionId validation', () => { getCurrentMessageWork: mockGetCurrentMessageWork, })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], + USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], @@ -1130,6 +1133,7 @@ describe('router sessionId validation', () => { getCurrentMessageWork: mockGetCurrentMessageWork, })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], + USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], @@ -1379,6 +1383,7 @@ describe('router sessionId validation', () => { getLatestAssistantMessage: mockGetLatestAssistantMessage, })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], + USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], diff --git a/services/cloud-agent-next/src/router/handlers/session-execution.ts b/services/cloud-agent-next/src/router/handlers/session-execution.ts index 6c55f6db6e..ac684d57b9 100644 --- a/services/cloud-agent-next/src/router/handlers/session-execution.ts +++ b/services/cloud-agent-next/src/router/handlers/session-execution.ts @@ -20,7 +20,7 @@ import { LegacyExecutionResponse, } from '../schemas.js'; import type { SessionId } from '../../types/ids.js'; -import { queueMessage, replayMessageIfAlreadyAdmitted } from '../../session/queue-message.js'; +import { preflightAndQueuePromptMessage, queueMessage } from '../../session/queue-message.js'; import { admitLegacyPreparedInitialMessage, replayLegacyPreparedInitialMessageIfAlreadyAdmitted, @@ -31,10 +31,7 @@ import type { TurnFinalization, } from '../../execution/types.js'; import type { QueueAckResponse } from '../schemas.js'; -import { - preflightExistingPromptModel, - preflightPreparedInitialPromptModel, -} from '../../session/model-preflight.js'; +import { preflightPreparedInitialPromptModel } from '../../session/model-preflight.js'; function withLegacyExecutionId(ack: QueueAckResponse): LegacyExecutionResponse { return { @@ -131,20 +128,14 @@ export function createSessionExecutionV2Handlers() { } satisfies TurnFinalization, }; const admissionContext = { env: ctx.env, userId: ctx.userId, botId: ctx.botId }; - if (turn.type === 'prompt') { - const replay = await replayMessageIfAlreadyAdmitted(queuedMessage, admissionContext); - if (replay) return withLegacyExecutionId(replay); - - await preflightExistingPromptModel({ - env: ctx.env, - userId: ctx.userId, - cloudAgentSessionId: input.cloudAgentSessionId, - requestedModel: agent?.model, - procedure: 'sendMessageV2', - }); - } - - const ack = await queueMessage(queuedMessage, admissionContext); + const ack = + turn.type === 'prompt' + ? await preflightAndQueuePromptMessage( + queuedMessage, + admissionContext, + 'sendMessageV2' + ) + : await queueMessage(queuedMessage, admissionContext); return withLegacyExecutionId(ack); }); }), diff --git a/services/cloud-agent-next/src/router/handlers/session-send.ts b/services/cloud-agent-next/src/router/handlers/session-send.ts index 91d4cbe780..676688b135 100644 --- a/services/cloud-agent-next/src/router/handlers/session-send.ts +++ b/services/cloud-agent-next/src/router/handlers/session-send.ts @@ -10,9 +10,8 @@ import { protectedProcedure } from '../auth.js'; import { logger, withLogTags } from '../../logger.js'; import { SendMessageInput, ExecutionResponse } from '../schemas.js'; -import { queueMessage, replayMessageIfAlreadyAdmitted } from '../../session/queue-message.js'; +import { preflightAndQueuePromptMessage } from '../../session/queue-message.js'; import type { SessionId } from '../../types/ids.js'; -import { preflightExistingPromptModel } from '../../session/model-preflight.js'; type SessionSendHandlers = { send: typeof sendMessageHandler; @@ -42,18 +41,6 @@ const sendMessageHandler = protectedProcedure finalization: input.finalization, }; const admissionContext = { env: ctx.env, userId: ctx.userId, botId: ctx.botId }; - const replay = await replayMessageIfAlreadyAdmitted(queuedMessage, admissionContext); - if (replay) return replay; - - await preflightExistingPromptModel({ - env: ctx.env, - userId: ctx.userId, - cloudAgentSessionId: input.cloudAgentSessionId, - requestedModel: input.agent?.model, - procedure: 'send', - }); - - const ack = await queueMessage(queuedMessage, admissionContext); - return ack; + return preflightAndQueuePromptMessage(queuedMessage, admissionContext, 'send'); }); }); diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index 16f8a96da4..9532450ef9 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -25,6 +25,8 @@ vi.mock('./logger.js', () => { return { logger, withLogTags: async (_tags: unknown, fn: () => Promise) => fn(), + WithLogTags: () => (_target: unknown, _propertyKey: string, descriptor: PropertyDescriptor) => + descriptor, }; }); @@ -39,6 +41,12 @@ vi.mock('./agent-sandbox/factory.js', () => ({ })), })); +vi.mock('cloudflare:workers', () => ({ + DurableObject: class DurableObject { + constructor(_state: unknown, _env: unknown) {} + }, +})); + vi.mock('./router.js', () => ({ appRouter: {}, })); @@ -82,6 +90,10 @@ type MockEnv = { idFromName: ReturnType; get: ReturnType; }; + USER_KILO_FACADE: { + idFromName: ReturnType; + get: ReturnType; + }; }; function createEnv(): MockEnv { @@ -93,6 +105,10 @@ function createEnv(): MockEnv { idFromName: vi.fn(), get: vi.fn(), }, + USER_KILO_FACADE: { + idFromName: vi.fn(), + get: vi.fn(), + }, }; } @@ -100,6 +116,17 @@ function fetchWorker(request: Request, env: MockEnv): Promise | Respon return worker.fetch(request, env as unknown as Env, {} as ExecutionContext); } +function signKiloToken(userId = 'usr_1'): string { + return jwt.sign( + { + version: 3, + kiloUserId: userId, + }, + secret, + { algorithm: 'HS256' } + ); +} + beforeEach(() => { getRunningTerminalClientMock.mockReset(); consumeCloudAgentReportBatchMock.mockClear(); @@ -325,3 +352,157 @@ describe('server /terminal', () => { expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); }); }); + +describe('server /kilo facade route', () => { + it('returns 401 before facade dispatch when auth is missing', async () => { + const env = createEnv(); + + const response = await fetchWorker(new Request('http://worker.test/kilo/event'), env); + + expect(response.status).toBe(401); + expect(env.USER_KILO_FACADE.idFromName).not.toHaveBeenCalled(); + expect(env.USER_KILO_FACADE.get).not.toHaveBeenCalled(); + }); + + it('routes valid bearer-authenticated requests to the per-user facade without public credentials', async () => { + const env = createEnv(); + const facadeFetch = vi.fn<(request: Request) => Promise>( + async () => new Response('facade response', { status: 209 }) + ); + env.USER_KILO_FACADE.idFromName.mockReturnValue('facade-id'); + env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); + const token = signKiloToken('usr_facade'); + + const response = await fetchWorker( + new Request('http://worker.test/kilo/session/ses_12345678901234567890123456/message', { + method: 'POST', + headers: { + Authorization: `Bearer ${token}`, + Cookie: 'session=secret', + 'Content-Type': 'application/json', + 'x-kilo-facade-user-id': 'usr_attacker', + 'x-kilo-facade-auth-token': 'attacker-token', + }, + body: JSON.stringify({ ok: true }), + }), + env + ); + + expect(response.status).toBe(209); + expect(env.USER_KILO_FACADE.idFromName).toHaveBeenCalledWith('usr_facade'); + expect(env.USER_KILO_FACADE.get).toHaveBeenCalledWith('facade-id'); + expect(facadeFetch).toHaveBeenCalledOnce(); + + const forwarded = facadeFetch.mock.calls[0][0]; + expect(forwarded.headers.get('authorization')).toBeNull(); + expect(forwarded.headers.get('cookie')).toBeNull(); + expect(forwarded.headers.get('x-kilo-facade-user-id')).toBe('usr_facade'); + expect(forwarded.headers.get('x-kilo-facade-auth-token')).toBe(token); + expect(forwarded.headers.get('content-type')).toBe('application/json'); + await expect(forwarded.text()).resolves.toBe('{"ok":true}'); + }); +}); + +describe('server raw global feed route', () => { + it('validates producer fencing and forwards accepted producer WebSockets to the user facade', async () => { + const env = createEnv(); + const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ validateKiloGlobalFeedProducer }); + const facadeFetch = vi.fn<(request: Request) => Promise>( + async () => new Response('accepted', { status: 200 }) + ); + env.USER_KILO_FACADE.idFromName.mockReturnValue('facade-id'); + env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); + const token = signKiloToken('usr_feed'); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/kilo-global-ingest?kiloSessionId=ses_12345678901234567890123456&wrapperRunId=wr_1&wrapperGeneration=2&wrapperConnectionId=conn_1', + { + headers: { + Upgrade: 'websocket', + Authorization: `Bearer ${token}`, + Cookie: 'session=secret', + 'x-kilo-facade-user-id': 'usr_attacker', + 'x-kilo-facade-auth-token': 'attacker-token', + }, + } + ), + env + ); + + expect(response.status).toBe(200); + expect(env.CLOUD_AGENT_SESSION.idFromName).toHaveBeenCalledWith('usr_feed:agent_live'); + expect(validateKiloGlobalFeedProducer).toHaveBeenCalledWith({ + kiloSessionId: 'ses_12345678901234567890123456', + wrapperRunId: 'wr_1', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_1', + }); + + const forwarded = facadeFetch.mock.calls[0][0]; + const forwardedUrl = new URL(forwarded.url); + expect(forwardedUrl.pathname).toBe('/internal/kilo/global-feed'); + expect(forwardedUrl.searchParams.get('userId')).toBe('usr_feed'); + expect(forwardedUrl.searchParams.get('cloudAgentSessionId')).toBe('agent_live'); + expect(forwardedUrl.searchParams.get('kiloSessionId')).toBe('ses_12345678901234567890123456'); + expect(forwarded.headers.get('upgrade')).toBe('websocket'); + expect(forwarded.headers.get('authorization')).toBeNull(); + expect(forwarded.headers.get('cookie')).toBeNull(); + expect(forwarded.headers.get('x-kilo-facade-user-id')).toBeNull(); + expect(forwarded.headers.get('x-kilo-facade-auth-token')).toBeNull(); + }); + + it('rejects stale producer fencing before facade dispatch', async () => { + const env = createEnv(); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + validateKiloGlobalFeedProducer: vi.fn(async () => ({ + success: false as const, + status: 409, + message: 'Stale wrapper connection', + })), + }); + const facadeFetch = vi.fn(); + env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); + const token = signKiloToken('usr_feed'); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/kilo-global-ingest?kiloSessionId=ses_12345678901234567890123456&wrapperRunId=wr_1&wrapperGeneration=2&wrapperConnectionId=conn_1', + { + headers: { + Upgrade: 'websocket', + Authorization: `Bearer ${token}`, + }, + } + ), + env + ); + + expect(response.status).toBe(409); + expect(facadeFetch).not.toHaveBeenCalled(); + }); + + it('rejects malformed producer generation before session validation', async () => { + const env = createEnv(); + const token = signKiloToken('usr_feed'); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/kilo-global-ingest?kiloSessionId=ses_12345678901234567890123456&wrapperRunId=wr_1&wrapperGeneration=2abc&wrapperConnectionId=conn_1', + { + headers: { + Upgrade: 'websocket', + Authorization: `Bearer ${token}`, + }, + } + ), + env + ); + + expect(response.status).toBe(400); + expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index 79e55ac461..b26a197582 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -17,6 +17,12 @@ import { import { authMiddleware } from './middleware/auth.js'; import { balanceMiddleware } from './middleware/balance.js'; import { resolveTerminalWrapperClient } from './terminal/access.js'; +import { requestMethodAllowsBody } from './shared/http-proxy.js'; +import { + KILO_FACADE_AUTH_TOKEN_HEADER, + KILO_FACADE_GLOBAL_FEED_PATH, + KILO_FACADE_USER_ID_HEADER, +} from './kilo-facade/user-kilo-facade.js'; const app = new Hono(); @@ -126,6 +132,67 @@ app.get('/health', (c: Context) => { }); }); +function createSanitizedForwardRequest( + request: Request, + url: string | URL, + headers: Headers +): Request { + const init: RequestInit & { duplex?: 'half' } = { + method: request.method, + headers, + }; + if (request.body && requestMethodAllowsBody(request.method)) { + init.body = request.body; + init.duplex = 'half'; + } + return new Request(url, init); +} + +function stripPublicCredentialHeaders(headers: Headers): Headers { + const sanitized = new Headers(headers); + sanitized.delete('Authorization'); + sanitized.delete('Cookie'); + sanitized.delete(KILO_FACADE_USER_ID_HEADER); + sanitized.delete(KILO_FACADE_AUTH_TOKEN_HEADER); + return sanitized; +} + +async function routeToUserKiloFacade( + c: Context, + userId: string, + authToken: string +): Promise { + const doId = c.env.USER_KILO_FACADE.idFromName(userId); + const stub = c.env.USER_KILO_FACADE.get(doId); + const headers = stripPublicCredentialHeaders(c.req.raw.headers); + headers.set(KILO_FACADE_USER_ID_HEADER, userId); + headers.set(KILO_FACADE_AUTH_TOKEN_HEADER, authToken); + const request = createSanitizedForwardRequest(c.req.raw, c.req.url, headers); + return stub.fetch(request); +} + +app.all('/kilo', async (c: Context) => { + const authResult = await validateKiloToken( + c.req.header('Authorization') ?? null, + c.env.NEXTAUTH_SECRET + ); + if (!authResult.success) { + return c.text(authResult.error, 401); + } + return routeToUserKiloFacade(c, authResult.userId, authResult.token); +}); + +app.all('/kilo/*', async (c: Context) => { + const authResult = await validateKiloToken( + c.req.header('Authorization') ?? null, + c.env.NEXTAUTH_SECRET + ); + if (!authResult.success) { + return c.text(authResult.error, 401); + } + return routeToUserKiloFacade(c, authResult.userId, authResult.token); +}); + // TODO: I think this and /terminal share a bit of code. Could be worth extracting to middleware or just a common method? app.get('/stream', async (c: Context) => { const upgradeHeader = c.req.header('Upgrade'); @@ -185,6 +252,82 @@ app.get('/terminal', async (c: Context) => { return handleTerminalWebSocket(c.req.raw, c.env); }); +app.all('/sessions/:userId/:sessionId/kilo-global-ingest', async (c: Context) => { + const upgradeHeader = c.req.header('Upgrade'); + if (upgradeHeader?.toLowerCase() !== 'websocket') { + return c.text('Expected WebSocket upgrade', 426); + } + + const rawUserId = c.req.param('userId'); + const cloudAgentSessionId = c.req.param('sessionId'); + if (!rawUserId || !cloudAgentSessionId) { + return c.text('Missing route params', 400); + } + + let userId: string; + try { + userId = decodeURIComponent(rawUserId); + } catch { + return c.text('Invalid userId encoding', 400); + } + + const authResult = await validateKiloToken( + c.req.header('Authorization') ?? null, + c.env.NEXTAUTH_SECRET + ); + if (!authResult.success) { + return c.text(authResult.error, 401); + } + if (authResult.userId !== userId) { + return c.text('Token does not match session user', 403); + } + + const url = new URL(c.req.url); + const kiloSessionId = url.searchParams.get('kiloSessionId'); + const wrapperRunId = url.searchParams.get('wrapperRunId'); + const wrapperGenerationParam = url.searchParams.get('wrapperGeneration'); + const wrapperConnectionId = url.searchParams.get('wrapperConnectionId'); + const wrapperGeneration = wrapperGenerationParam ? Number(wrapperGenerationParam) : NaN; + + if ( + !kiloSessionId || + !wrapperRunId || + !Number.isInteger(wrapperGeneration) || + wrapperGeneration < 0 || + !wrapperConnectionId + ) { + return c.text('Invalid global feed producer identity', 400); + } + + const sessionDoId = c.env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${cloudAgentSessionId}`); + const sessionStub = c.env.CLOUD_AGENT_SESSION.get(sessionDoId); + const validation = await sessionStub.validateKiloGlobalFeedProducer({ + kiloSessionId, + wrapperRunId, + wrapperGeneration, + wrapperConnectionId, + }); + if (!validation.success) { + return new Response(validation.message, { status: validation.status }); + } + + const facadeId = c.env.USER_KILO_FACADE.idFromName(userId); + const facadeStub = c.env.USER_KILO_FACADE.get(facadeId); + const facadeUrl = new URL(c.req.url); + facadeUrl.pathname = KILO_FACADE_GLOBAL_FEED_PATH; + facadeUrl.search = ''; + facadeUrl.searchParams.set('userId', userId); + facadeUrl.searchParams.set('cloudAgentSessionId', cloudAgentSessionId); + facadeUrl.searchParams.set('kiloSessionId', kiloSessionId); + facadeUrl.searchParams.set('wrapperRunId', wrapperRunId); + facadeUrl.searchParams.set('wrapperGeneration', String(wrapperGeneration)); + facadeUrl.searchParams.set('wrapperConnectionId', wrapperConnectionId); + + const headers = stripPublicCredentialHeaders(c.req.raw.headers); + const request = createSanitizedForwardRequest(c.req.raw, facadeUrl, headers); + return facadeStub.fetch(request); +}); + app.all('/sessions/:userId/:sessionId/ingest', async (c: Context) => { const upgradeHeader = c.req.header('Upgrade'); if (upgradeHeader?.toLowerCase() !== 'websocket') { @@ -350,3 +493,4 @@ export default { export { Sandbox } from '@cloudflare/sandbox'; export { CloudAgentSession } from './persistence/CloudAgentSession.js'; +export { UserKiloFacade } from './kilo-facade/user-kilo-facade.js'; diff --git a/services/cloud-agent-next/src/session-ingest-binding.ts b/services/cloud-agent-next/src/session-ingest-binding.ts index de1d0e2fc8..b0894916d5 100644 --- a/services/cloud-agent-next/src/session-ingest-binding.ts +++ b/services/cloud-agent-next/src/session-ingest-binding.ts @@ -1,29 +1,14 @@ /** - * RPC method types for the SESSION_INGEST service binding. + * Local Cloudflare adapter for the SESSION_INGEST service binding. * - * `wrangler types` only sees `Fetcher` for service bindings; the actual RPC - * shape comes from the session-ingest worker's WorkerEntrypoint and is - * declared here so the generated file can be freely regenerated. - * - * Keep in sync with: cloudflare-session-ingest/src/session-ingest-rpc.ts + * `wrangler types` only sees `Fetcher` for service bindings. The shared RPC + * schemas and method types live in `@kilocode/session-ingest-contracts` so the + * producer and consumers compile against one contract while this adapter can + * be regenerated independently from Cloudflare binding types. */ -export type CreateSessionForCloudAgentParams = { - sessionId: string; - kiloUserId: string; - cloudAgentSessionId: string; - organizationId?: string; - createdOnPlatform: string; - title?: string; -}; +import type { SessionIngestRpcMethods } from '@kilocode/session-ingest-contracts'; -export type DeleteSessionForCloudAgentParams = { - sessionId: string; - kiloUserId: string; - onlyIfEmpty?: boolean; -}; +export type * from '@kilocode/session-ingest-contracts'; -export type SessionIngestBinding = Fetcher & { - createSessionForCloudAgent(params: CreateSessionForCloudAgentParams): Promise; - deleteSessionForCloudAgent(params: DeleteSessionForCloudAgentParams): Promise; -}; +export type SessionIngestBinding = Fetcher & SessionIngestRpcMethods; diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index 2a0175d147..a8cf78f6c8 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -1061,10 +1061,11 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(result.ready.devcontainer).toBeUndefined(); }); - it('materializes workspace setup and prompt delivery into separate wrapper requests', async () => { + it('materializes workspace setup and prompt delivery without using provider token overrides for ingest auth', async () => { const service = new SessionService(); const env = createEnv(); env.WORKER_URL = 'https://cloud-agent.example.com'; + env.KILOCODE_TOKEN_OVERRIDE = 'provider-override-token'; const metadata = createMetadata({ setupCommands: ['pnpm install'], envVars: { PUBLIC_VALUE: 'visible' }, @@ -1144,7 +1145,10 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(result.promptRequest).not.toHaveProperty('workspace'); expect(result.promptRequest).not.toHaveProperty('materialized'); expect(result.readyRequest.materialized.env.PUBLIC_VALUE).toBe('visible'); - expect(result.readyRequest.materialized.env.KILOCODE_TOKEN).toBe('kilo-token'); + expect(result.readyRequest.materialized.env.KILOCODE_TOKEN).toBe('provider-override-token'); + expect(JSON.parse(result.readyRequest.materialized.env.KILO_AUTH_CONTENT)).toEqual({ + kilo: { type: 'api', key: 'kilo-token' }, + }); expect(result.readyRequest.materialized.env.GITLAB_TOKEN).toBe('resolved-gitlab-token'); expect(result.readyRequest.materialized.env.GITLAB_HOST).toBe('gitlab.com'); expect(result.readyRequest.materialized.env.GLAB_IS_OAUTH2).toBe('true'); diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index bccc9ed925..94bd3d3b22 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -515,11 +515,9 @@ export async function runSetupCommands( logger.info('Setup commands completed'); } -// Write Kilo auth file so the CLI's KiloSessions can call session ingest. -// The CLI reads ~/.local/share/kilo/auth.json via Auth.get("kilo") but we -// never run `kilo auth login` — credentials are injected purely via env vars -// for config (KILO_CONFIG_CONTENT). The session ingest code path ignores the -// provider config and only reads the auth file. +// Persist Kilo auth for workspace preparation paths that do not receive the +// wrapper process environment; wrapper-launched Kilo receives the same auth +// shape through KILO_AUTH_CONTENT. export async function writeAuthFile( sandbox: SandboxInstance, sessionHome: string, @@ -1027,6 +1025,7 @@ export class SessionService { SESSION_HOME: sessionHome, // Inject Kilocode credentials (with override support) KILOCODE_TOKEN: kilocodeToken, + KILO_AUTH_CONTENT: JSON.stringify({ kilo: { type: 'api', key: originalToken } }), // Platform identifier - defaults to 'cloud-agent' if not specified KILO_PLATFORM: createdOnPlatform ?? 'cloud-agent', KILO_DISABLE_AUTOUPDATE: 'true', diff --git a/services/cloud-agent-next/src/session/queue-message.ts b/services/cloud-agent-next/src/session/queue-message.ts index ba13edac02..189c8b8b8e 100644 --- a/services/cloud-agent-next/src/session/queue-message.ts +++ b/services/cloud-agent-next/src/session/queue-message.ts @@ -20,6 +20,7 @@ import type { CloudAgentSession } from '../persistence/CloudAgentSession.js'; import type { QueueAckResponse } from '../router/schemas.js'; import { withDORetry } from '../utils/do-retry.js'; import { logger } from '../logger.js'; +import { preflightExistingPromptModel } from './model-preflight.js'; /** Retryable error codes that should map to 503 Service Unavailable. */ const RETRYABLE_CODES: readonly RetryableResultCode[] = [ @@ -95,23 +96,44 @@ export function projectAdmissionToPublicAck( }; } -export async function replayMessageIfAlreadyAdmitted( - input: QueueMessageInput, - ctx: QueueMessageContext -): Promise { +async function hasMessageAdmission(input: QueueMessageInput, ctx: QueueMessageContext) { const messageId = input.turn.id; - if (messageId === undefined || messageId === null) return undefined; + if (messageId === undefined || messageId === null) return false; const sessionId = input.cloudAgentSessionId as SessionId; const doId = ctx.env.CLOUD_AGENT_SESSION.idFromName(`${ctx.userId}:${sessionId}`); - const alreadyAdmitted = await withDORetry, boolean>( + return withDORetry, boolean>( () => ctx.env.CLOUD_AGENT_SESSION.get(doId), stub => stub.hasMessageAdmission(messageId), 'hasMessageAdmission' ); - if (!alreadyAdmitted) return undefined; +} + +export async function preflightAndAdmitPromptMessage( + input: QueueMessageInput, + ctx: QueueMessageContext, + procedure: string, + admit: (input: QueueMessageInput, ctx: QueueMessageContext) => Promise +): Promise { + if (await hasMessageAdmission(input, ctx)) return admit(input, ctx); + + await preflightExistingPromptModel({ + env: ctx.env, + userId: ctx.userId, + cloudAgentSessionId: input.cloudAgentSessionId, + requestedModel: input.agent?.model, + procedure, + }); + + return admit(input, ctx); +} - return queueMessage(input, ctx); +export function preflightAndQueuePromptMessage( + input: QueueMessageInput, + ctx: QueueMessageContext, + procedure: string +): Promise { + return preflightAndAdmitPromptMessage(input, ctx, procedure, queueMessage); } export async function queueMessage( diff --git a/services/cloud-agent-next/src/session/wrapper-global-feed-validation.test.ts b/services/cloud-agent-next/src/session/wrapper-global-feed-validation.test.ts new file mode 100644 index 0000000000..132a3dcd04 --- /dev/null +++ b/services/cloud-agent-next/src/session/wrapper-global-feed-validation.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from 'vitest'; +import type { SessionMetadata } from '../persistence/session-metadata'; +import type { WrapperRuntimeState } from './wrapper-runtime-state'; +import { validateKiloGlobalFeedProducerIdentity } from './wrapper-global-feed-validation'; + +const producer = { + kiloSessionId: 'ses_12345678901234567890123456', + wrapperRunId: 'wr_current', + wrapperGeneration: 3, + wrapperConnectionId: 'conn_current', +}; + +const metadata = { + auth: { kiloSessionId: producer.kiloSessionId }, +} as SessionMetadata; + +const runtimeState = { + wrapperGeneration: producer.wrapperGeneration, + wrapperConnectionId: producer.wrapperConnectionId, + wrapperRunId: producer.wrapperRunId, +} as WrapperRuntimeState; + +describe('validateKiloGlobalFeedProducerIdentity', () => { + it('accepts the current wrapper feed for the metadata root Kilo session', () => { + expect( + validateKiloGlobalFeedProducerIdentity({ + metadata, + runtimeState, + producer, + }) + ).toEqual({ success: true }); + }); + + it('rejects claimed root Kilo session IDs that do not match metadata', () => { + expect( + validateKiloGlobalFeedProducerIdentity({ + metadata, + runtimeState, + producer: { ...producer, kiloSessionId: 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa' }, + }) + ).toEqual({ success: false, status: 403, message: 'Root Kilo session mismatch' }); + }); + + it('rejects stale wrapper runs', () => { + expect( + validateKiloGlobalFeedProducerIdentity({ + metadata, + runtimeState, + producer: { ...producer, wrapperRunId: 'wr_old' }, + }) + ).toEqual({ success: false, status: 409, message: 'Stale wrapper run' }); + }); + + it('rejects stale wrapper generations or connection IDs', () => { + expect( + validateKiloGlobalFeedProducerIdentity({ + metadata, + runtimeState, + producer: { ...producer, wrapperConnectionId: 'conn_old' }, + }) + ).toEqual({ success: false, status: 409, message: 'Stale wrapper connection' }); + }); +}); diff --git a/services/cloud-agent-next/src/session/wrapper-global-feed-validation.ts b/services/cloud-agent-next/src/session/wrapper-global-feed-validation.ts new file mode 100644 index 0000000000..d0f9f4bdb6 --- /dev/null +++ b/services/cloud-agent-next/src/session/wrapper-global-feed-validation.ts @@ -0,0 +1,45 @@ +import * as z from 'zod'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; +import type { WrapperRuntimeState } from './wrapper-runtime-state.js'; + +export const kiloGlobalFeedValidationSchema = z.object({ + kiloSessionId: z.string().min(1), + wrapperRunId: z.string().min(1), + wrapperGeneration: z.number().int().nonnegative(), + wrapperConnectionId: z.string().min(1), +}); + +export type KiloGlobalFeedValidationParams = z.infer; + +export type KiloGlobalFeedValidationResult = + | { success: true } + | { success: false; status: number; message: string }; + +export function validateKiloGlobalFeedProducerIdentity(params: { + metadata: SessionMetadata | null; + runtimeState: WrapperRuntimeState; + producer: KiloGlobalFeedValidationParams; +}): KiloGlobalFeedValidationResult { + const { metadata, runtimeState, producer } = params; + + if (!metadata) { + return { success: false, status: 404, message: 'Session metadata not found' }; + } + + if (metadata.auth.kiloSessionId !== producer.kiloSessionId) { + return { success: false, status: 403, message: 'Root Kilo session mismatch' }; + } + + if (runtimeState.wrapperRunId !== producer.wrapperRunId) { + return { success: false, status: 409, message: 'Stale wrapper run' }; + } + + if ( + runtimeState.wrapperGeneration !== producer.wrapperGeneration || + runtimeState.wrapperConnectionId !== producer.wrapperConnectionId + ) { + return { success: false, status: 409, message: 'Stale wrapper connection' }; + } + + return { success: true }; +} diff --git a/services/cloud-agent-next/src/shared/http-proxy.ts b/services/cloud-agent-next/src/shared/http-proxy.ts new file mode 100644 index 0000000000..58e8ad3a5d --- /dev/null +++ b/services/cloud-agent-next/src/shared/http-proxy.ts @@ -0,0 +1,52 @@ +const BODYLESS_METHODS = new Set(['GET', 'HEAD']); + +const STRIPPED_PROXY_REQUEST_HEADERS = new Set([ + 'authorization', + 'cookie', + 'host', + 'connection', + 'content-length', + 'keep-alive', + 'proxy-authenticate', + 'proxy-authorization', + 'te', + 'trailer', + 'transfer-encoding', + 'upgrade', + 'sec-websocket-accept', + 'sec-websocket-extensions', + 'sec-websocket-key', + 'sec-websocket-protocol', + 'sec-websocket-version', + 'x-kilo-facade-user-id', + 'x-kilo-facade-auth-token', +]); + +export function requestMethodAllowsBody(method: string): boolean { + return !BODYLESS_METHODS.has(method.toUpperCase()); +} + +export function cloneProxyRequestHeaders(headers: Headers): Headers { + const forwarded = new Headers(); + headers.forEach((value, key) => { + if (STRIPPED_PROXY_REQUEST_HEADERS.has(key.toLowerCase())) { + return; + } + forwarded.append(key, value); + }); + return forwarded; +} + +export function createProxyRequest(sourceRequest: Request, targetUrl: string | URL): Request { + const init: RequestInit & { duplex?: 'half' } = { + method: sourceRequest.method, + headers: cloneProxyRequestHeaders(sourceRequest.headers), + }; + + if (sourceRequest.body && requestMethodAllowsBody(sourceRequest.method)) { + init.body = sourceRequest.body; + init.duplex = 'half'; + } + + return new Request(targetUrl, init); +} diff --git a/services/cloud-agent-next/src/types.ts b/services/cloud-agent-next/src/types.ts index a38e5e8290..73fc1b9439 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -1,6 +1,7 @@ import type { getSandbox, ExecutionSession, Sandbox } from '@cloudflare/sandbox'; import type { CloudAgentSession } from './persistence/CloudAgentSession.js'; import type { CloudAgentQueueReport } from '@kilocode/worker-utils/cloud-agent-queue-report'; +import type { UserKiloFacade } from './kilo-facade/user-kilo-facade.js'; import type { CallbackJob } from './callbacks/index.js'; import type { NotificationsBinding } from './notifications-binding.js'; import type { SessionIngestBinding } from './session-ingest-binding.js'; @@ -192,6 +193,8 @@ export type Env = { SandboxDIND: DurableObjectNamespace; /** Durable Object namespace for CloudAgentSession metadata (SQLite-backed) with RPC support */ CLOUD_AGENT_SESSION: DurableObjectNamespace; + /** Durable Object namespace for per-user Kilo SDK facade coordination */ + USER_KILO_FACADE: DurableObjectNamespace; /** Service binding for the session ingest worker */ SESSION_INGEST: SessionIngestBinding; /** Shared secret for internal service-to-service authentication */ diff --git a/services/cloud-agent-next/test/e2e/README.md b/services/cloud-agent-next/test/e2e/README.md index 78796e47a4..71ee95ed72 100644 --- a/services/cloud-agent-next/test/e2e/README.md +++ b/services/cloud-agent-next/test/e2e/README.md @@ -43,7 +43,23 @@ cloud-agent-next refactor. ## Running -Single scenario: +Official SDK basic-chat acceptance (pinned `@kilocode/sdk/v2` `7.2.52`): + +```bash +pnpm --filter cloud-agent-next exec tsx test/e2e/sdk-basic-chat.ts +``` + +This uses a funded ephemeral local user and sends only `Authorization: Bearer ...` +to `/kilo`; prompt mutations therefore pass through real public balance +validation rather than the legacy lifecycle driver's tRPC bypass header. Because +`client.session.create()` is deliberately unsupported by the basic facade, the +driver first materializes one owned root through the existing lifecycle setup, +then proves SDK attach/chat behavior: warm and cold projected reads, cold event +wake-up plus `promptAsync()`, intentional `prompt()` rejection, active `abort()`, +stable warm/cold message pagination, and selector rejection without transcript mutation. +It stops owned sandbox families and releases any fake-LLM gate in cleanup. + +Focused lifecycle scenario: ```bash tsx services/cloud-agent-next/test/e2e/run.ts [--api=unified|legacy] @@ -126,6 +142,15 @@ The fake gateway serves the Kilo routes used in this harness: - `POST /api/openrouter/models/validate` - Worker-side fail-fast model validation. - `POST /api/openrouter/chat/completions` - deterministic streamed completion scenarios. +### SDK coverage boundary + +`sdk-basic-chat.ts` intentionally avoids timing-sensitive assertions already +covered by focused unit or Workers-runtime fixtures: multi-root mapping +ordering and zero-DO list projection, R2 replacement races, private-path +optional fixture variants, and SSE heartbeat/comment parsing. The normal acceptance +scenario asserts that blocking `prompt()` remains intentionally unsupported; +chat admission and wake-up are tested exclusively through `promptAsync()`. + ## Conversation directives A conversation directive is embedded in the user-visible prompt as diff --git a/services/cloud-agent-next/test/e2e/auth.ts b/services/cloud-agent-next/test/e2e/auth.ts index 072ffa95b7..51501f0a8e 100644 --- a/services/cloud-agent-next/test/e2e/auth.ts +++ b/services/cloud-agent-next/test/e2e/auth.ts @@ -19,6 +19,7 @@ import jwt from 'jsonwebtoken'; import { computeDatabaseUrl, createDrizzleClient, kilocode_users, sql } from '@kilocode/db'; export const DRIVER_USER_EMAIL_SUFFIX = '@cloud-agent-next-e2e.example.com'; +export const FUNDED_DRIVER_BALANCE_MICRODOLLARS = 10_000_000; const JWT_TOKEN_VERSION = 3; // --------------------------------------------------------------------------- @@ -100,7 +101,8 @@ export type TestUser = { */ export async function ensureTestUser( databaseUrl: string | undefined, - email: string + email: string, + options?: { funded?: boolean } ): Promise { const resolvedUrl = databaseUrl ?? computeDatabaseUrl(); const driver = createDrizzleClient({ @@ -110,6 +112,12 @@ export async function ensureTestUser( try { const apiTokenPepper = createHash('sha256').update(email).digest('hex').slice(0, 32); const userId = 'usr_e2e_' + createHash('sha256').update(email).digest('hex').slice(0, 16); + const fundedValues = options?.funded + ? { + microdollars_used: 0, + total_microdollars_acquired: FUNDED_DRIVER_BALANCE_MICRODOLLARS, + } + : {}; // Upsert via INSERT ... ON CONFLICT DO UPDATE so we can return the row. const db = driver.db; @@ -123,11 +131,13 @@ export async function ensureTestUser( stripe_customer_id: 'cus_e2e_' + userId, api_token_pepper: apiTokenPepper, is_admin: false, + ...fundedValues, }) .onConflictDoUpdate({ target: kilocode_users.id, set: { api_token_pepper: apiTokenPepper, + ...fundedValues, updated_at: sql`now()`, }, }); diff --git a/services/cloud-agent-next/test/e2e/sandbox-control.ts b/services/cloud-agent-next/test/e2e/sandbox-control.ts index 9143152f38..9671f105ab 100644 --- a/services/cloud-agent-next/test/e2e/sandbox-control.ts +++ b/services/cloud-agent-next/test/e2e/sandbox-control.ts @@ -5,8 +5,9 @@ * containers with synthesized names. The exact naming convention isn't pinned * by this repo, so we match on a stable substring (`Sandbox`) plus the worker * name (`cloud-agent-next-dev`) when present. Lifecycle tests snapshot the - * current set before starting a session when they need to identify that session's - * newly created sandbox. + * current set before starting a session when they need to identify a newly + * created sandbox. Scenarios that may overlap other sandbox creation use the + * wrapper log filename to prove a container belongs to their Cloud Agent root. */ import { execFile } from 'node:child_process'; @@ -14,6 +15,13 @@ import { promisify } from 'node:util'; const execFileAsync = promisify(execFile); +export type DockerCommandExecutor = (args: string[]) => Promise<{ stdout: string }>; + +const executeDockerCommand: DockerCommandExecutor = async args => { + const { stdout } = await execFileAsync('docker', args); + return { stdout }; +}; + export type SandboxContainer = { id: string; name: string; @@ -25,12 +33,10 @@ export type SandboxContainer = { * List running sandbox containers. Returns proxy containers separately so * callers can kill them together with their primary. */ -export async function listSandboxContainers(): Promise { - const { stdout } = await execFileAsync('docker', [ - 'ps', - '--format', - '{{.ID}}\t{{.Names}}\t{{.Image}}', - ]); +export async function listSandboxContainers( + executeDocker: DockerCommandExecutor = executeDockerCommand +): Promise { + const { stdout } = await executeDocker(['ps', '--format', '{{.ID}}\t{{.Names}}\t{{.Image}}']); const result: SandboxContainer[] = []; for (const line of stdout.trim().split('\n')) { if (!line) continue; @@ -54,9 +60,12 @@ export async function listSandboxContainers(): Promise { * Kill a container by ID. Swallows "no such container" errors so callers can * be defensive without try/catch. */ -export async function killContainer(idOrName: string): Promise { +export async function killContainer( + idOrName: string, + executeDocker: DockerCommandExecutor = executeDockerCommand +): Promise { try { - await execFileAsync('docker', ['kill', idOrName]); + await executeDocker(['kill', idOrName]); } catch (err) { const msg = err instanceof Error ? err.message : String(err); if (msg.includes('No such container') || msg.includes('is not running')) return; @@ -79,19 +88,83 @@ export async function waitForNewSandboxPresent( return null; } +async function sandboxHasWrapperLogForAgentSession( + containerId: string, + agentSessionId: string, + executeDocker: DockerCommandExecutor +): Promise { + try { + await executeDocker([ + 'exec', + containerId, + 'sh', + '-c', + 'for log in /tmp/kilocode-wrapper-"$1"-*.log; do test -e "$log" && exit 0; done; exit 1', + 'sandbox-wrapper-log-match', + agentSessionId, + ]); + return true; + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + if (msg.includes('No such container') || msg.includes('is not running')) return false; + if (typeof err === 'object' && err !== null && 'code' in err && err.code === 1) return false; + throw err; + } +} + +/** Return primary sandboxes proven to belong to `agentSessionId` by wrapper log filename. */ +export async function listSandboxesForAgentSession( + agentSessionId: string, + executeDocker: DockerCommandExecutor = executeDockerCommand +): Promise { + const containers = await listSandboxContainers(executeDocker); + const matches: SandboxContainer[] = []; + for (const container of containers) { + if (container.isProxy) continue; + if (await sandboxHasWrapperLogForAgentSession(container.id, agentSessionId, executeDocker)) { + matches.push(container); + } + } + return matches; +} + +/** + * Block until a running primary sandbox proves it belongs to `agentSessionId`. + * Unmatched containers are never returned, even when they appeared recently. + */ +export async function waitForSandboxForAgentSession( + agentSessionId: string, + timeoutMs: number +): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + const [sandbox] = await listSandboxesForAgentSession(agentSessionId); + if (sandbox) return sandbox; + await new Promise(r => setTimeout(r, 500)); + } + return null; +} + +export function sandboxFamilyKey(sandbox: SandboxContainer): string { + return sandbox.isProxy ? sandbox.name.replace(/-proxy$/, '') : sandbox.name; +} + function sandboxFamilyNames(sandbox: SandboxContainer): Set { const primaryName = sandbox.isProxy ? sandbox.name.replace(/-proxy$/, '') : sandbox.name; return new Set([primaryName, `${primaryName}-proxy`]); } /** Kill one sandbox container plus its proxy sibling when present. */ -export async function killSandboxFamily(sandbox: SandboxContainer): Promise { +export async function killSandboxFamily( + sandbox: SandboxContainer, + executeDocker: DockerCommandExecutor = executeDockerCommand +): Promise { const familyNames = sandboxFamilyNames(sandbox); - const containers = await listSandboxContainers(); + const containers = await listSandboxContainers(executeDocker); const killed: string[] = []; for (const container of containers) { if (!familyNames.has(container.name)) continue; - await killContainer(container.id); + await killContainer(container.id, executeDocker); killed.push(container.name); } return killed; diff --git a/services/cloud-agent-next/test/e2e/sdk-basic-chat.ts b/services/cloud-agent-next/test/e2e/sdk-basic-chat.ts new file mode 100644 index 0000000000..0e9e0e719f --- /dev/null +++ b/services/cloud-agent-next/test/e2e/sdk-basic-chat.ts @@ -0,0 +1,577 @@ +import assert from 'node:assert/strict'; +import { randomUUID } from 'node:crypto'; +import path from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { + createKiloClient, + type GlobalEvent, + type Message, + type Part, + type Session, +} from '@kilocode/sdk/v2'; +import { + DRIVER_USER_EMAIL_SUFFIX, + ensureTestUser, + loadDevVars, + loadRepoEnvFiles, + mintApiToken, +} from './auth.js'; +import { + DEFAULT_CONFIG, + openStream, + releaseGate, + startSession, + waitForGateEngaged, + type DriverConfig, +} from './client.js'; +import { createMessageId } from '../../src/session/message-id.js'; +import { + killSandboxFamily, + listSandboxesForAgentSession, + sandboxFamilyKey, + waitForSandboxFamilyGone, + waitForSandboxForAgentSession, + type SandboxContainer, +} from './sandbox-control.js'; + +const SERVICE_PACKAGE_DIR = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '../..'); +const READ_TIMEOUT_MS = 30_000; +const TURN_TIMEOUT_MS = 90_000; +const PUBLIC_DIRECTORY_PREFIX = '/cloud-agent/sessions/'; + +type MessageEntry = { + info: Message; + parts: Part[]; +}; + +type SdkResult = { + data?: T; + error?: unknown; + response: Response; +}; + +function requireData(result: SdkResult, action: string): T { + assert.equal( + result.error, + undefined, + `${action} returned an SDK error: ${JSON.stringify(result.error)}` + ); + if (result.data === undefined) { + throw new Error(`${action} returned no data`); + } + return result.data; +} + +function publicDirectory(sessionID: string): string { + return `${PUBLIC_DIRECTORY_PREFIX}${sessionID}`; +} + +function assertSafeProjection(value: unknown, action: string): void { + const encoded = JSON.stringify(value); + assert.equal(encoded.includes('/workspace/'), false, `${action} leaked /workspace/ path state`); + assert.equal(encoded.includes('/home/agent_'), false, `${action} leaked sandbox home path state`); +} + +function assertProjectedSession(session: Session, sessionID: string, action: string): void { + assert.equal(session.id, sessionID, `${action} returned an unexpected root session`); + assert.equal( + session.directory, + publicDirectory(sessionID), + `${action} returned private directory state` + ); + assert.equal(session.path, undefined, `${action} returned an optional private Session.path`); + assertSafeProjection(session, action); +} + +function textFromParts(parts: Part[]): string { + return parts + .filter(part => part.type === 'text') + .map(part => part.text) + .join(''); +} + +function findAssistantWithText( + messages: MessageEntry[], + expectedText: string +): MessageEntry | undefined { + return messages.find( + entry => entry.info.role === 'assistant' && textFromParts(entry.parts).includes(expectedText) + ); +} + +function assertProjectedAssistant(entry: MessageEntry, directory: string, action: string): void { + assert.equal(entry.info.role, 'assistant', `${action} did not return an assistant message`); + if (entry.info.role !== 'assistant') return; + assert.equal(entry.info.path.cwd, directory, `${action} exposed assistant cwd`); + assert.equal(entry.info.path.root, directory, `${action} exposed assistant root`); + assertSafeProjection(entry, action); +} + +async function pollFor( + action: () => Promise, + label: string, + timeoutMs = TURN_TIMEOUT_MS +): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + const result = await action(); + if (result !== undefined) return result; + await new Promise(resolve => setTimeout(resolve, 200)); + } + throw new Error(`${label} did not become available within ${timeoutMs}ms`); +} + +async function nextEvent( + stream: AsyncGenerator, + label: string, + timeoutMs = READ_TIMEOUT_MS +): Promise { + let timeout: ReturnType | undefined; + try { + const result = await Promise.race([ + stream.next(), + new Promise((_resolve, reject) => { + timeout = setTimeout( + () => reject(new Error(`${label} timed out after ${timeoutMs}ms`)), + timeoutMs + ); + }), + ]); + if (result.done) throw new Error(`${label} closed before yielding an event`); + return result.value; + } finally { + if (timeout !== undefined) clearTimeout(timeout); + } +} + +async function waitForEvent( + stream: AsyncGenerator, + label: string, + predicate: (event: T) => boolean, + timeoutMs = TURN_TIMEOUT_MS +): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + const event = await nextEvent(stream, label, deadline - Date.now()); + if (predicate(event)) return event; + } + throw new Error(`${label} did not yield a matching event within ${timeoutMs}ms`); +} + +function isCloudAgentQueuedExtensionEvent( + event: unknown, + sessionID: string, + messageID: string +): boolean { + if (typeof event !== 'object' || event === null) return false; + const candidate = event as { type?: unknown; properties?: unknown }; + if (candidate.type !== 'cloud.message.queued') return false; + if (typeof candidate.properties !== 'object' || candidate.properties === null) return false; + const properties = candidate.properties as { sessionID?: unknown; messageId?: unknown }; + return properties.sessionID === sessionID && properties.messageId === messageID; +} + +function isCorrelatedAssistantWakeEvent( + event: GlobalEvent['payload'], + sessionID: string, + messageID: string +): boolean { + if (event.type !== 'message.updated') return false; + return ( + event.properties.sessionID === sessionID && + event.properties.info.role === 'assistant' && + event.properties.info.parentID === messageID + ); +} + +function isProjectedGlobalWakeEvent( + event: GlobalEvent, + sessionID: string, + messageID: string +): boolean { + if (event.directory !== publicDirectory(sessionID)) return false; + return isCorrelatedAssistantWakeEvent(event.payload, sessionID, messageID); +} + +function assertProjectedWakeEvent( + event: GlobalEvent['payload'], + sessionID: string, + action: string +): void { + assert.equal(event.type, 'message.updated', `${action} did not carry a projected assistant`); + if (event.type !== 'message.updated') return; + assertProjectedAssistant( + { info: event.properties.info, parts: [] }, + publicDirectory(sessionID), + action + ); +} + +function trackSandboxFamily( + ownedSandboxFamilies: Map, + sandbox: SandboxContainer +): void { + ownedSandboxFamilies.set(sandboxFamilyKey(sandbox), sandbox); +} + +async function discoverOwnedSandboxFamilies( + cloudAgentSessionID: string, + ownedSandboxFamilies: Map +): Promise { + const sandboxes = await listSandboxesForAgentSession(cloudAgentSessionID); + for (const sandbox of sandboxes) trackSandboxFamily(ownedSandboxFamilies, sandbox); +} + +async function stopOwnedSandboxFamilies( + cloudAgentSessionID: string | undefined, + ownedSandboxFamilies: Map +): Promise { + if (cloudAgentSessionID) { + await discoverOwnedSandboxFamilies(cloudAgentSessionID, ownedSandboxFamilies); + } + for (const [key, sandbox] of ownedSandboxFamilies) { + await killSandboxFamily(sandbox); + assert.equal( + await waitForSandboxFamilyGone(sandbox, READ_TIMEOUT_MS), + true, + `sandbox family ${key} did not stop` + ); + ownedSandboxFamilies.delete(key); + } +} + +async function main(): Promise { + loadRepoEnvFiles(SERVICE_PACKAGE_DIR); + const devVars = loadDevVars(SERVICE_PACKAGE_DIR); + const email = `kilo-sdk-basic-chat-${Date.now()}${DRIVER_USER_EMAIL_SUFFIX}`; + const user = await ensureTestUser(process.env.DATABASE_URL, email, { funded: true }); + const config: DriverConfig = { + ...DEFAULT_CONFIG, + user, + nextAuthSecret: devVars.NEXTAUTH_SECRET ?? '', + workerUrl: process.env.WORKER_URL ?? DEFAULT_CONFIG.workerUrl, + gitUrl: process.env.E2E_GIT_URL ?? DEFAULT_CONFIG.gitUrl, + model: process.env.E2E_MODEL ?? DEFAULT_CONFIG.model, + }; + const bearer = mintApiToken(user, config.nextAuthSecret); + const client = createKiloClient({ + baseUrl: `${config.workerUrl.replace(/\/$/, '')}/kilo`, + headers: { Authorization: `Bearer ${bearer}` }, + }); + const ownedSandboxFamilies = new Map(); + const gateTag = `sdk-abort-${randomUUID()}`; + const rejectedDirectoryMessageID = createMessageId(); + const rejectedWorkspaceMessageID = createMessageId(); + let cloudAgentSessionID: string | undefined; + let sessionID: string | undefined; + let directory: string | undefined; + const scopedAbort = new AbortController(); + const globalAbort = new AbortController(); + + try { + const lifecycleStreamPromise = (async () => { + const created = await startSession(config, { prompt: '__fake__:echo:bootstrap' }); + cloudAgentSessionID = created.cloudAgentSessionId; + sessionID = created.kiloSessionId; + directory = publicDirectory(created.kiloSessionId); + const stream = openStream(config, created.cloudAgentSessionId); + const terminal = await stream.waitForTerminal(TURN_TIMEOUT_MS); + stream.close(); + assert.notEqual(terminal, null, 'lifecycle setup turn did not terminalize'); + return created; + })(); + const root = await lifecycleStreamPromise; + const initialSandbox = await waitForSandboxForAgentSession( + root.cloudAgentSessionId, + TURN_TIMEOUT_MS + ); + assert.notEqual(initialSandbox, null, 'lifecycle setup did not start an owned sandbox'); + if (initialSandbox) trackSandboxFamily(ownedSandboxFamilies, initialSandbox); + assert.notEqual(sessionID, undefined); + assert.notEqual(directory, undefined); + if (!sessionID || !directory) throw new Error('lifecycle setup did not provide root identity'); + const rootSessionID = sessionID; + const rootDirectory = directory; + + const listedRoot = await pollFor(async () => { + const result = await client.session.list({ limit: 10 }); + const sessions = requireData(result, 'session.list()'); + return sessions.find(session => session.id === rootSessionID); + }, 'materialized root session list row'); + assertProjectedSession(listedRoot, rootSessionID, 'session.list()'); + + const warmSession = requireData( + await client.session.get({ sessionID: rootSessionID }), + 'warm session.get()' + ); + assertProjectedSession(warmSession, rootSessionID, 'warm session.get()'); + + const warmBootstrapMessages = await pollFor(async () => { + const messages = requireData( + await client.session.messages({ sessionID: rootSessionID, limit: 20 }), + 'warm session.messages()' + ); + return findAssistantWithText(messages, 'bootstrap') ? messages : undefined; + }, 'warm bootstrap transcript'); + const warmBootstrapAssistant = findAssistantWithText(warmBootstrapMessages, 'bootstrap'); + assert.notEqual(warmBootstrapAssistant, undefined); + if (warmBootstrapAssistant) { + assertProjectedAssistant(warmBootstrapAssistant, rootDirectory, 'warm session.messages()'); + } + + await stopOwnedSandboxFamilies(root.cloudAgentSessionId, ownedSandboxFamilies); + + const coldSession = requireData( + await client.session.get({ sessionID: rootSessionID }), + 'cold session.get()' + ); + assertProjectedSession(coldSession, rootSessionID, 'cold session.get()'); + const coldBootstrapMessages = requireData( + await client.session.messages({ sessionID: rootSessionID, limit: 20 }), + 'cold session.messages()' + ); + const coldBootstrapAssistant = findAssistantWithText(coldBootstrapMessages, 'bootstrap'); + assert.notEqual( + coldBootstrapAssistant, + undefined, + 'cold transcript omitted bootstrap response' + ); + if (coldBootstrapAssistant) { + assertProjectedAssistant(coldBootstrapAssistant, rootDirectory, 'cold session.messages()'); + } + + const scopedEvents = await client.event.subscribe( + { directory: rootDirectory }, + { signal: scopedAbort.signal } + ); + const globalEvents = await client.global.event({ signal: globalAbort.signal }); + const scopedConnected = await nextEvent(scopedEvents.stream, 'scoped event connection'); + const globalConnected = await nextEvent(globalEvents.stream, 'global event connection'); + assert.equal(scopedConnected.type, 'server.connected'); + assert.equal(globalConnected.payload.type, 'server.connected'); + assert.equal(globalConnected.directory, '/cloud-agent'); + + const asyncMessageID = createMessageId(); + const scopedQueuedEvent = waitForEvent( + scopedEvents.stream, + 'scoped Cloud Agent queued extension stream', + event => isCloudAgentQueuedExtensionEvent(event, rootSessionID, asyncMessageID) + ); + const globalQueuedEvent = waitForEvent( + globalEvents.stream, + 'global Cloud Agent queued extension stream', + event => + event.directory === rootDirectory && + isCloudAgentQueuedExtensionEvent(event.payload, rootSessionID, asyncMessageID) + ); + const asyncAdmission = await client.session.promptAsync({ + sessionID: rootSessionID, + directory: rootDirectory, + messageID: asyncMessageID, + parts: [{ type: 'text', text: '__fake__:echo:async-result' }], + }); + assert.equal(asyncAdmission.response.status, 204, 'promptAsync() was not accepted'); + assert.equal(asyncAdmission.error, undefined, 'promptAsync() returned an SDK error'); + await Promise.all([scopedQueuedEvent, globalQueuedEvent]); + const scopedWakeEvent = waitForEvent(scopedEvents.stream, 'scoped event wake stream', event => + isCorrelatedAssistantWakeEvent(event, rootSessionID, asyncMessageID) + ); + const globalWakeEvent = waitForEvent(globalEvents.stream, 'global event wake stream', event => + isProjectedGlobalWakeEvent(event, rootSessionID, asyncMessageID) + ); + const [scopedProjectedEvent, globalProjectedEvent] = await Promise.all([ + scopedWakeEvent, + globalWakeEvent, + ]); + assertProjectedWakeEvent(scopedProjectedEvent, rootSessionID, 'scoped wake event'); + assertProjectedWakeEvent(globalProjectedEvent.payload, rootSessionID, 'global wake event'); + assertSafeProjection(scopedProjectedEvent, 'scoped wake event'); + assertSafeProjection(globalProjectedEvent, 'global wake event'); + + const asyncMessages = await pollFor(async () => { + const messages = requireData( + await client.session.messages({ sessionID: rootSessionID, limit: 20 }), + 'promptAsync result messages()' + ); + return findAssistantWithText(messages, 'async-result') ? messages : undefined; + }, 'promptAsync assistant result'); + const asyncAssistant = findAssistantWithText(asyncMessages, 'async-result'); + assert.notEqual(asyncAssistant, undefined); + if (asyncAssistant) { + assert.equal(asyncAssistant.info.role, 'assistant'); + if (asyncAssistant.info.role === 'assistant') { + assert.equal( + asyncAssistant.info.parentID, + asyncMessageID, + 'promptAsync transcript lost correlation' + ); + } + assertProjectedAssistant(asyncAssistant, rootDirectory, 'promptAsync result'); + } + + const unsupportedSyncMessageID = createMessageId(); + const unsupportedSyncPrompt = await client.session.prompt({ + sessionID: rootSessionID, + messageID: unsupportedSyncMessageID, + parts: [{ type: 'text', text: '__fake__:echo:unsupported-sync-result' }], + }); + assert.equal( + unsupportedSyncPrompt.response.status, + 501, + 'blocking prompt() was not rejected as unsupported' + ); + assert.notEqual( + unsupportedSyncPrompt.error, + undefined, + 'blocking prompt() did not return an SDK error' + ); + const afterUnsupportedSync = requireData( + await client.session.messages({ sessionID: rootSessionID, limit: 20 }), + 'messages after unsupported blocking prompt()' + ); + assert.equal( + findAssistantWithText(afterUnsupportedSync, 'unsupported-sync-result'), + undefined, + 'blocking prompt() unexpectedly produced a reply' + ); + assert.equal( + afterUnsupportedSync.some(message => message.info.id === unsupportedSyncMessageID), + false, + 'blocking prompt() unexpectedly admitted its user message' + ); + + const warmFirstPageResult = await client.session.messages({ + sessionID: rootSessionID, + limit: 2, + }); + const warmFirstPage = requireData(warmFirstPageResult, 'warm first messages page'); + const nextCursor = warmFirstPageResult.response.headers.get('x-next-cursor'); + assert.notEqual(nextCursor, null, 'warm first page did not expose a continuation cursor'); + const warmSecondPage = requireData( + await client.session.messages({ + sessionID: rootSessionID, + limit: 2, + before: nextCursor ?? undefined, + }), + 'warm second messages page' + ); + const warmPageIDs = [...warmFirstPage, ...warmSecondPage].map(entry => entry.info.id); + + await stopOwnedSandboxFamilies(root.cloudAgentSessionId, ownedSandboxFamilies); + const coldFirstPageResult = await client.session.messages({ + sessionID: rootSessionID, + limit: 2, + }); + const coldFirstPage = requireData(coldFirstPageResult, 'cold first messages page'); + const coldCursor = coldFirstPageResult.response.headers.get('x-next-cursor'); + assert.equal(coldCursor, nextCursor, 'cold first page returned a different cursor'); + const coldSecondPage = requireData( + await client.session.messages({ + sessionID: rootSessionID, + limit: 2, + before: coldCursor ?? undefined, + }), + 'cold second messages page' + ); + assert.deepEqual( + [...coldFirstPage, ...coldSecondPage].map(entry => entry.info.id), + warmPageIDs, + 'cold pagination returned a different message chronology' + ); + + const matchingSelectorRead = requireData( + await client.session.get({ sessionID: rootSessionID, directory: rootDirectory }), + 'matching directory session.get()' + ); + assertProjectedSession(matchingSelectorRead, rootSessionID, 'matching directory session.get()'); + + const rejectedDirectory = await client.session.promptAsync({ + sessionID: rootSessionID, + directory: `${PUBLIC_DIRECTORY_PREFIX}ses_not_the_owned_root`, + messageID: rejectedDirectoryMessageID, + parts: [{ type: 'text', text: '__fake__:echo:must-not-run' }], + }); + assert.equal( + rejectedDirectory.response.status, + 400, + 'mismatching directory prompt did not fail' + ); + assert.notEqual( + rejectedDirectory.error, + undefined, + 'mismatching directory prompt had no SDK error' + ); + const rejectedWorkspace = await client.session.promptAsync({ + sessionID: rootSessionID, + workspace: 'private-workspace', + messageID: rejectedWorkspaceMessageID, + parts: [{ type: 'text', text: '__fake__:echo:must-not-run' }], + }); + assert.equal(rejectedWorkspace.response.status, 400, 'workspace prompt did not fail'); + assert.notEqual(rejectedWorkspace.error, undefined, 'workspace prompt had no SDK error'); + + const gateMessageID = createMessageId(); + const gateAdmission = await client.session.promptAsync({ + sessionID: rootSessionID, + directory: rootDirectory, + messageID: gateMessageID, + parts: [{ type: 'text', text: `__fake__:gate:${gateTag}` }], + }); + assert.equal(gateAdmission.response.status, 204, 'gated prompt was not accepted'); + assert.equal( + await waitForGateEngaged(config, gateTag, TURN_TIMEOUT_MS), + true, + 'gated prompt did not become active' + ); + const rejectedAbort = await client.session.abort({ + sessionID: rootSessionID, + workspace: 'private-workspace', + }); + assert.equal(rejectedAbort.response.status, 400, 'mismatching workspace abort did not fail'); + assert.equal( + await waitForGateEngaged(config, gateTag, 2_000), + true, + 'mismatching abort interrupted the active turn' + ); + const aborted = requireData( + await client.session.abort({ sessionID: rootSessionID, directory: rootDirectory }), + 'matching directory abort()' + ); + assert.equal(aborted, true, 'abort() did not report successful interruption'); + + const finalMessages = await pollFor(async () => { + const messages = requireData( + await client.session.messages({ sessionID: rootSessionID, limit: 100 }), + 'final session.messages()' + ); + return messages.some(entry => entry.info.id === gateMessageID) ? messages : undefined; + }, 'aborted turn transcript visibility'); + assert.equal( + finalMessages.some(entry => entry.info.id === rejectedDirectoryMessageID), + false, + 'mismatching directory prompt mutated the transcript' + ); + assert.equal( + finalMessages.some(entry => entry.info.id === rejectedWorkspaceMessageID), + false, + 'workspace prompt mutated the transcript' + ); + + console.log(`sdk-basic-chat: passed for ${root.kiloSessionId} (${user.email})`); + } finally { + scopedAbort.abort(); + globalAbort.abort(); + await releaseGate(config.fakeLlmUrl, gateTag).catch(() => {}); + if (sessionID && directory) { + await client.session.abort({ sessionID, directory }).catch(() => undefined); + } + await stopOwnedSandboxFamilies(cloudAgentSessionID, ownedSandboxFamilies).catch(error => { + console.warn(`sdk-basic-chat: sandbox cleanup failed: ${String(error)}`); + }); + } +} + +main().catch(error => { + console.error('sdk-basic-chat failed:', error); + process.exit(1); +}); diff --git a/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts b/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts new file mode 100644 index 0000000000..079ecc2de8 --- /dev/null +++ b/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts @@ -0,0 +1,356 @@ +import { SELF, env, runInDurableObject } from 'cloudflare:test'; +import { describe, expect, it } from 'vitest'; +import { registerReadySession } from '../helpers/session-setup.js'; + +const userId = 'user_facade_runtime'; +const cloudAgentSessionId = 'agent_facade_runtime'; +const kiloSessionId = 'ses_12345678901234567890123456'; +const publicDirectory = `/cloud-agent/sessions/${kiloSessionId}`; + +type WrapperIdentity = { + wrapperRunId: string; + wrapperGeneration: number; + wrapperConnectionId: string; +}; + +type SseDataReader = { + next(): Promise; + cancel(): Promise; +}; + +function createSseDataReader(response: Response): SseDataReader { + if (!response.body) { + throw new Error('Expected an SSE response body'); + } + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + + return { + async next(): Promise { + for (;;) { + const delimiter = buffer.indexOf('\n\n'); + if (delimiter >= 0) { + const frame = buffer.slice(0, delimiter); + buffer = buffer.slice(delimiter + 2); + const data = frame + .split('\n') + .filter(line => line.startsWith('data:')) + .map(line => line.slice('data:'.length).trimStart()) + .join('\n'); + if (data) return JSON.parse(data) as unknown; + continue; + } + const chunk = await reader.read(); + if (chunk.done) throw new Error('Expected another SSE data frame'); + buffer += decoder.decode(chunk.value, { stream: true }); + } + }, + async cancel(): Promise { + await reader.cancel(); + }, + }; +} + +async function configureCurrentProducer(identity: WrapperIdentity): Promise { + const session = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${cloudAgentSessionId}`) + ); + await runInDurableObject(session, async instance => { + const metadata = await instance.getMetadata(); + if (!metadata) { + await registerReadySession(instance, { + sessionId: cloudAgentSessionId, + userId, + kiloSessionId, + prompt: 'runtime facade fixture', + mode: 'code', + model: 'test-model', + kilocodeToken: 'test-token', + }); + } + await instance.ctx.storage.put('wrapper_runtime_state', identity); + }); +} + +async function connectProducer(identity: WrapperIdentity): Promise { + const query = new URLSearchParams({ + userId, + cloudAgentSessionId, + kiloSessionId, + wrapperRunId: identity.wrapperRunId, + wrapperGeneration: String(identity.wrapperGeneration), + wrapperConnectionId: identity.wrapperConnectionId, + }); + const response = await SELF.fetch(`http://worker.test/kilo-global-feed-test?${query}`, { + headers: { Upgrade: 'websocket' }, + }); + if (response.status !== 101 || !response.webSocket) { + throw new Error(`Unexpected producer upgrade response: ${response.status}`); + } + response.webSocket.accept(); + return response.webSocket; +} + +function closeEvent(ws: WebSocket): Promise { + return new Promise(resolve => { + ws.addEventListener('close', event => resolve(event), { once: true }); + }); +} + +async function readProducerMarker(): Promise { + const facade = env.USER_KILO_FACADE.get(env.USER_KILO_FACADE.idFromName(userId)); + return runInDurableObject(facade, async (_instance, state) => + state.storage.get(`producer:${cloudAgentSessionId}`) + ); +} + +describe('UserKiloFacade in the Workers runtime', () => { + it('routes a public facade request to the user Durable Object', async () => { + const response = await SELF.fetch( + `http://worker.test/kilo-test/kilo/global/event?userId=${userId}` + ); + expect(response.status).toBe(200); + expect(response.headers.get('content-type')).toContain('text/event-stream'); + + const stream = createSseDataReader(response); + try { + await expect(stream.next()).resolves.toMatchObject({ + directory: '/cloud-agent', + payload: { type: 'server.connected', properties: {} }, + }); + } finally { + await stream.cancel(); + } + }); + + it('fans redacted Cloud Agent lifecycle extensions out before any wrapper producer is live', async () => { + const extensionUserId = 'user_facade_extension'; + const extensionCloudAgentSessionId = 'agent_facade_extension'; + const extensionKiloSessionId = 'ses_abcdefghijklmnopqrstuvwxyz'; + const session = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${extensionUserId}:${extensionCloudAgentSessionId}`) + ); + await runInDurableObject(session, async instance => { + await registerReadySession(instance, { + sessionId: extensionCloudAgentSessionId, + userId: extensionUserId, + kiloSessionId: extensionKiloSessionId, + prompt: 'extension fixture', + mode: 'code', + model: 'test-model', + kilocodeToken: 'test-token', + }); + }); + const scopedResponse = await SELF.fetch( + `http://worker.test/kilo-scoped-event-test?userId=${extensionUserId}&kiloSessionId=${extensionKiloSessionId}` + ); + const scopedStream = createSseDataReader(scopedResponse); + const messageId = 'msg_018f1e2d3c4bFacadeQueueABC'; + + try { + await expect(scopedStream.next()).resolves.toMatchObject({ type: 'server.connected' }); + await runInDurableObject(session, async instance => { + const emit = (streamEventType: string, payload: unknown) => { + instance.broadcastEvent({ + id: 0 as never, + execution_id: messageId, + session_id: extensionCloudAgentSessionId, + stream_event_type: streamEventType, + payload: JSON.stringify(payload), + timestamp: Date.now(), + }); + }; + emit('cloud.message.queued', { + messageId, + content: 'wake from cold facade', + delivery: 'queued', + }); + emit('cloud.status', { + cloudStatus: { + type: 'preparing', + step: 'kilo_session', + message: 'Restoring /workspace/private/session', + }, + }); + emit('cloud.message.sent', { messageId, delivery: 'sent' }); + emit('cloud.message.completed', { + messageId, + status: 'completed', + delivery: 'sent', + accepted: true, + assistantMessageId: 'msg_private_assistant', + }); + }); + await expect(scopedStream.next()).resolves.toEqual({ + type: 'cloud.message.queued', + properties: { + sessionID: extensionKiloSessionId, + messageId, + delivery: 'queued', + }, + }); + await expect(scopedStream.next()).resolves.toEqual({ + type: 'cloud.status', + properties: { + sessionID: extensionKiloSessionId, + cloudStatus: { type: 'preparing', step: 'kilo_session' }, + }, + }); + await expect(scopedStream.next()).resolves.toEqual({ + type: 'cloud.message.sent', + properties: { sessionID: extensionKiloSessionId, messageId, delivery: 'sent' }, + }); + await expect(scopedStream.next()).resolves.toEqual({ + type: 'cloud.message.completed', + properties: { + sessionID: extensionKiloSessionId, + messageId, + status: 'completed', + delivery: 'sent', + accepted: true, + }, + }); + } finally { + await scopedStream.cancel(); + } + }); + + it('fans a projected producer event out to global and session-scoped public SSE', async () => { + const identity = { + wrapperRunId: 'wr_facade_runtime_1', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_facade_runtime_1', + } satisfies WrapperIdentity; + await configureCurrentProducer(identity); + + const globalResponse = await SELF.fetch( + `http://worker.test/kilo-test/kilo/global/event?userId=${userId}` + ); + const scopedResponse = await SELF.fetch( + `http://worker.test/kilo-scoped-event-test?userId=${userId}&kiloSessionId=${kiloSessionId}` + ); + const globalStream = createSseDataReader(globalResponse); + const scopedStream = createSseDataReader(scopedResponse); + const producer = await connectProducer(identity); + + try { + await expect(globalStream.next()).resolves.toMatchObject({ + payload: { type: 'server.connected' }, + }); + await expect(scopedStream.next()).resolves.toMatchObject({ type: 'server.connected' }); + + producer.send( + JSON.stringify({ + directory: '/workspace/private/session', + payload: { type: 'file.edited', properties: {} }, + }) + ); + producer.send( + JSON.stringify({ + directory: '/workspace/private/session', + payload: { + type: 'session.idle', + properties: { sessionID: 'ses_22222222222222222222222222' }, + }, + }) + ); + producer.send( + JSON.stringify({ + directory: '/workspace/private/session', + payload: { + id: 'evt_runtime_message', + type: 'message.updated', + properties: { + sessionID: kiloSessionId, + info: { + id: 'msg_runtime_assistant', + sessionID: kiloSessionId, + role: 'assistant', + path: { cwd: '/workspace/private/session', root: '/workspace/private' }, + }, + }, + }, + }) + ); + + const globalEvent = await globalStream.next(); + const scopedEvent = await scopedStream.next(); + expect(globalEvent).toEqual({ + directory: publicDirectory, + payload: { + id: 'evt_runtime_message', + type: 'message.updated', + properties: { + sessionID: kiloSessionId, + info: { + id: 'msg_runtime_assistant', + sessionID: kiloSessionId, + role: 'assistant', + path: { cwd: publicDirectory, root: publicDirectory }, + }, + }, + }, + }); + expect(scopedEvent).toEqual({ + id: 'evt_runtime_message', + type: 'message.updated', + properties: { + sessionID: kiloSessionId, + info: { + id: 'msg_runtime_assistant', + sessionID: kiloSessionId, + role: 'assistant', + path: { cwd: publicDirectory, root: publicDirectory }, + }, + }, + }); + // Periodic liveness uses a 10 second comment cadence; the timed comment assertion remains in + // focused facade tests, while runtime data delivery must never introduce a heartbeat variant. + expect(JSON.stringify([globalEvent, scopedEvent])).not.toContain('server.heartbeat'); + } finally { + producer.close(1000, 'test complete'); + await globalStream.cancel(); + await scopedStream.cancel(); + } + }); + + it('replaces an older producer while keeping the replacement marker current', async () => { + const firstIdentity = { + wrapperRunId: 'wr_facade_replaced_1', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_facade_replaced_1', + } satisfies WrapperIdentity; + const nextIdentity = { + wrapperRunId: 'wr_facade_replaced_2', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_facade_replaced_2', + } satisfies WrapperIdentity; + await configureCurrentProducer(firstIdentity); + const first = await connectProducer(firstIdentity); + const replaced = closeEvent(first); + + await configureCurrentProducer(nextIdentity); + const next = await connectProducer(nextIdentity); + const replacementClose = await replaced; + expect(replacementClose.code).toBe(1000); + expect(replacementClose.reason).toBe('Replaced by newer global feed'); + + const staleQuery = new URLSearchParams({ + userId, + cloudAgentSessionId, + kiloSessionId, + wrapperRunId: firstIdentity.wrapperRunId, + wrapperGeneration: String(firstIdentity.wrapperGeneration), + wrapperConnectionId: firstIdentity.wrapperConnectionId, + }); + const staleResponse = await SELF.fetch( + `http://worker.test/kilo-global-feed-test?${staleQuery}`, + { headers: { Upgrade: 'websocket' } } + ); + expect(staleResponse.status).toBe(409); + await expect(readProducerMarker()).resolves.toEqual(nextIdentity); + + next.close(1000, 'test complete'); + }); +}); diff --git a/services/cloud-agent-next/test/test-worker.ts b/services/cloud-agent-next/test/test-worker.ts index de6b9de63a..fc07f4b214 100644 --- a/services/cloud-agent-next/test/test-worker.ts +++ b/services/cloud-agent-next/test/test-worker.ts @@ -4,18 +4,22 @@ * This is a separate worker entry for integration tests that excludes * the Sandbox DO (which requires @cloudflare/containers at runtime). * - * The tests only need CloudAgentSession for WebSocket testing. - * This worker intentionally does NOT import any sandbox-related code - * to avoid the @cloudflare/sandbox import chain. + * The tests expose only the Durable Objects needed for runtime coverage. */ import type { CloudAgentQueueReport } from '@kilocode/worker-utils/cloud-agent-queue-report'; -import { CloudAgentSession as RealCloudAgentSession } from '../src/persistence/CloudAgentSession'; +import type { UserKiloFacade } from '../src/kilo-facade/user-kilo-facade.js'; +import { + KILO_FACADE_AUTH_TOKEN_HEADER, + KILO_FACADE_GLOBAL_FEED_PATH, + KILO_FACADE_USER_ID_HEADER, +} from '../src/kilo-facade/user-kilo-facade.js'; import type { NotificationsBinding, SendCloudAgentSessionNotificationParams, SendCloudAgentSessionNotificationResult, } from '../src/notifications-binding.js'; +import { CloudAgentSession as RealCloudAgentSession } from '../src/persistence/CloudAgentSession'; type RecordedPushCall = SendCloudAgentSessionNotificationParams; @@ -59,17 +63,26 @@ export class CloudAgentSession extends RealCloudAgentSession { } } -// Minimal Env type for tests +export { UserKiloFacade } from '../src/kilo-facade/user-kilo-facade.js'; + type TestEnv = { CLOUD_AGENT_SESSION: DurableObjectNamespace; CLOUD_AGENT_REPORT_QUEUE: Queue; + USER_KILO_FACADE: DurableObjectNamespace; }; +function routeToUserKiloFacade(request: Request, env: TestEnv, userId: string): Promise { + const facade = env.USER_KILO_FACADE.get(env.USER_KILO_FACADE.idFromName(userId)); + const headers = new Headers(request.headers); + headers.set(KILO_FACADE_USER_ID_HEADER, userId); + headers.set(KILO_FACADE_AUTH_TOKEN_HEADER, 'test-token'); + return facade.fetch(new Request(request, { headers })); +} + export default { async fetch(request: Request, env: TestEnv): Promise { const url = new URL(request.url); - // Handle /stream WebSocket endpoint if (url.pathname === '/stream') { const upgradeHeader = request.headers.get('Upgrade'); if (upgradeHeader !== 'websocket') { @@ -103,6 +116,34 @@ export default { return Response.json([...recordedNotificationCalls]); } + if (url.pathname.startsWith('/kilo-test/')) { + const userId = url.searchParams.get('userId') ?? 'test_user'; + const facadeUrl = new URL(request.url); + facadeUrl.pathname = url.pathname.slice('/kilo-test'.length); + facadeUrl.searchParams.delete('userId'); + return routeToUserKiloFacade(new Request(facadeUrl, request), env, userId); + } + + if (url.pathname === '/kilo-scoped-event-test') { + const userId = url.searchParams.get('userId') ?? 'test_user'; + const kiloSessionId = url.searchParams.get('kiloSessionId'); + if (!kiloSessionId) { + return new Response('Missing kiloSessionId parameter', { status: 400 }); + } + const facade = env.USER_KILO_FACADE.get(env.USER_KILO_FACADE.idFromName(userId)); + return facade.openPublicSessionEventStream(kiloSessionId); + } + + if (url.pathname === '/kilo-global-feed-test') { + const userId = url.searchParams.get('userId') ?? 'test_user'; + const facadeUrl = new URL(request.url); + facadeUrl.pathname = KILO_FACADE_GLOBAL_FEED_PATH; + facadeUrl.searchParams.delete('userId'); + facadeUrl.searchParams.set('userId', userId); + const facade = env.USER_KILO_FACADE.get(env.USER_KILO_FACADE.idFromName(userId)); + return facade.fetch(new Request(facadeUrl, request)); + } + return new Response('Not Found', { status: 404 }); }, }; diff --git a/services/cloud-agent-next/test/unit/sandbox-control.test.ts b/services/cloud-agent-next/test/unit/sandbox-control.test.ts new file mode 100644 index 0000000000..3c235c62b0 --- /dev/null +++ b/services/cloud-agent-next/test/unit/sandbox-control.test.ts @@ -0,0 +1,112 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { + killSandboxFamily, + listSandboxesForAgentSession, + type DockerCommandExecutor, + type SandboxContainer, +} from '../e2e/sandbox-control.js'; + +const ownedPrimary: SandboxContainer = { + id: 'owned-primary-id', + name: 'cloud-agent-next-dev-Sandbox-owned', + image: 'cloudflare/sandbox:latest', + isProxy: false, +}; + +const ownedProxy: SandboxContainer = { + id: 'owned-proxy-id', + name: `${ownedPrimary.name}-proxy`, + image: 'cloudflare/sandbox:latest', + isProxy: true, +}; + +const unrelatedPrimary: SandboxContainer = { + id: 'unrelated-primary-id', + name: 'cloud-agent-next-dev-Sandbox-unrelated', + image: 'cloudflare/sandbox:latest', + isProxy: false, +}; + +function dockerPsOutput(containers: SandboxContainer[]): string { + return containers + .map(container => `${container.id}\t${container.name}\t${container.image}`) + .join('\n'); +} + +function createDockerExecutor( + containers: SandboxContainer[], + markerContainerIds: Set = new Set() +): DockerCommandExecutor { + return vi.fn(async args => { + if (args[0] === 'ps') return { stdout: dockerPsOutput(containers) }; + if (args[0] === 'kill') return { stdout: args[1] ?? '' }; + if (args[0] === 'exec' && args[1] && markerContainerIds.has(args[1])) return { stdout: '' }; + if (args[0] === 'exec') throw Object.assign(new Error('wrapper marker not found'), { code: 1 }); + throw new Error(`Unexpected docker command: ${args.join(' ')}`); + }); +} + +describe('listSandboxesForAgentSession', () => { + it('returns only the primary container with a root-correlated wrapper marker', async () => { + const executeDocker = createDockerExecutor( + [ownedPrimary, unrelatedPrimary, ownedProxy], + new Set([ownedPrimary.id]) + ); + + await expect(listSandboxesForAgentSession('agent_owned', executeDocker)).resolves.toEqual([ + ownedPrimary, + ]); + expect(executeDocker).toHaveBeenCalledWith([ + 'exec', + ownedPrimary.id, + 'sh', + '-c', + 'for log in /tmp/kilocode-wrapper-"$1"-*.log; do test -e "$log" && exit 0; done; exit 1', + 'sandbox-wrapper-log-match', + 'agent_owned', + ]); + expect(executeDocker).toHaveBeenCalledWith([ + 'exec', + unrelatedPrimary.id, + 'sh', + '-c', + 'for log in /tmp/kilocode-wrapper-"$1"-*.log; do test -e "$log" && exit 0; done; exit 1', + 'sandbox-wrapper-log-match', + 'agent_owned', + ]); + expect(executeDocker).not.toHaveBeenCalledWith(expect.arrayContaining(['exec', ownedProxy.id])); + }); + + it('returns no family when no primary has a root-correlated wrapper marker', async () => { + const executeDocker = createDockerExecutor([ownedPrimary, unrelatedPrimary, ownedProxy]); + + await expect(listSandboxesForAgentSession('agent_owned', executeDocker)).resolves.toEqual([]); + }); +}); + +describe('killSandboxFamily', () => { + it('kills only the selected family exact primary and proxy containers', async () => { + const similarlyNamedPrimary: SandboxContainer = { + id: 'similarly-named-primary-id', + name: `${ownedPrimary.name}-replacement`, + image: 'cloudflare/sandbox:latest', + isProxy: false, + }; + const executeDocker = createDockerExecutor([ + ownedPrimary, + ownedProxy, + unrelatedPrimary, + similarlyNamedPrimary, + ]); + + await expect(killSandboxFamily(ownedPrimary, executeDocker)).resolves.toEqual([ + ownedPrimary.name, + ownedProxy.name, + ]); + expect(executeDocker).toHaveBeenCalledWith(['kill', ownedPrimary.id]); + expect(executeDocker).toHaveBeenCalledWith(['kill', ownedProxy.id]); + expect(executeDocker).not.toHaveBeenCalledWith(['kill', unrelatedPrimary.id]); + expect(executeDocker).not.toHaveBeenCalledWith(['kill', similarlyNamedPrimary.id]); + }); +}); diff --git a/services/cloud-agent-next/test/unit/wrapper/server.test.ts b/services/cloud-agent-next/test/unit/wrapper/server.test.ts index e6eef59e63..6d559b0935 100644 --- a/services/cloud-agent-next/test/unit/wrapper/server.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/server.test.ts @@ -20,6 +20,8 @@ import { createPromptHandler, createCommandHandler, createSessionReadyHandler, + createKiloProxyHandler, + isKiloProxyPath, bindSessionContext, type ServerConfig, type SessionBinding, @@ -62,6 +64,7 @@ function createMockDeps(state: WrapperState) { readySession: vi.fn(), materializePromptAttachments: vi.fn(async prompt => prompt), configureCommitCoAuthor: vi.fn().mockResolvedValue(undefined), + onSessionBound: vi.fn(), }; } @@ -100,6 +103,112 @@ afterEach(async () => { ); }); +afterEach(() => { + vi.unstubAllGlobals(); +}); + +// --------------------------------------------------------------------------- +// Kilo Proxy Handler +// --------------------------------------------------------------------------- + +describe('createKiloProxyHandler', () => { + it('matches only the wrapper kilo proxy path family', () => { + expect(isKiloProxyPath('/kilo-proxy')).toBe(true); + expect(isKiloProxyPath('/kilo-proxy/session/ses_123/message')).toBe(true); + expect(isKiloProxyPath('/health')).toBe(false); + expect(isKiloProxyPath('/job/status')).toBe(false); + }); + + it('forwards path and query to the private Kilo server', async () => { + const state = new WrapperState(); + const deps = createMockDeps(state); + const fetchMock = vi.fn(async () => new Response('proxied')); + vi.stubGlobal('fetch', fetchMock); + const handler = createKiloProxyHandler(deps); + + await handler( + new Request('http://wrapper/kilo-proxy/session/ses_123/message?cursor=abc', { + method: 'GET', + }) + ); + + const upstreamRequest = fetchMock.mock.calls[0][0] as Request; + const upstreamUrl = new URL(upstreamRequest.url); + expect(upstreamUrl.origin).toBe('http://127.0.0.1:0'); + expect(upstreamUrl.pathname).toBe('/session/ses_123/message'); + expect(upstreamUrl.search).toBe('?cursor=abc'); + }); + + it('preserves mutating methods, JSON bodies, and safe request headers', async () => { + const state = new WrapperState(); + const deps = createMockDeps(state); + const fetchMock = vi.fn(async () => new Response('proxied')); + vi.stubGlobal('fetch', fetchMock); + const handler = createKiloProxyHandler(deps); + + await handler( + new Request('http://wrapper/kilo-proxy/session/ses_123', { + method: 'PATCH', + headers: { + Authorization: 'Bearer should-not-forward', + Cookie: 'session=secret', + 'Content-Type': 'application/json', + Accept: 'application/json', + }, + body: JSON.stringify({ title: 'Updated' }), + }) + ); + + const upstreamRequest = fetchMock.mock.calls[0][0] as Request; + expect(upstreamRequest.method).toBe('PATCH'); + expect(upstreamRequest.headers.get('authorization')).toBeNull(); + expect(upstreamRequest.headers.get('cookie')).toBeNull(); + expect(upstreamRequest.headers.get('content-type')).toBe('application/json'); + expect(upstreamRequest.headers.get('accept')).toBe('application/json'); + await expect(upstreamRequest.text()).resolves.toBe('{"title":"Updated"}'); + }); + + it('passes upstream status, headers, and body through without JSON parsing', async () => { + const state = new WrapperState(); + const deps = createMockDeps(state); + vi.stubGlobal( + 'fetch', + vi.fn( + async () => + new Response('plain upstream error', { + status: 500, + headers: { 'X-Kilo-Upstream': 'yes' }, + }) + ) + ); + const handler = createKiloProxyHandler(deps); + + const response = await handler(new Request('http://wrapper/kilo-proxy/session/ses_123')); + + expect(response.status).toBe(500); + expect(response.headers.get('x-kilo-upstream')).toBe('yes'); + await expect(response.text()).resolves.toBe('plain upstream error'); + }); + + it('returns 503 when the private Kilo runtime is unavailable', async () => { + const state = new WrapperState(); + const deps = createMockDeps(state); + Object.defineProperty(deps.kiloClient, 'serverUrl', { + get() { + throw new Error('Kilo server has not been bootstrapped'); + }, + }); + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + const handler = createKiloProxyHandler(deps); + + const response = await handler(new Request('http://wrapper/kilo-proxy/session/ses_123')); + + expect(response.status).toBe(503); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); + // --------------------------------------------------------------------------- // Answer Permission // --------------------------------------------------------------------------- @@ -369,6 +478,7 @@ describe('bindSessionContext', () => { expect(session!.wrapperRunId).toBe(completeBinding.wrapperRunId); expect(session!.wrapperGeneration).toBe(completeBinding.wrapperGeneration); expect(session!.wrapperConnectionId).toBe(completeBinding.wrapperConnectionId); + expect(deps.onSessionBound).toHaveBeenCalledOnce(); }); it('does not mutate state when validation fails', async () => { @@ -378,6 +488,29 @@ describe('bindSessionContext', () => { await bindSessionContext(bindingWithoutRunId as SessionBinding, defaultServerConfig, deps); expect(state.hasSession).toBe(false); + expect(deps.onSessionBound).not.toHaveBeenCalled(); + }); + + it('notifies only when an existing binding changes', async () => { + const state = new WrapperState(); + const deps = createMockDeps(state); + await bindSessionContext(completeBinding, defaultServerConfig, deps); + deps.onSessionBound.mockClear(); + + await bindSessionContext(completeBinding, defaultServerConfig, deps); + expect(deps.onSessionBound).not.toHaveBeenCalled(); + + await bindSessionContext( + { + ...completeBinding, + wrapperGeneration: completeBinding.wrapperGeneration + 1, + wrapperConnectionId: 'conn_789', + }, + defaultServerConfig, + deps + ); + + expect(deps.onSessionBound).toHaveBeenCalledOnce(); }); }); diff --git a/services/cloud-agent-next/wrangler.jsonc b/services/cloud-agent-next/wrangler.jsonc index 6883ddea91..623d73fed9 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -190,6 +190,10 @@ "class_name": "CloudAgentSession", "name": "CLOUD_AGENT_SESSION", }, + { + "class_name": "UserKiloFacade", + "name": "USER_KILO_FACADE", + }, ], }, "migrations": [ @@ -209,6 +213,10 @@ "new_sqlite_classes": ["SandboxDIND"], "tag": "v4", }, + { + "new_sqlite_classes": ["UserKiloFacade"], + "tag": "v5", + }, ], /** * Hyperdrive Bindings @@ -363,6 +371,10 @@ "class_name": "CloudAgentSession", "name": "CLOUD_AGENT_SESSION", }, + { + "class_name": "UserKiloFacade", + "name": "USER_KILO_FACADE", + }, ], }, "queues": { diff --git a/services/cloud-agent-next/wrangler.test.jsonc b/services/cloud-agent-next/wrangler.test.jsonc index 78ec4f4d20..c5e03db06b 100644 --- a/services/cloud-agent-next/wrangler.test.jsonc +++ b/services/cloud-agent-next/wrangler.test.jsonc @@ -11,6 +11,10 @@ "class_name": "CloudAgentSession", "name": "CLOUD_AGENT_SESSION", }, + { + "class_name": "UserKiloFacade", + "name": "USER_KILO_FACADE", + }, ], }, "migrations": [ @@ -18,6 +22,10 @@ "new_sqlite_classes": ["CloudAgentSession"], "tag": "v2", }, + { + "new_sqlite_classes": ["UserKiloFacade"], + "tag": "v5", + }, ], "queues": { "producers": [ diff --git a/services/cloud-agent-next/wrapper/src/global-feed-manager.test.ts b/services/cloud-agent-next/wrapper/src/global-feed-manager.test.ts new file mode 100644 index 0000000000..dea50f5cb7 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/global-feed-manager.test.ts @@ -0,0 +1,128 @@ +import { describe, expect, it } from 'bun:test'; +import { createGlobalFeedManager } from './global-feed-manager'; +import type { KiloGlobalFeedConnection } from './global-feed'; + +function createDeferredDone() { + let resolve: (() => void) | undefined; + const promise = new Promise(resolvePromise => { + resolve = resolvePromise; + }); + + return { + promise, + resolve: () => resolve?.(), + }; +} + +function createTestManager() { + const calls: string[] = []; + let canOpen = true; + let nextFeedId = 0; + + const manager = createGlobalFeedManager({ + canOpen: () => canOpen, + open: () => { + const feedId = ++nextFeedId; + calls.push(`open:${feedId}`); + return { + close: () => calls.push(`close:${feedId}`), + done: Promise.resolve(), + } satisfies KiloGlobalFeedConnection; + }, + onConnectionError: error => { + throw error; + }, + onOpenError: error => { + throw error; + }, + }); + + return { + calls, + manager, + setCanOpen: (value: boolean) => { + canOpen = value; + }, + }; +} + +describe('global feed readiness orchestration', () => { + it('closes the existing feed immediately and keeps it closed while bootstrap readiness is pending', () => { + const { calls, manager } = createTestManager(); + manager.onRuntimeReady(); + + manager.onSessionBound('close-until-runtime-ready'); + + expect(calls).toEqual(['open:1', 'close:1']); + }); + + it('opens exactly one replacement after runtime readiness succeeds', () => { + const { calls, manager } = createTestManager(); + manager.onRuntimeReady(); + manager.onSessionBound('close-until-runtime-ready'); + + manager.onRuntimeReady(); + + expect(calls).toEqual(['open:1', 'close:1', 'open:2']); + }); + + it('leaves the feed closed when runtime readiness fails', () => { + const { calls, manager } = createTestManager(); + manager.onRuntimeReady(); + manager.onSessionBound('close-until-runtime-ready'); + + expect(calls).toEqual(['open:1', 'close:1']); + }); + + it('immediately replaces the feed for legacy direct bindings', () => { + const { calls, manager } = createTestManager(); + manager.onRuntimeReady(); + + manager.onSessionBound('restart'); + + expect(calls).toEqual(['open:1', 'close:1', 'open:2']); + }); + + it('keeps a replacement feed tracked when the closed feed settles later', async () => { + const calls: string[] = []; + const doneByFeedId = new Map>(); + let nextFeedId = 0; + const manager = createGlobalFeedManager({ + canOpen: () => true, + open: () => { + const feedId = ++nextFeedId; + const done = createDeferredDone(); + doneByFeedId.set(feedId, done); + calls.push(`open:${feedId}`); + return { + close: () => calls.push(`close:${feedId}`), + done: done.promise, + } satisfies KiloGlobalFeedConnection; + }, + onConnectionError: error => { + throw error; + }, + onOpenError: error => { + throw error; + }, + }); + + manager.onRuntimeReady(); + manager.onRuntimeReady(); + doneByFeedId.get(1)?.resolve(); + await Promise.resolve(); + await Promise.resolve(); + manager.onRuntimeReady(); + + expect(calls).toEqual(['open:1', 'close:1', 'open:2', 'close:2', 'open:3']); + }); + + it('does not open a replacement until a Kilo client and session binding are available', () => { + const { calls, manager, setCanOpen } = createTestManager(); + setCanOpen(false); + + manager.onRuntimeReady(); + + expect(calls).toEqual([]); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/global-feed-manager.ts b/services/cloud-agent-next/wrapper/src/global-feed-manager.ts new file mode 100644 index 0000000000..b7dd957edf --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/global-feed-manager.ts @@ -0,0 +1,58 @@ +import type { KiloGlobalFeedConnection } from './global-feed.js'; + +export type SessionBoundFeedPolicy = 'restart' | 'close-until-runtime-ready'; + +export type GlobalFeedManager = { + close(): void; + onRuntimeReady(): void; + onSessionBound(feedPolicy: SessionBoundFeedPolicy): void; +}; + +type GlobalFeedManagerDependencies = { + canOpen(): boolean; + open(): KiloGlobalFeedConnection; + onConnectionError(error: unknown): void; + onOpenError(error: unknown): void; +}; + +export function createGlobalFeedManager(deps: GlobalFeedManagerDependencies): GlobalFeedManager { + let connection: KiloGlobalFeedConnection | undefined; + + function close(): void { + connection?.close(); + connection = undefined; + } + + function restart(): void { + close(); + if (!deps.canOpen()) return; + + try { + const nextConnection = deps.open(); + connection = nextConnection; + void nextConnection.done + .catch(error => { + deps.onConnectionError(error); + }) + .finally(() => { + if (connection === nextConnection) { + connection = undefined; + } + }); + } catch (error) { + deps.onOpenError(error); + } + } + + return { + close, + onRuntimeReady: restart, + onSessionBound: feedPolicy => { + if (feedPolicy === 'close-until-runtime-ready') { + close(); + return; + } + restart(); + }, + }; +} diff --git a/services/cloud-agent-next/wrapper/src/global-feed.test.ts b/services/cloud-agent-next/wrapper/src/global-feed.test.ts new file mode 100644 index 0000000000..9d84adc752 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/global-feed.test.ts @@ -0,0 +1,304 @@ +import { afterEach, describe, expect, it } from 'bun:test'; +import { WrapperState } from './state'; +import type { WrapperKiloClient } from './kilo-api'; +import { + buildKiloGlobalFeedWebSocketUrl, + openKiloGlobalFeed, + parseSseDataStream, +} from './global-feed'; + +type GlobalFeedWebSocketImpl = NonNullable< + Parameters[0]['WebSocketImpl'] +>; + +function streamFromChunks(chunks: string[]): ReadableStream { + const encoder = new TextEncoder(); + return new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(encoder.encode(chunk)); + } + controller.close(); + }, + }); +} + +function asFetch( + fn: (...args: Parameters) => ReturnType +): typeof fetch { + return Object.assign(fn, { preconnect: fetch.preconnect }); +} + +class FakeWebSocket { + static instances: FakeWebSocket[] = []; + static initialBufferedAmount = 0; + + readonly url: string; + readonly options?: { headers?: Record } | string | string[]; + readyState: number = WebSocket.CONNECTING; + bufferedAmount = FakeWebSocket.initialBufferedAmount; + sent: string[] = []; + onopen: ((event: Event) => void) | null = null; + onerror: ((event: Event) => void) | null = null; + onclose: ((event: CloseEvent) => void) | null = null; + + constructor(url: string, options?: { headers?: Record } | string | string[]) { + this.url = url; + this.options = options; + FakeWebSocket.instances.push(this); + queueMicrotask(() => { + this.readyState = WebSocket.OPEN; + this.onopen?.(new Event('open')); + }); + } + + send(data: string): void { + this.sent.push(data); + } + + close(): void { + this.readyState = WebSocket.CLOSED; + this.onclose?.( + new CloseEvent('close', { + code: 1000, + reason: '', + }) + ); + } +} + +function bindGlobalFeedSession(state: WrapperState): void { + state.bindSession({ + kiloSessionId: 'kilo_root_1', + ingestUrl: 'wss://worker.example.com/sessions/user_1/agent_1/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'wr_run_1', + wrapperGeneration: 7, + wrapperConnectionId: 'conn_1', + agentSessionId: 'agent_1', + }); +} + +function makeKiloClient(serverUrl = 'http://127.0.0.1:4321'): WrapperKiloClient { + return { serverUrl } as WrapperKiloClient; +} + +afterEach(() => { + FakeWebSocket.instances = []; + FakeWebSocket.initialBufferedAmount = 0; +}); + +describe('buildKiloGlobalFeedWebSocketUrl', () => { + it('uses the fenced global ingest path and identity query parameters', () => { + const state = new WrapperState(); + bindGlobalFeedSession(state); + + const url = new URL(buildKiloGlobalFeedWebSocketUrl(state.currentSession!)); + + expect(url.protocol).toBe('wss:'); + expect(url.pathname).toBe('/sessions/user_1/agent_1/kilo-global-ingest'); + expect(url.searchParams.get('kiloSessionId')).toBe('kilo_root_1'); + expect(url.searchParams.get('wrapperRunId')).toBe('wr_run_1'); + expect(url.searchParams.get('wrapperGeneration')).toBe('7'); + expect(url.searchParams.get('wrapperConnectionId')).toBe('conn_1'); + }); +}); + +describe('parseSseDataStream', () => { + it('yields data frames across chunks and supports multiline data fields', async () => { + const stream = streamFromChunks([ + 'event: message\ndata: {"one":', + '1}\n\n:data ignored\ndata: line-a\n', + 'data: line-b\n\n', + ]); + + const frames: string[] = []; + for await (const frame of parseSseDataStream(stream)) { + frames.push(frame); + } + + expect(frames).toEqual(['{"one":1}', 'line-a\nline-b']); + }); + + it('recognizes CRLF frame boundaries split after a carriage return', async () => { + const stream = streamFromChunks([ + 'data: line-a\r\n', + 'data: line-b\r', + '\n\r\n', + 'data: {"next":true}\r\n\r', + '\n', + ]); + + const frames: string[] = []; + for await (const frame of parseSseDataStream(stream)) { + frames.push(frame); + } + + expect(frames).toEqual(['line-a\nline-b', '{"next":true}']); + }); + + it('yields a final data frame when the stream ends without a blank-line delimiter', async () => { + const frames: string[] = []; + for await (const frame of parseSseDataStream(streamFromChunks(['data: final\r\n']))) { + frames.push(frame); + } + + expect(frames).toEqual(['final']); + }); +}); + +describe('openKiloGlobalFeed', () => { + it('streams substantive Kilo global events to the fenced worker WebSocket', async () => { + const state = new WrapperState(); + bindGlobalFeedSession(state); + const fetchedUrls: string[] = []; + const fetchImpl = asFetch(async input => { + fetchedUrls.push(input instanceof Request ? input.url : input.toString()); + return new Response( + streamFromChunks([ + 'data: {"payload":{"type":"server.connected","properties":{}}}\n\n', + 'data: not-json\n\n', + 'data: {"directory":"/workspace/root","payload":{"type":"message.updated","properties":{"id":"msg_1"}}}\n\n', + 'data: {"payload":{"type":"server.heartbeat","properties":{}}}\n\n', + 'data: {"directory":"/workspace/child","payload":{"type":"session.idle","properties":{"sessionID":"child"}}}\n\n', + ]), + { status: 200 } + ); + }); + + const connection = openKiloGlobalFeed({ + state, + kiloClient: makeKiloClient(), + fetchImpl, + WebSocketImpl: FakeWebSocket as unknown as GlobalFeedWebSocketImpl, + }); + + await new Promise(resolve => setTimeout(resolve, 0)); + connection.close(); + await connection.done; + + const ws = FakeWebSocket.instances[0]; + expect(fetchedUrls[0]).toBe('http://127.0.0.1:4321/global/event'); + expect(ws.url).toContain('/sessions/user_1/agent_1/kilo-global-ingest'); + expect((ws.options as { headers?: Record }).headers).toEqual({ + Authorization: 'Bearer worker-token', + }); + expect(ws.sent).toEqual([ + JSON.stringify({ + directory: '/workspace/root', + payload: { type: 'message.updated', properties: { id: 'msg_1' } }, + }), + JSON.stringify({ + directory: '/workspace/child', + payload: { type: 'session.idle', properties: { sessionID: 'child' } }, + }), + ]); + }); + + it('forwards substantive events from a CRLF feed split between delimiter bytes', async () => { + const state = new WrapperState(); + bindGlobalFeedSession(state); + const fetchImpl = asFetch(async () => { + return new Response( + streamFromChunks([ + 'data: {"payload":{"type":"server.connected","properties":{}}}\r', + '\n\r', + '\n', + 'data: {"directory":"/workspace/root","payload":{"type":"message.updated","properties":{"id":"msg_crlf"}}}\r', + '\n\r', + '\n', + ]), + { status: 200 } + ); + }); + + const connection = openKiloGlobalFeed({ + state, + kiloClient: makeKiloClient(), + fetchImpl, + WebSocketImpl: FakeWebSocket as unknown as GlobalFeedWebSocketImpl, + retryDelayMs: 60_000, + }); + + await new Promise(resolve => setTimeout(resolve, 0)); + connection.close(); + await connection.done; + + expect(FakeWebSocket.instances[0].sent).toEqual([ + JSON.stringify({ + directory: '/workspace/root', + payload: { type: 'message.updated', properties: { id: 'msg_crlf' } }, + }), + ]); + }); + + it('does not buffer frames when the worker WebSocket is backed up', async () => { + const state = new WrapperState(); + bindGlobalFeedSession(state); + FakeWebSocket.initialBufferedAmount = Number.MAX_SAFE_INTEGER; + const fetchImpl = asFetch(async () => { + return new Response( + streamFromChunks([ + 'data: {"directory":"/workspace/root","payload":{"type":"message.updated","properties":{"id":"msg_1"}}}\n\n', + ]), + { status: 200 } + ); + }); + + const connection = openKiloGlobalFeed({ + state, + kiloClient: makeKiloClient(), + fetchImpl, + WebSocketImpl: FakeWebSocket as unknown as GlobalFeedWebSocketImpl, + retryDelayMs: 60_000, + }); + + await new Promise(resolve => setTimeout(resolve, 0)); + connection.close(); + await connection.done; + + expect(FakeWebSocket.instances[0].sent).toEqual([]); + }); + + it('restarts the global feed after the private Kilo SSE stream ends', async () => { + const state = new WrapperState(); + bindGlobalFeedSession(state); + let fetchCount = 0; + let secondFetchStarted: (() => void) | undefined; + const secondFetch = new Promise(resolve => { + secondFetchStarted = resolve; + }); + const fetchImpl = asFetch(async () => { + fetchCount += 1; + if (fetchCount === 2) { + secondFetchStarted?.(); + } + return new Response( + streamFromChunks([ + `data: {"directory":"/workspace/${fetchCount}","payload":{"type":"message.updated","properties":{"id":"msg_${fetchCount}"}}}\n\n`, + ]), + { status: 200 } + ); + }); + + const connection = openKiloGlobalFeed({ + state, + kiloClient: makeKiloClient(), + fetchImpl, + WebSocketImpl: FakeWebSocket as unknown as GlobalFeedWebSocketImpl, + }); + + await secondFetch; + await new Promise(resolve => setTimeout(resolve, 0)); + connection.close(); + await connection.done; + + expect(fetchCount).toBeGreaterThanOrEqual(2); + expect(FakeWebSocket.instances.flatMap(instance => instance.sent)).toContain( + JSON.stringify({ + directory: '/workspace/2', + payload: { type: 'message.updated', properties: { id: 'msg_2' } }, + }) + ); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/global-feed.ts b/services/cloud-agent-next/wrapper/src/global-feed.ts new file mode 100644 index 0000000000..d7b6e98cc9 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/global-feed.ts @@ -0,0 +1,299 @@ +import type { WrapperState, SessionContext } from './state.js'; +import type { WrapperKiloClient } from './kilo-api.js'; +import { logToFile } from './utils.js'; + +type WebSocketCtor = new ( + url: string, + options?: { headers?: Record } | string | string[] +) => WebSocket; + +export type KiloGlobalFeedConnection = { + close(): void; + done: Promise; +}; + +type OpenKiloGlobalFeedOptions = { + state: WrapperState; + kiloClient: WrapperKiloClient; + fetchImpl?: typeof fetch; + WebSocketImpl?: WebSocketCtor; + retryDelayMs?: number; +}; + +const OPEN_READY_STATE = 1; +const GLOBAL_FEED_RETRY_DELAY_MS = 1_000; +const MAX_GLOBAL_FEED_WEBSOCKET_BUFFERED_BYTES = 1024 * 1024; +const encoder = new TextEncoder(); + +type RequiredGlobalFeedSession = SessionContext & { + agentSessionId: string; + kiloSessionId: string; + ingestUrl: string; + workerAuthToken: string; + wrapperRunId: string; + wrapperGeneration: number; + wrapperConnectionId: string; +}; + +function requireGlobalFeedSession(session: SessionContext | null): RequiredGlobalFeedSession { + if (!session) { + throw new Error('Cannot open Kilo global feed: no session context'); + } + if (!session.kiloSessionId) { + throw new Error('Cannot open Kilo global feed: missing kiloSessionId'); + } + if (!session.ingestUrl) { + throw new Error('Cannot open Kilo global feed: missing ingestUrl'); + } + if (!session.workerAuthToken) { + throw new Error('Cannot open Kilo global feed: missing workerAuthToken'); + } + if (!session.agentSessionId) { + throw new Error('Cannot open Kilo global feed: missing agentSessionId'); + } + if (!session.wrapperRunId) { + throw new Error('Cannot open Kilo global feed: missing wrapperRunId'); + } + if (session.wrapperGeneration === undefined) { + throw new Error('Cannot open Kilo global feed: missing wrapperGeneration'); + } + if (!session.wrapperConnectionId) { + throw new Error('Cannot open Kilo global feed: missing wrapperConnectionId'); + } + return session as RequiredGlobalFeedSession; +} + +export function buildKiloGlobalFeedWebSocketUrl(session: SessionContext): string { + const feedSession = requireGlobalFeedSession(session); + const url = new URL(feedSession.ingestUrl); + if (url.pathname.endsWith('/ingest')) { + url.pathname = `${url.pathname.slice(0, -'/ingest'.length)}/kilo-global-ingest`; + } else { + url.pathname = `${url.pathname.replace(/\/$/, '')}/kilo-global-ingest`; + } + url.searchParams.set('kiloSessionId', feedSession.kiloSessionId); + url.searchParams.set('wrapperRunId', feedSession.wrapperRunId); + url.searchParams.set('wrapperGeneration', String(feedSession.wrapperGeneration)); + url.searchParams.set('wrapperConnectionId', feedSession.wrapperConnectionId); + return url.toString(); +} + +function isSyntheticGlobalEnvelope(value: unknown): boolean { + if (typeof value !== 'object' || value === null || !('payload' in value)) { + return false; + } + const payload = (value as { payload?: unknown }).payload; + if (typeof payload !== 'object' || payload === null || !('type' in payload)) { + return false; + } + const type = (payload as { type?: unknown }).type; + return type === 'server.connected' || type === 'server.heartbeat'; +} + +function parseSseMessageBlock(block: string): string | null { + const dataLines: string[] = []; + for (const line of block.split(/\r\n|\n/)) { + if (line.startsWith('data:')) { + dataLines.push(line.slice('data:'.length).trimStart()); + } + } + return dataLines.length > 0 ? dataLines.join('\n') : null; +} + +function takeNextSseMessageBlock(buffer: string): { block: string; remainder: string } | null { + const lfBoundary = buffer.indexOf('\n\n'); + const crlfBoundary = buffer.indexOf('\r\n\r\n'); + if (lfBoundary === -1 && crlfBoundary === -1) { + return null; + } + + if (crlfBoundary !== -1 && (lfBoundary === -1 || crlfBoundary < lfBoundary)) { + return { + block: buffer.slice(0, crlfBoundary), + remainder: buffer.slice(crlfBoundary + '\r\n\r\n'.length), + }; + } + + return { + block: buffer.slice(0, lfBoundary), + remainder: buffer.slice(lfBoundary + '\n\n'.length), + }; +} + +export async function* parseSseDataStream( + stream: ReadableStream +): AsyncGenerator { + const decoder = new TextDecoder(); + const reader = stream.getReader(); + let buffer = ''; + + while (true) { + const { value, done } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); + + let nextBlock = takeNextSseMessageBlock(buffer); + while (nextBlock !== null) { + buffer = nextBlock.remainder; + const data = parseSseMessageBlock(nextBlock.block); + if (data !== null) { + yield data; + } + nextBlock = takeNextSseMessageBlock(buffer); + } + } + + buffer += decoder.decode(); + const trailing = parseSseMessageBlock(buffer); + if (trailing !== null) { + yield trailing; + } +} + +export function openKiloGlobalFeed(options: OpenKiloGlobalFeedOptions): KiloGlobalFeedConnection { + const session = requireGlobalFeedSession(options.state.currentSession); + const wsUrl = buildKiloGlobalFeedWebSocketUrl(session); + const fetchImpl = options.fetchImpl ?? fetch; + const WebSocketWithHeaders = (options.WebSocketImpl ?? WebSocket) as WebSocketCtor; + const retryDelayMs = options.retryDelayMs ?? GLOBAL_FEED_RETRY_DELAY_MS; + const closeController = new AbortController(); + let closed = false; + let attemptAbortController: AbortController | undefined; + let attemptWebSocket: WebSocket | undefined; + + function waitForRetry(): Promise { + if (closeController.signal.aborted) return Promise.resolve(); + return new Promise(resolve => { + const timer = setTimeout(() => { + closeController.signal.removeEventListener('abort', onAbort); + resolve(); + }, retryDelayMs); + const onAbort = () => { + clearTimeout(timer); + resolve(); + }; + closeController.signal.addEventListener('abort', onAbort, { once: true }); + }); + } + + function closeAttemptSocket(ws: WebSocket | undefined): void { + if (!ws) return; + try { + ws.close(); + } catch { + // Ignore close errors. + } + } + + async function runAttempt(): Promise { + const abortController = new AbortController(); + attemptAbortController = abortController; + const ws = new WebSocketWithHeaders(wsUrl, { + headers: { + Authorization: `Bearer ${session.workerAuthToken}`, + }, + }); + attemptWebSocket = ws; + + await new Promise((resolve, reject) => { + let opened = false; + ws.onopen = () => { + opened = true; + logToFile(`kilo global feed WS connected to: ${wsUrl}`); + resolve(); + }; + ws.onerror = () => { + if (!opened) { + reject(new Error(`Failed to connect Kilo global feed WebSocket: ${wsUrl}`)); + return; + } + abortController.abort(new Error('Kilo global feed WebSocket failed')); + }; + ws.onclose = event => { + if (!closed) { + logToFile( + `kilo global feed WS closed: code=${event.code} reason=${event.reason || '(none)'}` + ); + } + if (!opened) { + reject(new Error(`Kilo global feed WebSocket closed before open: ${wsUrl}`)); + return; + } + abortController.abort(new Error('Kilo global feed WebSocket closed')); + }; + }); + + if (closed) return; + + const globalEventUrl = new URL('/global/event', options.kiloClient.serverUrl); + const response = await fetchImpl(globalEventUrl, { + method: 'GET', + headers: { Accept: 'text/event-stream' }, + signal: abortController.signal, + }); + if (!response.ok || !response.body) { + throw new Error(`Kilo global event stream failed: ${response.status}`); + } + + for await (const data of parseSseDataStream(response.body)) { + if (closed || abortController.signal.aborted) break; + let parsed: unknown; + try { + parsed = JSON.parse(data); + } catch { + logToFile('kilo global feed ignored invalid JSON SSE frame'); + continue; + } + + if (isSyntheticGlobalEnvelope(parsed)) { + continue; + } + + if (ws.readyState === OPEN_READY_STATE) { + const serialized = JSON.stringify(parsed); + const pendingBytes = ws.bufferedAmount + encoder.encode(serialized).byteLength; + if (pendingBytes > MAX_GLOBAL_FEED_WEBSOCKET_BUFFERED_BYTES) { + throw new Error('Kilo global feed WebSocket buffer limit exceeded'); + } + ws.send(serialized); + } + } + } + + const done = (async () => { + while (!closed) { + try { + await runAttempt(); + if (!closed) { + logToFile('kilo global feed stream ended; reconnecting'); + } + } catch (error) { + if (!closed) { + logToFile( + `kilo global feed stopped: ${error instanceof Error ? error.message : String(error)}` + ); + } + } finally { + attemptAbortController?.abort(); + attemptAbortController = undefined; + closeAttemptSocket(attemptWebSocket); + attemptWebSocket = undefined; + } + + if (!closed) { + await waitForRetry(); + } + } + })(); + + return { + done, + close: () => { + closed = true; + closeController.abort(); + attemptAbortController?.abort(); + closeAttemptSocket(attemptWebSocket); + attemptWebSocket = undefined; + }, + }; +} diff --git a/services/cloud-agent-next/wrapper/src/main.ts b/services/cloud-agent-next/wrapper/src/main.ts index 3b14716780..4b1fe63100 100644 --- a/services/cloud-agent-next/wrapper/src/main.ts +++ b/services/cloud-agent-next/wrapper/src/main.ts @@ -19,6 +19,8 @@ import { createWrapperKiloClient, type WrapperKiloClient } from './kilo-api.js'; import { createConnectionManager, openIngestProgressChannel } from './connection.js'; import { createLifecycleManager } from './lifecycle.js'; import { bindSessionContext, createServer } from './server.js'; +import { openKiloGlobalFeed } from './global-feed.js'; +import { createGlobalFeedManager, type SessionBoundFeedPolicy } from './global-feed-manager.js'; import { logToFile } from './utils.js'; import type { WrapperCommand } from '../../src/shared/protocol.js'; import type { @@ -262,6 +264,8 @@ async function main() { readySession: readySession, updateRuntimeEnvironment: updateRuntimeEnvironment, materializePromptAttachments, + onSessionBound: (feedPolicy: SessionBoundFeedPolicy) => + globalFeedManager.onSessionBound(feedPolicy), }; async function verifyExistingKiloSession( @@ -287,6 +291,28 @@ async function main() { } } + const globalFeedManager = createGlobalFeedManager({ + canOpen: () => Boolean(kiloClient && state.currentSession), + open: () => { + if (!kiloClient) { + throw new Error('Cannot open Kilo global feed: no Kilo client'); + } + return openKiloGlobalFeed({ state, kiloClient }); + }, + onConnectionError: error => { + logToFile( + `kilo global feed failed: ${error instanceof Error ? error.message : String(error)}` + ); + }, + onOpenError: error => { + logToFile( + `failed to start kilo global feed: ${ + error instanceof Error ? error.message : String(error) + }` + ); + }, + }); + async function startKiloRuntime( workspacePath: string, expectedSessionId?: string, @@ -306,12 +332,14 @@ async function main() { `startKiloRuntime reused existing runtime without session rebinding sessionId=${kiloSessionId || '(unset)'}` ); } + globalFeedManager.onRuntimeReady(); return; } logToFile( `startKiloRuntime preparing new runtime workspacePath=${workspacePath} previousWorkspacePath=${runtimeWorkspacePath ?? '(unset)'} hadLifecycle=${Boolean(lifecycleManager)} hadConnection=${Boolean(connectionManager)} hadServer=${Boolean(closeKiloServer)}` ); + globalFeedManager.close(); lifecycleManager?.stop(); await connectionManager?.close(); if (closeKiloServer) { @@ -467,6 +495,7 @@ async function main() { } ); lifecycleManager.start(); + globalFeedManager.onRuntimeReady(); } async function updateRuntimeEnvironment(env: Record): Promise { @@ -492,7 +521,12 @@ async function main() { serverConfig.sessionId = request.kiloSessionId; serverConfig.platform = request.materialized.env.KILO_PLATFORM ?? process.env.KILO_PLATFORM; - const bindError = await bindSessionContext(request.session, serverConfig, serverDeps); + const bindError = await bindSessionContext( + request.session, + serverConfig, + serverDeps, + 'close-until-runtime-ready' + ); if (bindError) { const error = (await bindError.json()) as { error?: string; message?: string }; logToFile( @@ -603,6 +637,7 @@ async function main() { // Stop lifecycle timers lifecycleManager?.stop(); + globalFeedManager.close(); // Force exit after timeout setTimeout(() => { diff --git a/services/cloud-agent-next/wrapper/src/restore-session.test.ts b/services/cloud-agent-next/wrapper/src/restore-session.test.ts index b59805f265..b20f972dfa 100644 --- a/services/cloud-agent-next/wrapper/src/restore-session.test.ts +++ b/services/cloud-agent-next/wrapper/src/restore-session.test.ts @@ -56,6 +56,25 @@ function writeMockKilo(binDir: string, exitCode: number): void { fs.writeFileSync(kiloPath, script, { mode: 0o755 }); } +function writeSlowMockKilo(binDir: string): void { + const script = '#!/bin/sh\nsleep 1\nexit 0\n'; + const kiloPath = path.join(binDir, 'kilo'); + fs.writeFileSync(kiloPath, script, { mode: 0o755 }); +} + +function writeSignalTerminatedMockKilo(binDir: string): void { + const script = '#!/bin/sh\nkill -TERM $$\n'; + const kiloPath = path.join(binDir, 'kilo'); + fs.writeFileSync(kiloPath, script, { mode: 0o755 }); +} + +function writeSignalIgnoringDescendantMockKilo(binDir: string, descendantMarker: string): void { + const readyMarker = `${descendantMarker}.ready`; + const script = `#!/bin/sh\ntrap 'exit 0' TERM\nnode -e 'const fs = require("node:fs"); process.on("SIGTERM", () => {}); fs.writeFileSync(process.argv[1], "ready"); setTimeout(() => fs.writeFileSync(process.argv[2], "alive"), 800); setInterval(() => {}, 1000);' "${readyMarker}" "${descendantMarker}" /dev/null 2>&1 &\nwhile [ ! -f "${readyMarker}" ]; do sleep 0.01; done\nsleep 2\n`; + const kiloPath = path.join(binDir, 'kilo'); + fs.writeFileSync(kiloPath, script, { mode: 0o755 }); +} + // --------------------------------------------------------------------------- // Test suite // --------------------------------------------------------------------------- @@ -307,6 +326,54 @@ describe('restoreSession', () => { } }); + it('terminates and returns import error when kilo import exceeds its deadline', async () => { + mockFetchOk(makeSnapshot([])); + writeSlowMockKilo(binDir); + + const result = await restoreSession(SESSION_ID, workspace, undefined, { importTimeoutMs: 50 }); + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.step).toBe('import'); + expect(result.error).toContain('kilo import timed out'); + } + expect(fs.existsSync(TMP_PATH)).toBe(false); + }); + + it('returns import error when kilo import is terminated by a signal', async () => { + mockFetchOk(makeSnapshot([])); + writeSignalTerminatedMockKilo(binDir); + + const result = await restoreSession(SESSION_ID, workspace); + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.step).toBe('import'); + expect(result.error).toContain('kilo import failed'); + } + }); + + it('kills signal-ignoring descendants after a kilo import timeout grace period', async () => { + mockFetchOk(makeSnapshot([])); + const descendantMarker = path.join(tmpDir, 'import-descendant-survived'); + writeSignalIgnoringDescendantMockKilo(binDir, descendantMarker); + const startedAt = Date.now(); + + const result = await restoreSession(SESSION_ID, workspace, undefined, { + importTimeoutMs: 500, + importTerminationGraceMs: 150, + }); + const elapsedMs = Date.now() - startedAt; + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.step).toBe('import'); + expect(result.error).toContain('kilo import timed out'); + } + expect(elapsedMs).toBeGreaterThanOrEqual(600); + expect(fs.existsSync(descendantMarker)).toBe(false); + }); + // ---- Happy paths ---- it('downloads snapshot, imports, and applies diffs', async () => { diff --git a/services/cloud-agent-next/wrapper/src/restore-session.ts b/services/cloud-agent-next/wrapper/src/restore-session.ts index 0dd2abb641..e98bf11d12 100644 --- a/services/cloud-agent-next/wrapper/src/restore-session.ts +++ b/services/cloud-agent-next/wrapper/src/restore-session.ts @@ -1,6 +1,6 @@ import fs from 'node:fs'; import path from 'node:path'; -import { logToFile } from './utils.js'; +import { logToFile, runProcess } from './utils.js'; // --------------------------------------------------------------------------- // Types @@ -21,6 +21,13 @@ type SnapshotDiff = { status: string; }; +export type RestoreSessionOptions = { + importTimeoutMs?: number; + importTerminationGraceMs?: number; +}; + +const KILO_IMPORT_TIMEOUT_MS = 120_000; + // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- @@ -411,10 +418,12 @@ async function extractDiffsWithBun(snapshotPath: string): Promise { const tmpPath = filePath ?? `/tmp/kilo-session-export-${kiloSessionId}.json`; const downloaded = !filePath; + const importTimeoutMs = options.importTimeoutMs ?? KILO_IMPORT_TIMEOUT_MS; log( `starting kiloSessionId=${kiloSessionId} workspace=${workspacePath} input=${downloaded ? 'downloaded' : 'provided'} tmpPath=${tmpPath} home=${process.env.HOME ?? '(unset)'}` @@ -506,23 +515,28 @@ export async function restoreSession( log( `running kilo import kiloSessionId=${kiloSessionId} input=${downloaded ? 'downloaded' : 'provided'} cwd=${workspacePath} home=${process.env.HOME ?? '(unset)'} tmpPath=${tmpPath}` ); - const importProc = Bun.spawn(['kilo', 'import', tmpPath], { - stdout: 'pipe', - stderr: 'pipe', + const importResult = await runProcess('kilo', ['import', tmpPath], { cwd: workspacePath, - env: process.env, + timeoutMs: importTimeoutMs, + terminationGraceMs: options.importTerminationGraceMs, }); - const exitCode = await importProc.exited; const importElapsedMs = Date.now() - importStartedAt; - if (exitCode !== 0) { + if (importResult.terminationReason === 'timeout') { + log( + `kilo import finished outcome=timeout kiloSessionId=${kiloSessionId} input=${downloaded ? 'downloaded' : 'provided'} cwd=${workspacePath} home=${process.env.HOME ?? '(unset)'} elapsedMs=${importElapsedMs} timeoutMs=${importTimeoutMs}` + ); + return fail(`kilo import timed out after ${importTimeoutMs}ms`, null, 'import'); + } + + if (importResult.exitCode !== 0) { log( - `kilo import finished outcome=error exitCode=${exitCode} kiloSessionId=${kiloSessionId} input=${downloaded ? 'downloaded' : 'provided'} cwd=${workspacePath} home=${process.env.HOME ?? '(unset)'} elapsedMs=${importElapsedMs}` + `kilo import finished outcome=error exitCode=${importResult.exitCode} kiloSessionId=${kiloSessionId} input=${downloaded ? 'downloaded' : 'provided'} cwd=${workspacePath} home=${process.env.HOME ?? '(unset)'} elapsedMs=${importElapsedMs}` ); - return fail(`kilo import failed exitCode=${exitCode}`, null, 'import'); + return fail(`kilo import failed exitCode=${importResult.exitCode}`, null, 'import'); } log( - `kilo import finished outcome=ok exitCode=${exitCode} kiloSessionId=${kiloSessionId} input=${downloaded ? 'downloaded' : 'provided'} cwd=${workspacePath} home=${process.env.HOME ?? '(unset)'} elapsedMs=${importElapsedMs}` + `kilo import finished outcome=ok exitCode=${importResult.exitCode} kiloSessionId=${kiloSessionId} input=${downloaded ? 'downloaded' : 'provided'} cwd=${workspacePath} home=${process.env.HOME ?? '(unset)'} elapsedMs=${importElapsedMs}` ); // ---- Step 3: Apply diffs ---- diff --git a/services/cloud-agent-next/wrapper/src/server.test.ts b/services/cloud-agent-next/wrapper/src/server.test.ts index d689429833..94400c5748 100644 --- a/services/cloud-agent-next/wrapper/src/server.test.ts +++ b/services/cloud-agent-next/wrapper/src/server.test.ts @@ -279,7 +279,202 @@ describe('wrapper runtime environment', () => { }); }); +describe('wrapper Kilo proxy route', () => { + it('requests an identity response from private Kilo even when the client accepts gzip', async () => { + const upstreamPort = await getFreePort(); + const wrapperPort = await getFreePort(); + const upstreamAcceptEncodings: Array = []; + const upstream = Bun.serve({ + port: upstreamPort, + fetch(req) { + upstreamAcceptEncodings.push(req.headers.get('accept-encoding')); + return new Response('proxied'); + }, + }); + const kiloClient = { + serverUrl: `http://127.0.0.1:${upstreamPort}`, + } as unknown as WrapperKiloClient; + const wrapper = createServer( + { + port: wrapperPort, + workspacePath: '/workspace/repo', + version: 'test', + sessionId: 'kilo_sess_test', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + userId: 'user_test', + }, + { + state: new WrapperState(), + kiloClient, + openConnection: async () => {}, + closeConnection: async () => {}, + setAborted: () => {}, + resetLifecycle: () => {}, + }, + () => {} + ); + + try { + const response = await fetch(`http://127.0.0.1:${wrapperPort}/kilo-proxy/session/ses_123`, { + headers: { 'Accept-Encoding': 'gzip' }, + }); + + expect(response.status).toBe(200); + expect(upstreamAcceptEncodings).toEqual(['identity']); + } finally { + await wrapper.server.stop(true); + await upstream.stop(true); + } + }); +}); + describe('wrapper session binding', () => { + it('keeps bootstrap rebindings close-only until runtime readiness is verified', async () => { + const state = new WrapperState(); + state.bindSession({ + kiloSessionId: 'kilo_sess_test', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_1', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + }); + state.setConnections({ readyState: WebSocket.OPEN } as WebSocket, new AbortController()); + + const closeOrder: string[] = []; + const response = await bindSessionContext( + { + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_2', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_2', + }, + { + port: 5000, + workspacePath: '/workspace/repo', + version: 'test', + sessionId: 'kilo_sess_test', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + userId: 'user_test', + }, + { + state, + kiloClient: {} as WrapperKiloClient, + openConnection: async () => {}, + closeConnection: async () => { + closeOrder.push('ingest'); + }, + setAborted: () => {}, + resetLifecycle: () => {}, + onSessionBound: feedPolicy => { + closeOrder.push(feedPolicy); + }, + }, + 'close-until-runtime-ready' + ); + + expect(response).toBeNull(); + expect(closeOrder).toEqual(['close-until-runtime-ready', 'ingest']); + }); + + it('closes the bootstrap feed for an unchanged binding until runtime readiness is verified', async () => { + const state = new WrapperState(); + const sessionBinding = { + kiloSessionId: 'kilo_sess_test', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_1', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + }; + state.bindSession(sessionBinding); + + const feedPolicies: string[] = []; + let closeConnectionCalls = 0; + let resetLifecycleCalls = 0; + const response = await bindSessionContext( + sessionBinding, + { + port: 5000, + workspacePath: '/workspace/repo', + version: 'test', + sessionId: 'kilo_sess_test', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + userId: 'user_test', + }, + { + state, + kiloClient: {} as WrapperKiloClient, + openConnection: async () => {}, + closeConnection: async () => { + closeConnectionCalls += 1; + }, + setAborted: () => {}, + resetLifecycle: () => { + resetLifecycleCalls += 1; + }, + onSessionBound: feedPolicy => { + feedPolicies.push(feedPolicy); + }, + }, + 'close-until-runtime-ready' + ); + + expect(response).toBeNull(); + expect(feedPolicies).toEqual(['close-until-runtime-ready']); + expect(closeConnectionCalls).toBe(0); + expect(resetLifecycleCalls).toBe(0); + }); + + it('keeps restart behavior for legacy direct rebindings', async () => { + const state = new WrapperState(); + state.bindSession({ + kiloSessionId: 'kilo_sess_test', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_1', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + }); + + const feedPolicies: string[] = []; + const response = await bindSessionContext( + { + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_2', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_2', + }, + { + port: 5000, + workspacePath: '/workspace/repo', + version: 'test', + sessionId: 'kilo_sess_test', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + userId: 'user_test', + }, + { + state, + kiloClient: {} as WrapperKiloClient, + openConnection: async () => {}, + closeConnection: async () => {}, + setAborted: () => {}, + resetLifecycle: () => {}, + onSessionBound: feedPolicy => { + feedPolicies.push(feedPolicy); + }, + } + ); + + expect(response).toBeNull(); + expect(feedPolicies).toEqual(['restart']); + }); + it('resets lifecycle state when warm rebinding an existing connected session', async () => { const state = new WrapperState(); state.bindSession({ diff --git a/services/cloud-agent-next/wrapper/src/server.ts b/services/cloud-agent-next/wrapper/src/server.ts index 3d27ca9a1c..526ed78942 100644 --- a/services/cloud-agent-next/wrapper/src/server.ts +++ b/services/cloud-agent-next/wrapper/src/server.ts @@ -28,6 +28,8 @@ import { type WrapperSessionReadyRequest, type WrapperSessionReadyResponse, } from '../../src/shared/wrapper-bootstrap.js'; +import { createProxyRequest } from '../../src/shared/http-proxy.js'; +import type { SessionBoundFeedPolicy } from './global-feed-manager.js'; // --------------------------------------------------------------------------- // Types @@ -75,6 +77,8 @@ export type ServerDependencies = { workspacePath: string, commitCoAuthor: WrapperCommitCoAuthor | undefined ) => Promise; + /** Called after a fresh or changed session binding has been stored. */ + onSessionBound?: (feedPolicy: SessionBoundFeedPolicy) => void | Promise; }; export type SessionBinding = { @@ -220,10 +224,25 @@ function parsePtyPath(path: string): { ptyId: string; action?: 'connect' } | nul }; } +async function notifySessionBound( + deps: ServerDependencies, + feedPolicy: SessionBoundFeedPolicy +): Promise { + if (!deps.onSessionBound) return; + try { + await deps.onSessionBound(feedPolicy); + } catch (error) { + logToFile( + `session-bound callback failed: ${error instanceof Error ? error.message : String(error)}` + ); + } +} + export async function bindSessionContext( binding: SessionBinding | undefined, config: ServerConfig, - deps: ServerDependencies + deps: ServerDependencies, + feedPolicy: SessionBoundFeedPolicy = 'restart' ): Promise { const { state } = deps; @@ -293,6 +312,7 @@ export async function bindSessionContext( state.setLogUploader(logUploader); logUploader.start(); logToFile(`session bound: sessionId=${config.sessionId}`); + await notifySessionBound(deps, feedPolicy); return null; } @@ -308,6 +328,9 @@ export async function bindSessionContext( agentSessionId: config.agentSessionId, }; const result = state.bindSession(sessionContext); + if (feedPolicy === 'close-until-runtime-ready') { + await notifySessionBound(deps, feedPolicy); + } if (result.changed) { logToFile( `session binding refreshed: generation=${binding.wrapperGeneration ?? 'none'} connectionId=${binding.wrapperConnectionId ?? 'none'}` @@ -316,6 +339,9 @@ export async function bindSessionContext( await deps.closeConnection(); } deps.resetLifecycle(); + if (feedPolicy === 'restart') { + await notifySessionBound(deps, feedPolicy); + } } return null; } @@ -959,6 +985,36 @@ export function createRuntimeEnvironmentHandler(deps: ServerDependencies) { }; } +export function isKiloProxyPath(path: string): boolean { + return path === '/kilo-proxy' || path.startsWith('/kilo-proxy/'); +} + +export function createKiloProxyHandler(deps: ServerDependencies) { + return async (req: Request): Promise => { + let serverUrl: string; + try { + serverUrl = deps.kiloClient.serverUrl; + } catch { + return errorResponse('KILO_RUNTIME_UNAVAILABLE', 'Kilo runtime is not bootstrapped', 503); + } + + if (!serverUrl) { + return errorResponse('KILO_RUNTIME_UNAVAILABLE', 'Kilo runtime is not bootstrapped', 503); + } + + const url = new URL(req.url); + const upstreamUrl = new URL(serverUrl); + const strippedPath = + url.pathname === '/kilo-proxy' ? '/' : url.pathname.slice('/kilo-proxy'.length); + upstreamUrl.pathname = strippedPath || '/'; + upstreamUrl.search = url.search; + + const upstreamRequest = createProxyRequest(req, upstreamUrl); + upstreamRequest.headers.set('Accept-Encoding', 'identity'); + return fetch(upstreamRequest); + }; +} + // --------------------------------------------------------------------------- // Server Creation // --------------------------------------------------------------------------- @@ -987,6 +1043,7 @@ export function createFetchHandler( const ptyCreateHandler = createPtyCreateHandler(config, deps); const sessionReadyHandler = createSessionReadyHandler(deps); const runtimeEnvironmentHandler = createRuntimeEnvironmentHandler(deps); + const kiloProxyHandler = createKiloProxyHandler(deps); // Route table type RouteHandler = (req: Request) => Response | Promise; @@ -1040,6 +1097,14 @@ export function createFetchHandler( } } + if (isKiloProxyPath(path)) { + return kiloProxyHandler(req).catch(error => { + const msg = error instanceof Error ? error.message : String(error); + logToFile(`kilo proxy error: ${msg}`); + return errorResponse('KILO_PROXY_ERROR', msg, 502); + }); + } + // Look up route const methodRoutes = routes[method]; if (!methodRoutes) { diff --git a/services/cloud-agent-next/wrapper/src/utils.ts b/services/cloud-agent-next/wrapper/src/utils.ts index 90126b92e4..6b6f49af40 100644 --- a/services/cloud-agent-next/wrapper/src/utils.ts +++ b/services/cloud-agent-next/wrapper/src/utils.ts @@ -26,6 +26,7 @@ export type TimeoutAbortOptions = { const EXEC_TIMEOUT_EXIT_CODE = 124; const EXEC_TERMINATION_GRACE_MS = 2_000; +const EXEC_TERMINATION_POLL_MS = 25; const EXEC_TIMEOUT_MESSAGE = 'exec timeout reached'; const EXEC_ABORTED_MESSAGE = 'exec aborted'; @@ -60,6 +61,7 @@ export function runProcess( let settled = false; let terminationReason: TerminationReason | null = null; let terminationTimer: ReturnType | undefined; + let terminationPollTimer: ReturnType | undefined; function abortHandler(): void { terminate('abort'); @@ -68,6 +70,7 @@ export function runProcess( const clearTimers = () => { if (timer) clearTimeout(timer); if (terminationTimer) clearTimeout(terminationTimer); + if (terminationPollTimer) clearTimeout(terminationPollTimer); }; const removeAbortHandler = () => { @@ -107,6 +110,25 @@ export function runProcess( } }; + const processGroupExists = (): boolean => { + if (proc.pid === undefined) return false; + try { + process.kill(-proc.pid, 0); + return true; + } catch { + return false; + } + }; + + const waitForTerminatedGroup = (): void => { + if (settled || terminationReason === null) return; + if (!processGroupExists()) { + resolveTermination(); + return; + } + terminationPollTimer = setTimeout(waitForTerminatedGroup, EXEC_TERMINATION_POLL_MS); + }; + const terminate = (reason: TerminationReason): void => { if (settled || terminationReason !== null) return; terminationReason = reason; @@ -133,16 +155,16 @@ export function runProcess( opts.signal.addEventListener('abort', abortHandler, { once: true }); } } - proc.on('close', code => { + proc.on('close', (code, signal) => { if (settled) return; if (terminationReason !== null) { - resolveTermination(); + waitForTerminatedGroup(); return; } settled = true; clearTimers(); removeAbortHandler(); - resolve({ stdout, stderr, exitCode: code ?? 0 }); + resolve({ stdout, stderr, exitCode: code ?? (signal === null ? 0 : 1) }); }); proc.on('error', err => { if (!settled) { diff --git a/services/session-ingest/package.json b/services/session-ingest/package.json index d5bf317d84..85a27ed894 100644 --- a/services/session-ingest/package.json +++ b/services/session-ingest/package.json @@ -14,6 +14,7 @@ }, "dependencies": { "@kilocode/db": "workspace:*", + "@kilocode/session-ingest-contracts": "workspace:*", "@kilocode/worker-utils": "workspace:*", "@streamparser/json": "0.0.22", "drizzle-orm": "catalog:", diff --git a/services/session-ingest/src/dos/SessionIngestDO.ts b/services/session-ingest/src/dos/SessionIngestDO.ts index df53e76101..6128905999 100644 --- a/services/session-ingest/src/dos/SessionIngestDO.ts +++ b/services/session-ingest/src/dos/SessionIngestDO.ts @@ -1,5 +1,5 @@ import { DurableObject } from 'cloudflare:workers'; -import { eq, ne, gt, and, sql, inArray, isNotNull } from 'drizzle-orm'; +import { eq, ne, gt, gte, lt, and, inArray, isNotNull } from 'drizzle-orm'; import { drizzle, type DrizzleSqliteDODatabase } from 'drizzle-orm/durable-sqlite'; import { migrate } from 'drizzle-orm/durable-sqlite/migrator'; @@ -7,7 +7,7 @@ import { ingestItems, ingestMeta } from '../db/sqlite-schema'; import type { Env } from '../env'; import type { IngestBatch } from '../types/session-sync'; import type { SessionDataItem } from '../types/session-sync'; -import { getItemIdentity } from '../util/compaction'; +import { getItemIdentity, getPartItemIdentityRange } from '../util/compaction'; import { extractNormalizedGitBranchFromItem, extractNormalizedGitUrlFromItem, @@ -24,6 +24,11 @@ import { type TerminationReason, } from './session-metrics'; import migrations from '../../drizzle/migrations'; +import { + readKiloSdkMessages, + readKiloSdkSessionSnapshot, + type KiloSdkSessionSnapshotRead, +} from './kilo-sdk-materialization'; type IngestMetaKey = | ExtractableMetaKey @@ -254,6 +259,14 @@ export class SessionIngestDO extends DurableObject { }; } + async readKiloSdkSessionSnapshot(): Promise { + return readKiloSdkSessionSnapshot(this.db, this.env.SESSION_INGEST_R2); + } + + async readKiloSdkMessages(params: { limit?: number; before?: string }) { + return readKiloSdkMessages(this.db, this.env.SESSION_INGEST_R2, params); + } + async getAllStream(): Promise> { const db = this.db; const r2 = this.env.SESSION_INGEST_R2; @@ -312,8 +325,7 @@ export class SessionIngestDO extends DurableObject { // parts for this message: item_id = '{msgId}/{partId}' const msgId = msgRow.item_id.slice('message/'.length); - // Escape LIKE wildcards (% and _) so they match literally, with ESCAPE clause - const likePattern = msgId.replace(/[%_\\]/g, '\\$&') + '/%'; + const partRange = getPartItemIdentityRange(msgId); controller.enqueue(encoder.encode(',"parts":[')); let partCursor = 0; let firstPart = true; @@ -329,7 +341,8 @@ export class SessionIngestDO extends DurableObject { .where( and( eq(ingestItems.item_type, 'part'), - sql`${ingestItems.item_id} LIKE ${likePattern} ESCAPE '\\'`, + gte(ingestItems.item_id, partRange.start), + lt(ingestItems.item_id, partRange.end), gt(ingestItems.id, partCursor) ) ) diff --git a/services/session-ingest/src/dos/kilo-sdk-materialization.ts b/services/session-ingest/src/dos/kilo-sdk-materialization.ts new file mode 100644 index 0000000000..b34f5f11e5 --- /dev/null +++ b/services/session-ingest/src/dos/kilo-sdk-materialization.ts @@ -0,0 +1,829 @@ +import { and, desc, eq, gt, gte, lt, sql } from 'drizzle-orm'; +import type { DrizzleSqliteDODatabase } from 'drizzle-orm/durable-sqlite'; +import { z } from 'zod'; + +import { + decodeKiloSdkMessagesCursor, + encodeKiloSdkMessagesCursor, + MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE, + messageIdSchema, + partIdSchema, + type KiloSdkMessagesCursor, + type KiloSdkMessagesLegacyCursor, +} from '@kilocode/session-ingest-contracts'; + +import { ingestItems } from '../db/sqlite-schema'; +import { getPartItemIdentityRange } from '../util/compaction'; + +export const MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES = 8 * 1024 * 1024; +export const MAX_KILO_SDK_SESSION_SNAPSHOT_BYTES = 8 * 1024 * 1024; +export const KILO_SDK_HISTORY_CANDIDATE_OVERHEAD_BYTES = 256; +const KILO_SDK_HISTORY_ENUMERATION_BATCH_SIZE = 64; +// Bound cold positive-limit scans to two SQLite batches before failing unsafe continuation. +export const KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP = + 2 * KILO_SDK_HISTORY_ENUMERATION_BATCH_SIZE; + +type KiloSdkHistoryReadPhase = 'message_scan' | 'page_parts'; + +type KiloSdkHistoryTooLarge = { + kind: 'too_large'; + maximumBytes: number; + phase: KiloSdkHistoryReadPhase; +}; + +type KiloSdkHistoryRetryableFailure = { + kind: 'retryable_failure'; + phase: KiloSdkHistoryReadPhase; +}; + +type KiloSdkInvalidData = { kind: 'invalid_data' }; +type KiloSdkIntrinsicallyTooLarge = { kind: 'intrinsically_too_large' }; +type KiloSdkHistoryReadFailure = + | KiloSdkHistoryTooLarge + | KiloSdkHistoryRetryableFailure + | KiloSdkInvalidData; + +export type KiloSdkSessionSnapshotRead = + | { kind: 'pending' } + | { kind: 'value'; info: Record; byteLength: number } + | { kind: 'too_large'; maximumBytes: number } + | { kind: 'retryable_failure' } + | KiloSdkInvalidData; + +type KiloSdkMessagesRead = + | { + messages: Array<{ info: Record; parts: Record[] }>; + nextCursor: string | null; + omittedItemCount: number; + } + | KiloSdkHistoryReadFailure + | null; + +type KiloSdkHistoryReadBudget = { + maximumBytes: number; + consumedBytes: number; +}; + +export function createKiloSdkHistoryReadBudget( + maximumBytes = MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES +): KiloSdkHistoryReadBudget { + return { maximumBytes, consumedBytes: 0 }; +} + +type ItemDataRef = Pick; + +type MaterializedKiloSdkMessage = { + info: Record; + identity: KiloSdkMessagesLegacyCursor; +}; + +type PersistedKiloSdkMessageStorageIdentity = { + itemId: string; + messageId: string; +}; + +type KiloSdkMessageInfoEnumeration = { + pageNewestFirst: MaterializedKiloSdkMessage[]; + nextCursor: string | null; + omittedItemCount: number; +}; + +type KiloSdkOversizedItemPolicy = 'skip' | 'fail'; + +type KiloSdkHistoryCandidateRead = + | { kind: 'value'; value: Record } + | { kind: 'missing' } + | KiloSdkIntrinsicallyTooLarge + | KiloSdkHistoryReadFailure; + +type KiloSdkHistoryCandidateOutcome = + | { kind: 'skip' } + | { kind: 'value'; value: Record } + | KiloSdkHistoryReadFailure; + +export async function readKiloSdkSessionSnapshot( + db: DrizzleSqliteDODatabase, + r2: R2Bucket +): Promise { + const resolveSession = () => + db + .select({ + item_data: ingestItems.item_data, + item_data_r2_key: ingestItems.item_data_r2_key, + }) + .from(ingestItems) + .where(eq(ingestItems.item_type, 'session')) + .limit(1) + .get(); + const sessionRow = resolveSession(); + + if (!sessionRow) { + return { kind: 'pending' }; + } + + return readKiloSdkSessionItem( + sessionRow, + resolveSession, + r2, + MAX_KILO_SDK_SESSION_SNAPSHOT_BYTES + ); +} + +export async function readKiloSdkMessages( + db: DrizzleSqliteDODatabase, + r2: R2Bucket, + params: { limit?: number; before?: string } +): Promise { + const budget = createKiloSdkHistoryReadBudget(); + const before = + params.before === undefined ? undefined : decodeKiloSdkMessagesCursor(params.before); + const requestedLimit = params.limit ?? 0; + const limit = + requestedLimit > 0 + ? Math.min(requestedLimit, MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE) + : requestedLimit; + if (limit <= 0 && before !== undefined) { + return { kind: 'invalid_data' }; + } + const enumerated = + limit > 0 + ? await enumerateBoundedKiloSdkMessageInfos(db, r2, budget, limit, before) + : await enumerateUnboundedKiloSdkMessageInfos(db, r2, budget); + if (isKiloSdkHistoryReadFailure(enumerated)) { + return enumerated; + } + + if (enumerated.pageNewestFirst.length === 0) { + if (enumerated.omittedItemCount > 0 || enumerated.nextCursor) { + return { + messages: [], + nextCursor: enumerated.nextCursor, + omittedItemCount: enumerated.omittedItemCount, + }; + } + const sessionRow = db + .select({ id: ingestItems.id }) + .from(ingestItems) + .where(eq(ingestItems.item_type, 'session')) + .limit(1) + .get(); + const messageRow = db + .select({ id: ingestItems.id }) + .from(ingestItems) + .where(eq(ingestItems.item_type, 'message')) + .limit(1) + .get(); + return sessionRow || messageRow + ? { messages: [], nextCursor: null, omittedItemCount: 0 } + : null; + } + const storedMessagesNewestFirst: Array<{ + info: Record; + parts: Record[]; + }> = []; + let omittedItemCount = enumerated.omittedItemCount; + let partHydrationStopped = false; + + for (const message of enumerated.pageNewestFirst) { + if (partHydrationStopped) { + const omittedParts = countKiloSdkMessagePartRows(db, message.identity); + if (omittedParts.kind === 'invalid_data') { + return omittedParts; + } + omittedItemCount += omittedParts.count; + storedMessagesNewestFirst.push({ info: message.info, parts: [] }); + continue; + } + const hydratedParts = await hydrateKiloSdkMessageParts( + db, + r2, + budget, + message.identity, + limit > 0 ? 'skip' : 'fail' + ); + if (isKiloSdkHistoryReadFailure(hydratedParts)) { + return hydratedParts; + } + omittedItemCount += hydratedParts.omittedItemCount; + partHydrationStopped = hydratedParts.aggregateBudgetExhausted; + storedMessagesNewestFirst.push({ info: message.info, parts: hydratedParts.parts }); + } + + return { + messages: storedMessagesNewestFirst.reverse(), + nextCursor: enumerated.nextCursor, + omittedItemCount, + }; +} + +function messageItemId(messageId: string): string { + return `message/${messageId}`; +} + +function parsePersistedKiloSdkMessageStorageIdentity( + itemId: string +): PersistedKiloSdkMessageStorageIdentity | null { + const prefix = 'message/'; + if (!itemId.startsWith(prefix)) return null; + const parsed = messageIdSchema.safeParse(itemId.slice(prefix.length)); + return parsed.success && !parsed.data.includes('/') ? { itemId, messageId: parsed.data } : null; +} + +function readItemReference(db: DrizzleSqliteDODatabase, rowId: number): ItemDataRef | undefined { + return db + .select({ item_data: ingestItems.item_data, item_data_r2_key: ingestItems.item_data_r2_key }) + .from(ingestItems) + .where(eq(ingestItems.id, rowId)) + .get(); +} + +function isKiloSdkHistoryReadFailure( + result: KiloSdkHistoryReadFailure | object +): result is KiloSdkHistoryReadFailure { + return 'kind' in result; +} + +function resolveKiloSdkHistoryCandidateOutcome( + materialized: KiloSdkHistoryCandidateRead, + oversizedItemPolicy: KiloSdkOversizedItemPolicy, + phase: KiloSdkHistoryReadPhase +): KiloSdkHistoryCandidateOutcome { + if (materialized.kind === 'missing') { + return { kind: 'skip' }; + } + if (materialized.kind === 'intrinsically_too_large') { + return oversizedItemPolicy === 'skip' + ? { kind: 'skip' } + : { kind: 'too_large', maximumBytes: MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, phase }; + } + return materialized; +} + +function hasOlderKiloSdkMessageRows( + db: DrizzleSqliteDODatabase, + beforeMessageItemId: string +): boolean { + return Boolean( + db + .select({ id: ingestItems.id }) + .from(ingestItems) + .where( + and(eq(ingestItems.item_type, 'message'), lt(ingestItems.item_id, beforeMessageItemId)) + ) + .limit(1) + .get() + ); +} + +function finishBoundedKiloSdkMessageInfoEnumeration( + db: DrizzleSqliteDODatabase, + pageNewestFirst: MaterializedKiloSdkMessage[], + omittedItemCount: number, + lastConsumedMessageStorageIdentity: PersistedKiloSdkMessageStorageIdentity | undefined +): KiloSdkMessageInfoEnumeration | KiloSdkHistoryTooLarge { + if ( + !lastConsumedMessageStorageIdentity || + !hasOlderKiloSdkMessageRows(db, lastConsumedMessageStorageIdentity.itemId) + ) { + return { pageNewestFirst, nextCursor: null, omittedItemCount }; + } + const oldestReturnedMessage = pageNewestFirst[pageNewestFirst.length - 1]; + if ( + !oldestReturnedMessage || + messageItemId(oldestReturnedMessage.identity.id) !== lastConsumedMessageStorageIdentity.itemId + ) { + return { + kind: 'too_large', + maximumBytes: MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, + phase: 'message_scan', + }; + } + return { + pageNewestFirst, + nextCursor: encodeKiloSdkMessagesCursor(oldestReturnedMessage.identity), + omittedItemCount, + }; +} + +async function enumerateBoundedKiloSdkMessageInfos( + db: DrizzleSqliteDODatabase, + r2: R2Bucket, + budget: KiloSdkHistoryReadBudget, + limit: number, + before: KiloSdkMessagesCursor | undefined +): Promise { + const pageNewestFirst: MaterializedKiloSdkMessage[] = []; + let omittedItemCount = 0; + let consumedRowCount = 0; + let lastConsumedMessageStorageIdentity: PersistedKiloSdkMessageStorageIdentity | undefined; + let scanBeforeItemId = before ? messageItemId(before.id) : undefined; + while ( + pageNewestFirst.length < limit && + consumedRowCount < KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP + ) { + const pageRows = db + .select({ id: ingestItems.id, item_id: ingestItems.item_id }) + .from(ingestItems) + .where( + and( + eq(ingestItems.item_type, 'message'), + scanBeforeItemId ? lt(ingestItems.item_id, scanBeforeItemId) : undefined + ) + ) + .orderBy(desc(ingestItems.item_id)) + .limit( + Math.min( + KILO_SDK_HISTORY_ENUMERATION_BATCH_SIZE, + KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP - consumedRowCount + ) + ) + .all(); + if (pageRows.length === 0) break; + for (const row of pageRows) { + const storageIdentity = parsePersistedKiloSdkMessageStorageIdentity(row.item_id); + if (!storageIdentity) { + return { kind: 'invalid_data' }; + } + const materialized = await readKiloSdkHistoryCandidate( + row.id, + rowId => readItemReference(db, rowId), + r2, + budget, + 'message_scan' + ); + if (materialized.kind === 'too_large') { + if (pageNewestFirst.length === 0) { + return materialized; + } + return finishBoundedKiloSdkMessageInfoEnumeration( + db, + pageNewestFirst, + omittedItemCount, + lastConsumedMessageStorageIdentity + ); + } else { + const outcome = resolveKiloSdkHistoryCandidateOutcome(materialized, 'skip', 'message_scan'); + if (outcome.kind === 'skip') { + if (materialized.kind === 'intrinsically_too_large') { + const omittedRows = countOmittedKiloSdkMessageRows(db, row.item_id); + if (omittedRows.kind === 'invalid_data') { + return omittedRows; + } + omittedItemCount += omittedRows.count; + } else { + omittedItemCount += 1; + } + } else if (outcome.kind !== 'value') { + return outcome; + } else { + const identity = readMessageIdentity(outcome.value); + if (!identity || identity.id !== storageIdentity.messageId) { + return { kind: 'invalid_data' }; + } + pageNewestFirst.push({ info: outcome.value, identity }); + } + } + lastConsumedMessageStorageIdentity = storageIdentity; + scanBeforeItemId = storageIdentity.itemId; + consumedRowCount += 1; + if (pageNewestFirst.length === limit) break; + } + if ( + pageNewestFirst.length === limit || + pageRows.length < KILO_SDK_HISTORY_ENUMERATION_BATCH_SIZE + ) { + break; + } + } + return finishBoundedKiloSdkMessageInfoEnumeration( + db, + pageNewestFirst, + omittedItemCount, + lastConsumedMessageStorageIdentity + ); +} + +async function enumerateUnboundedKiloSdkMessageInfos( + db: DrizzleSqliteDODatabase, + r2: R2Bucket, + budget: KiloSdkHistoryReadBudget +): Promise { + const messages: MaterializedKiloSdkMessage[] = []; + let lastMessageRowId = 0; + for (;;) { + const messageRowIds = db + .select({ id: ingestItems.id, item_id: ingestItems.item_id }) + .from(ingestItems) + .where(and(eq(ingestItems.item_type, 'message'), gt(ingestItems.id, lastMessageRowId))) + .orderBy(ingestItems.id) + .limit(KILO_SDK_HISTORY_ENUMERATION_BATCH_SIZE) + .all(); + if (messageRowIds.length === 0) break; + for (const messageRowId of messageRowIds) { + const storageIdentity = parsePersistedKiloSdkMessageStorageIdentity(messageRowId.item_id); + if (!storageIdentity) { + return { kind: 'invalid_data' }; + } + const materialized = await readKiloSdkHistoryCandidate( + messageRowId.id, + rowId => readItemReference(db, rowId), + r2, + budget, + 'message_scan' + ); + const outcome = resolveKiloSdkHistoryCandidateOutcome(materialized, 'fail', 'message_scan'); + if (outcome.kind !== 'skip' && outcome.kind !== 'value') { + return outcome; + } + if (outcome.kind === 'value') { + const identity = readMessageIdentity(outcome.value); + if (!identity || identity.id !== storageIdentity.messageId) { + return { kind: 'invalid_data' }; + } + messages.push({ info: outcome.value, identity }); + } + lastMessageRowId = messageRowId.id; + } + if (messageRowIds.length < KILO_SDK_HISTORY_ENUMERATION_BATCH_SIZE) break; + } + messages.sort(compareKiloSdkMessageInfo).reverse(); + return { pageNewestFirst: messages, nextCursor: null, omittedItemCount: 0 }; +} + +type KiloSdkMessagePartRowCount = { kind: 'count'; count: number } | KiloSdkInvalidData; + +type PersistedKiloSdkPartStorageIdentity = { + messageId: string; + partId: string; +}; + +function parsePersistedKiloSdkPartStorageIdentity( + itemId: string, + expectedMessageId: string +): PersistedKiloSdkPartStorageIdentity | null { + const segments = itemId.split('/'); + if (segments.length !== 2) return null; + const [messageId, partId] = segments; + const parsed = z.object({ messageId: messageIdSchema, partId: partIdSchema }).safeParse({ + messageId, + partId, + }); + return parsed.success && parsed.data.messageId === expectedMessageId ? parsed.data : null; +} + +function countKiloSdkMessagePartRowsByMessageId( + db: DrizzleSqliteDODatabase, + messageId: string, + afterRowId = 0 +): KiloSdkMessagePartRowCount { + if (!messageIdSchema.safeParse(messageId).success) return { kind: 'invalid_data' }; + const partRange = getPartItemIdentityRange(messageId); + const broadCount = + db + .select({ count: sql`count(*)` }) + .from(ingestItems) + .where( + and( + eq(ingestItems.item_type, 'part'), + gte(ingestItems.item_id, partRange.start), + lt(ingestItems.item_id, partRange.end), + gt(ingestItems.id, afterRowId) + ) + ) + .get()?.count ?? 0; + const directCount = + db + .select({ count: sql`count(*)` }) + .from(ingestItems) + .where( + and( + eq(ingestItems.item_type, 'part'), + gte(ingestItems.item_id, partRange.start), + lt(ingestItems.item_id, partRange.end), + sql`substr(${ingestItems.item_id}, length(${partRange.start}) + 1, 3) = 'prt'`, + sql`instr(substr(${ingestItems.item_id}, length(${partRange.start}) + 1), '/') = 0`, + gt(ingestItems.id, afterRowId) + ) + ) + .get()?.count ?? 0; + const nulCount = + db + .select({ count: sql`count(*)` }) + .from(ingestItems) + .where( + and( + eq(ingestItems.item_type, 'part'), + gte(ingestItems.item_id, partRange.start), + lt(ingestItems.item_id, partRange.end), + sql`instr(CAST(${ingestItems.item_id} AS BLOB), X'00') > 0`, + gt(ingestItems.id, afterRowId) + ) + ) + .get()?.count ?? 0; + return broadCount === directCount && nulCount === 0 + ? { kind: 'count', count: directCount } + : { kind: 'invalid_data' }; +} + +function countKiloSdkMessagePartRows( + db: DrizzleSqliteDODatabase, + identity: KiloSdkMessagesLegacyCursor, + afterRowId = 0 +): KiloSdkMessagePartRowCount { + return countKiloSdkMessagePartRowsByMessageId(db, identity.id, afterRowId); +} + +function countOmittedKiloSdkMessageRows( + db: DrizzleSqliteDODatabase, + messageItemId: string +): KiloSdkMessagePartRowCount { + const storageIdentity = parsePersistedKiloSdkMessageStorageIdentity(messageItemId); + if (!storageIdentity) return { kind: 'invalid_data' }; + const parts = countKiloSdkMessagePartRowsByMessageId(db, storageIdentity.messageId); + return parts.kind === 'invalid_data' ? parts : { kind: 'count', count: 1 + parts.count }; +} + +async function hydrateKiloSdkMessageParts( + db: DrizzleSqliteDODatabase, + r2: R2Bucket, + budget: KiloSdkHistoryReadBudget, + identity: KiloSdkMessagesLegacyCursor, + oversizedItemPolicy: KiloSdkOversizedItemPolicy +): Promise< + | { + parts: Array>; + omittedItemCount: number; + aggregateBudgetExhausted: boolean; + } + | KiloSdkHistoryReadFailure +> { + const partRange = getPartItemIdentityRange(identity.id); + const parts: Array> = []; + let omittedItemCount = 0; + let lastPartRowId = 0; + for (;;) { + const partRowIds = db + .select({ id: ingestItems.id, item_id: ingestItems.item_id }) + .from(ingestItems) + .where( + and( + eq(ingestItems.item_type, 'part'), + gte(ingestItems.item_id, partRange.start), + lt(ingestItems.item_id, partRange.end), + gt(ingestItems.id, lastPartRowId) + ) + ) + .orderBy(ingestItems.id) + .limit(KILO_SDK_HISTORY_ENUMERATION_BATCH_SIZE) + .all(); + if (partRowIds.length === 0) break; + for (const partRowId of partRowIds) { + const storageIdentity = parsePersistedKiloSdkPartStorageIdentity( + partRowId.item_id, + identity.id + ); + if (!storageIdentity) { + return { kind: 'invalid_data' }; + } + const materialized = await readKiloSdkHistoryCandidate( + partRowId.id, + rowId => readItemReference(db, rowId), + r2, + budget, + 'page_parts' + ); + if (materialized.kind === 'too_large') { + if (oversizedItemPolicy === 'fail') return materialized; + const remainingParts = countKiloSdkMessagePartRows(db, identity, partRowId.id); + if (remainingParts.kind === 'invalid_data') { + return remainingParts; + } + return { + parts: parts.sort(compareKiloSdkPart), + omittedItemCount: omittedItemCount + 1 + remainingParts.count, + aggregateBudgetExhausted: true, + }; + } + const outcome = resolveKiloSdkHistoryCandidateOutcome( + materialized, + oversizedItemPolicy, + 'page_parts' + ); + if (outcome.kind === 'skip') { + omittedItemCount += 1; + } else if (outcome.kind !== 'value') { + return outcome; + } else { + const bodyIdentity = readPartIdentity(outcome.value); + if ( + !bodyIdentity || + bodyIdentity.messageId !== identity.id || + bodyIdentity.messageId !== storageIdentity.messageId || + bodyIdentity.partId !== storageIdentity.partId + ) { + return { kind: 'invalid_data' }; + } + parts.push(outcome.value); + } + lastPartRowId = partRowId.id; + } + if (partRowIds.length < KILO_SDK_HISTORY_ENUMERATION_BATCH_SIZE) break; + } + parts.sort(compareKiloSdkPart); + return { parts, omittedItemCount, aggregateBudgetExhausted: false }; +} + +type BoundedItemRead = + | { kind: 'value'; value: Record; byteLength: number } + | { kind: 'too_large'; byteLength: number } + | { kind: 'invalid_data' } + | { kind: 'r2_missing' }; + +function isSameItemReference( + left: ItemDataRef | undefined, + right: ItemDataRef | undefined +): boolean { + return left?.item_data === right?.item_data && left?.item_data_r2_key === right?.item_data_r2_key; +} + +async function readBoundedItemData( + ref: ItemDataRef, + r2: R2Bucket, + maximumBytes: number +): Promise { + if (!ref.item_data_r2_key) { + const byteLength = new TextEncoder().encode(ref.item_data).byteLength; + if (byteLength > maximumBytes) { + return { kind: 'too_large', byteLength }; + } + const value = parseItemObject(ref.item_data); + return value ? { kind: 'value', value, byteLength } : { kind: 'invalid_data' }; + } + + const metadata = await r2.head(ref.item_data_r2_key); + if (!metadata) { + return { kind: 'r2_missing' }; + } + if (metadata.size > maximumBytes) { + return { kind: 'too_large', byteLength: metadata.size }; + } + + const object = await r2.get(ref.item_data_r2_key); + if (!object || !('body' in object)) { + return { kind: 'r2_missing' }; + } + if (object.size > maximumBytes) { + await object.body.cancel().catch(() => undefined); + return { kind: 'too_large', byteLength: object.size }; + } + + const data = await object.text(); + const byteLength = new TextEncoder().encode(data).byteLength; + if (byteLength > maximumBytes) { + return { kind: 'too_large', byteLength }; + } + const value = parseItemObject(data); + return value ? { kind: 'value', value, byteLength } : { kind: 'invalid_data' }; +} + +export async function readKiloSdkSessionItem( + ref: ItemDataRef, + resolveCurrent: () => ItemDataRef | undefined, + r2: R2Bucket, + maximumBytes: number +): Promise { + const result = await readBoundedItemData(ref, r2, maximumBytes); + if (result.kind === 'value') { + return { kind: 'value', info: result.value, byteLength: result.byteLength }; + } + if (result.kind === 'too_large') { + return { kind: 'too_large', maximumBytes }; + } + if (result.kind === 'invalid_data') { + return { kind: 'invalid_data' }; + } + + const current = resolveCurrent(); + if (!current || isSameItemReference(ref, current)) { + return { kind: 'retryable_failure' }; + } + const retry = await readBoundedItemData(current, r2, maximumBytes); + if (retry.kind === 'value') { + return { kind: 'value', info: retry.value, byteLength: retry.byteLength }; + } + if (retry.kind === 'too_large') { + return { kind: 'too_large', maximumBytes }; + } + if (retry.kind === 'invalid_data') { + return { kind: 'invalid_data' }; + } + return { kind: 'retryable_failure' }; +} + +export async function readKiloSdkHistoryCandidate( + rowId: number, + resolveItem: (rowId: number) => ItemDataRef | undefined, + r2: R2Bucket, + budget: KiloSdkHistoryReadBudget, + phase: KiloSdkHistoryReadPhase +): Promise { + if (budget.consumedBytes + KILO_SDK_HISTORY_CANDIDATE_OVERHEAD_BYTES > budget.maximumBytes) { + return { kind: 'too_large', maximumBytes: budget.maximumBytes, phase }; + } + budget.consumedBytes += KILO_SDK_HISTORY_CANDIDATE_OVERHEAD_BYTES; + const item = resolveItem(rowId); + if (!item) { + return { kind: 'missing' }; + } + const materialized = await readKiloSdkHistoryItem(item, r2, budget, phase); + if (materialized.kind !== 'retryable_failure') { + return materialized; + } + + const current = resolveItem(rowId); + if (!current || isSameItemReference(item, current)) { + return materialized; + } + return readKiloSdkHistoryItem(current, r2, budget, phase); +} + +export async function readKiloSdkHistoryItem( + ref: ItemDataRef, + r2: R2Bucket, + budget: KiloSdkHistoryReadBudget, + phase: KiloSdkHistoryReadPhase +): Promise< + | { kind: 'value'; value: Record } + | KiloSdkIntrinsicallyTooLarge + | KiloSdkInvalidData + | KiloSdkHistoryTooLarge + | KiloSdkHistoryRetryableFailure +> { + const materialized = await readBoundedItemData( + ref, + r2, + budget.maximumBytes - budget.consumedBytes + ); + if (materialized.kind === 'too_large') { + return materialized.byteLength > budget.maximumBytes + ? { kind: 'intrinsically_too_large' } + : { kind: 'too_large', maximumBytes: budget.maximumBytes, phase }; + } + if (materialized.kind === 'invalid_data') { + return { kind: 'invalid_data' }; + } + if (materialized.kind === 'r2_missing') { + return { kind: 'retryable_failure', phase }; + } + budget.consumedBytes += materialized.byteLength; + return { kind: 'value', value: materialized.value }; +} + +function parseItemObject(value: string): Record | null { + try { + const parsed = z.record(z.string(), z.unknown()).safeParse(JSON.parse(value)); + return parsed.success ? parsed.data : null; + } catch { + return null; + } +} + +function readMessageIdentity(message: Record): KiloSdkMessagesLegacyCursor | null { + const parsed = z + .object({ + id: messageIdSchema, + time: z.object({ created: z.number().nonnegative() }), + }) + .transform(value => ({ id: value.id, time: value.time.created })) + .safeParse(message); + return parsed.success ? parsed.data : null; +} + +function readPartIdentity( + part: Record +): PersistedKiloSdkPartStorageIdentity | null { + const parsed = z + .object({ id: partIdSchema, messageID: messageIdSchema }) + .transform(value => ({ partId: value.id, messageId: value.messageID })) + .safeParse(part); + return parsed.success ? parsed.data : null; +} + +function compareKiloSdkMessageInfo( + left: MaterializedKiloSdkMessage, + right: MaterializedKiloSdkMessage +): number { + return left.identity.time - right.identity.time || compareId(left.identity.id, right.identity.id); +} + +function compareId(left: string, right: string): number { + return left < right ? -1 : left > right ? 1 : 0; +} + +function compareKiloSdkPart(left: Record, right: Record): number { + const idSchema = z.object({ id: z.string() }); + const leftId = idSchema.safeParse(left); + const rightId = idSchema.safeParse(right); + if (!leftId.success || !rightId.success) return 0; + return compareId(leftId.data.id, rightId.data.id); +} diff --git a/services/session-ingest/src/dos/session-ingest-history-budget.test.ts b/services/session-ingest/src/dos/session-ingest-history-budget.test.ts new file mode 100644 index 0000000000..0a8ff0d04b --- /dev/null +++ b/services/session-ingest/src/dos/session-ingest-history-budget.test.ts @@ -0,0 +1,224 @@ +import { describe, expect, it, vi } from 'vitest'; + +vi.mock('cloudflare:workers', () => ({ + DurableObject: class DurableObject { + constructor(_state: unknown, _env: unknown) {} + }, +})); + +import { + KILO_SDK_HISTORY_CANDIDATE_OVERHEAD_BYTES, + createKiloSdkHistoryReadBudget, + readKiloSdkHistoryCandidate, + readKiloSdkHistoryItem, + readKiloSdkSessionItem, +} from './kilo-sdk-materialization'; + +function r2Body(value: Record) { + const data = JSON.stringify(value); + return { + size: new TextEncoder().encode(data).byteLength, + text: vi.fn(async () => data), + body: new ReadableStream(), + }; +} + +describe('Kilo SDK session snapshot materialization', () => { + it('returns invalid_data for malformed snapshots instead of throwing parse errors', async () => { + await expect( + readKiloSdkSessionItem( + { item_data: 'not-json', item_data_r2_key: null }, + () => undefined, + {} as R2Bucket, + 32 + ) + ).resolves.toEqual({ kind: 'invalid_data' }); + }); + + it('rejects oversized inline snapshots before parsing', async () => { + await expect( + readKiloSdkSessionItem( + { item_data: JSON.stringify({ title: 'x'.repeat(40) }), item_data_r2_key: null }, + () => undefined, + {} as R2Bucket, + 32 + ) + ).resolves.toEqual({ kind: 'too_large', maximumBytes: 32 }); + }); + + it('re-reads a replaced R2-backed snapshot once under the same maximum', async () => { + const replacement = { id: 'ses_new', title: 'latest' }; + const resolveCurrent = vi.fn(() => ({ item_data: '{}', item_data_r2_key: 'snapshot-new' })); + const head = vi + .fn() + .mockResolvedValueOnce(null) + .mockResolvedValueOnce({ + size: new TextEncoder().encode(JSON.stringify(replacement)).byteLength, + }); + const get = vi.fn(async () => r2Body(replacement)); + const bucket = { head, get } as unknown as R2Bucket; + + await expect( + readKiloSdkSessionItem( + { item_data: '{}', item_data_r2_key: 'snapshot-old' }, + resolveCurrent, + bucket, + 4096 + ) + ).resolves.toEqual({ + kind: 'value', + info: replacement, + byteLength: new TextEncoder().encode(JSON.stringify(replacement)).byteLength, + }); + expect(resolveCurrent).toHaveBeenCalledTimes(1); + expect(get).toHaveBeenCalledWith('snapshot-new'); + }); + + it('fails deliberately after one unresolved R2 snapshot attempt', async () => { + const resolveCurrent = vi.fn(() => ({ item_data: '{}', item_data_r2_key: 'snapshot-new' })); + const head = vi.fn(async () => null); + const get = vi.fn(); + const bucket = { head, get } as unknown as R2Bucket; + + await expect( + readKiloSdkSessionItem( + { item_data: '{}', item_data_r2_key: 'snapshot-old' }, + resolveCurrent, + bucket, + 4096 + ) + ).resolves.toEqual({ kind: 'retryable_failure' }); + expect(resolveCurrent).toHaveBeenCalledTimes(1); + expect(head).toHaveBeenCalledTimes(2); + expect(get).not.toHaveBeenCalled(); + }); +}); + +describe('Kilo SDK history materialization budget', () => { + it('rejects an inline item that exceeds the remaining cumulative budget', async () => { + const budget = createKiloSdkHistoryReadBudget(32); + + await expect( + readKiloSdkHistoryItem( + { item_data: JSON.stringify({ text: 'x'.repeat(40) }), item_data_r2_key: null }, + {} as R2Bucket, + budget, + 'message_scan' + ) + ).resolves.toEqual({ kind: 'intrinsically_too_large' }); + }); + + it('applies one cumulative budget across separately materialized items', async () => { + const budget = createKiloSdkHistoryReadBudget(40); + const first = JSON.stringify({ id: 'first', text: 'one' }); + const second = JSON.stringify({ id: 'second', text: 'two' }); + + await expect( + readKiloSdkHistoryItem( + { item_data: first, item_data_r2_key: null }, + {} as R2Bucket, + budget, + 'message_scan' + ) + ).resolves.toMatchObject({ kind: 'value' }); + await expect( + readKiloSdkHistoryItem( + { item_data: second, item_data_r2_key: null }, + {} as R2Bucket, + budget, + 'page_parts' + ) + ).resolves.toEqual({ + kind: 'too_large', + maximumBytes: 40, + phase: 'page_parts', + }); + }); + + it('rejects candidate growth before resolving more tiny persisted bodies', async () => { + const budget = createKiloSdkHistoryReadBudget( + KILO_SDK_HISTORY_CANDIDATE_OVERHEAD_BYTES + new TextEncoder().encode('{}').byteLength + ); + const resolveItem = vi.fn(() => ({ item_data: '{}', item_data_r2_key: null })); + + await expect( + readKiloSdkHistoryCandidate(1, resolveItem, {} as R2Bucket, budget, 'message_scan') + ).resolves.toMatchObject({ kind: 'value' }); + await expect( + readKiloSdkHistoryCandidate(2, resolveItem, {} as R2Bucket, budget, 'message_scan') + ).resolves.toEqual({ + kind: 'too_large', + maximumBytes: + KILO_SDK_HISTORY_CANDIDATE_OVERHEAD_BYTES + new TextEncoder().encode('{}').byteLength, + phase: 'message_scan', + }); + expect(resolveItem).toHaveBeenCalledTimes(1); + }); + + it('rejects an oversized R2-backed part using metadata without consuming its body', async () => { + const head = vi.fn(async () => ({ size: 33 })); + const get = vi.fn(async () => { + throw new Error('oversized R2 body must not be requested'); + }); + const bucket = { head, get } as unknown as R2Bucket; + + await expect( + readKiloSdkHistoryItem( + { item_data: '{}', item_data_r2_key: 'large-part' }, + bucket, + createKiloSdkHistoryReadBudget(32), + 'page_parts' + ) + ).resolves.toEqual({ kind: 'intrinsically_too_large' }); + expect(head).toHaveBeenCalledWith('large-part'); + expect(get).not.toHaveBeenCalled(); + }); + + it('re-reads a replaced R2-backed message once instead of returning empty data', async () => { + const replacement = { id: 'msg_new', time: { created: 100 } }; + const resolveItem = vi + .fn<(_rowId: number) => { item_data: string; item_data_r2_key: string }>() + .mockReturnValueOnce({ item_data: '{}', item_data_r2_key: 'missing-old' }) + .mockReturnValueOnce({ item_data: '{}', item_data_r2_key: 'current-new' }); + const head = vi + .fn() + .mockResolvedValueOnce(null) + .mockResolvedValueOnce({ + size: new TextEncoder().encode(JSON.stringify(replacement)).byteLength, + }); + const get = vi.fn(async () => r2Body(replacement)); + const bucket = { head, get } as unknown as R2Bucket; + + await expect( + readKiloSdkHistoryCandidate( + 1, + resolveItem, + bucket, + createKiloSdkHistoryReadBudget(4096), + 'message_scan' + ) + ).resolves.toEqual({ kind: 'value', value: replacement }); + expect(resolveItem).toHaveBeenCalledTimes(2); + expect(get).toHaveBeenCalledTimes(1); + expect(get).toHaveBeenCalledWith('current-new'); + }); + + it('returns a retryable failure when an R2-backed part still has no materialized object', async () => { + const resolveItem = vi.fn(() => ({ item_data: '{}', item_data_r2_key: 'missing-part' })); + const head = vi.fn(async () => null); + const get = vi.fn(); + const bucket = { head, get } as unknown as R2Bucket; + + await expect( + readKiloSdkHistoryCandidate( + 2, + resolveItem, + bucket, + createKiloSdkHistoryReadBudget(4096), + 'page_parts' + ) + ).resolves.toEqual({ kind: 'retryable_failure', phase: 'page_parts' }); + expect(resolveItem).toHaveBeenCalledTimes(2); + expect(get).not.toHaveBeenCalled(); + }); +}); diff --git a/services/session-ingest/src/queue-consumer.test.ts b/services/session-ingest/src/queue-consumer.test.ts index 65604d0df5..531ba289b8 100644 --- a/services/session-ingest/src/queue-consumer.test.ts +++ b/services/session-ingest/src/queue-consumer.test.ts @@ -35,10 +35,10 @@ vi.mock('./session-events', async importOriginal => { }; }); -// Mock ingest-limits so we can use a small MAX_SINGLE_ITEM_BYTES in oversize tests +// Mock ingest-limits so we can exercise both streaming and SQLite-row compaction thresholds. vi.mock('./util/ingest-limits', () => ({ - MAX_INGEST_ITEM_BYTES: 2 * 1024 * 1024, - MAX_SINGLE_ITEM_BYTES: 50, + MAX_INGEST_ITEM_BYTES: 100, + MAX_SINGLE_ITEM_BYTES: 500, })); import { getWorkerDb } from '@kilocode/db/client'; @@ -80,11 +80,11 @@ describe('createItemExtractor', () => { }); it('skips oversized items (byte budget)', () => { - // MAX_SINGLE_ITEM_BYTES is mocked to 50 + // MAX_SINGLE_ITEM_BYTES is mocked to 500 const ext = createItemExtractor('test-key'); - // Create an item that exceeds 50 bytes - const bigValue = 'x'.repeat(60); + // Create an item that exceeds 500 bytes + const bigValue = 'x'.repeat(600); const payload = JSON.stringify({ data: [ { type: 'big', data: { content: bigValue } }, @@ -100,12 +100,12 @@ describe('createItemExtractor', () => { }); it('clears skippingItem when oversize item ends on closing brace', () => { - // MAX_SINGLE_ITEM_BYTES is mocked to 50 + // MAX_SINGLE_ITEM_BYTES is mocked to 500 const ext = createItemExtractor('test-key'); // A flat object (no nested braces) that exceeds budget — the closing } // is the token that triggers the budget check AND ends the item at depth=2 - const bigValue = 'y'.repeat(60); + const bigValue = 'y'.repeat(600); const payload = JSON.stringify({ data: [{ big: bigValue }, { type: 'after', ok: true }], }); @@ -142,6 +142,71 @@ describe('createItemExtractor', () => { }); }); +describe('queue', () => { + it('passes full parsed oversized message data and its R2 reference into ingest', async () => { + const ingest = vi.fn(async () => ({ changes: [] })); + vi.mocked(getSessionIngestDO).mockReturnValue({ ingest } as never); + const limit = vi.fn(async () => [{ session_id: 'ses_compacted' }]); + const where = vi.fn(() => ({ limit })); + const from = vi.fn(() => ({ where })); + vi.mocked(getWorkerDb).mockReturnValue({ select: vi.fn(() => ({ from })) } as never); + const data = { + id: 'msg_compacted', + sessionID: 'ses_compacted', + time: { created: 123 }, + content: 'x'.repeat(150), + }; + const body = JSON.stringify({ data: [{ type: 'message', data }] }); + const put = vi.fn(async () => undefined); + const deleteObject = vi.fn(async () => undefined); + const env = { + HYPERDRIVE: { connectionString: 'postgres://unused' }, + SESSION_INGEST_R2: { + get: vi.fn(async () => new Response(body)), + put, + delete: deleteObject, + }, + } as never; + const ack = vi.fn(); + const retry = vi.fn(); + const ctx = { waitUntil: vi.fn() } as unknown as ExecutionContext; + + await queue( + { + messages: [ + { + body: { + r2Key: 'staging/items', + kiloUserId: 'usr_compacted', + sessionId: 'ses_compacted', + ingestVersion: 1, + ingestedAt: 456, + }, + ack, + retry, + }, + ], + } as never, + env, + ctx + ); + + const expectedR2Key = 'items/usr_compacted/ses_compacted/message/msg_compacted/456'; + expect(put).toHaveBeenCalledWith(expectedR2Key, JSON.stringify(data)); + expect(ingest).toHaveBeenCalledWith( + [{ type: 'message', data }], + 'usr_compacted', + 'ses_compacted', + 1, + 456, + { 'message/msg_compacted': expectedR2Key } + ); + expect(deleteObject).toHaveBeenCalledWith('staging/items'); + expect(ack).toHaveBeenCalledTimes(1); + expect(retry).not.toHaveBeenCalled(); + }); +}); + describe('computeSessionMetadataUpdates', () => { const fixedNow = () => '2026-05-05T00:00:00.000Z'; diff --git a/services/session-ingest/src/session-ingest-rpc.test.ts b/services/session-ingest/src/session-ingest-rpc.test.ts new file mode 100644 index 0000000000..429424962f --- /dev/null +++ b/services/session-ingest/src/session-ingest-rpc.test.ts @@ -0,0 +1,822 @@ +import type * as DrizzleOrm from 'drizzle-orm'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +vi.mock('cloudflare:workers', () => ({ + DurableObject: class DurableObject { + constructor(_state: unknown, _env: unknown) {} + }, + WorkerEntrypoint: class WorkerEntrypoint { + env: unknown; + ctx: ExecutionContext; + + constructor(ctx: ExecutionContext, env: unknown) { + this.ctx = ctx; + this.env = env; + } + }, +})); + +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: vi.fn(), +})); + +vi.mock('drizzle-orm', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + desc: vi.fn(actual.desc), + gte: vi.fn(actual.gte), + isNotNull: vi.fn(actual.isNotNull), + or: vi.fn(actual.or), + }; +}); + +vi.mock('./dos/SessionIngestDO', () => ({ + getSessionIngestDO: vi.fn(), +})); + +vi.mock('./dos/SessionAccessCacheDO', () => ({ + getSessionAccessCacheDO: vi.fn(), +})); + +import { getWorkerDb } from '@kilocode/db/client'; +import { cli_sessions_v2, organization_memberships } from '@kilocode/db/schema'; +import { + decodeKiloSdkMessagesCursor, + encodeKiloSdkMessagesCursor, + MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE, + messageIdSchema, + partIdSchema, + validateKiloSdkMessagesCursor, +} from '@kilocode/session-ingest-contracts'; +import { desc, gte, isNotNull, or } from 'drizzle-orm'; +import { getSessionIngestDO } from './dos/SessionIngestDO'; +import { SessionIngestRPC } from './session-ingest-rpc'; + +const sdkSessionInfoFixture = { + id: 'ses_12345678901234567890123456', + slug: 'quiet-forest', + projectID: 'project-cloud-agent', + directory: '/workspace/cloud-agent', + title: 'SDK attach session', + agent: 'build', + model: { id: 'anthropic/claude-sonnet-4', providerID: 'openrouter' }, + version: '7.2.52', + time: { created: 1761000000000, updated: 1761000001000 }, +}; + +const sdkUserMessageFixture = { + id: 'msg_user_01', + sessionID: sdkSessionInfoFixture.id, + role: 'user' as const, + time: { created: 1761000000100 }, + agent: 'build', + model: { providerID: 'openrouter', modelID: 'anthropic/claude-sonnet-4' }, +}; +const sdkTextPartFixture = { + id: 'prt_user_01', + sessionID: sdkSessionInfoFixture.id, + messageID: sdkUserMessageFixture.id, + type: 'text' as const, + text: 'Attach to this persisted turn', +}; +const sdkStoredMessageFixture = { info: sdkUserMessageFixture, parts: [sdkTextPartFixture] }; + +type MappingRow = { + kiloSessionId?: string; + cloudAgentSessionId: string | null; + title?: string | null; + createdAt?: string; + updatedAt?: string; +}; + +function makeDbFakes(rows: MappingRow[]) { + const selectResult = vi.fn(async () => rows); + const select = { + from: vi.fn(() => select), + leftJoin: vi.fn(() => select), + where: vi.fn(() => select), + orderBy: vi.fn(() => select), + limit: vi.fn(() => select), + then: vi.fn((resolve: (value: unknown) => unknown) => resolve(selectResult())), + }; + const db = { + select: vi.fn(() => select), + }; + return { db, select, selectResult }; +} + +function makeRpc(db: ReturnType['db']) { + vi.mocked(getWorkerDb).mockReturnValue(db as never); + const ctx = { + waitUntil: () => undefined, + passThroughOnException: () => undefined, + } as unknown as ConstructorParameters[0]; + const env = { + HYPERDRIVE: { connectionString: 'postgres://test' }, + } as unknown as ConstructorParameters[1]; + return new SessionIngestRPC(ctx, env); +} + +describe('Kilo SDK persisted identity schemas', () => { + it('accepts generated message IDs and rejects non-message, slash-bearing, or NUL-bearing IDs', () => { + expect(messageIdSchema.safeParse('msg_storage').success).toBe(true); + expect(messageIdSchema.safeParse('other_storage').success).toBe(false); + expect(messageIdSchema.safeParse('msg_storage/child').success).toBe(false); + expect(messageIdSchema.safeParse('msg_storage\u0000child').success).toBe(false); + }); + + it('accepts generated part IDs and rejects non-part, slash-bearing, or NUL-bearing IDs', () => { + expect(partIdSchema.safeParse('prt_storage').success).toBe(true); + expect(partIdSchema.safeParse('other_storage').success).toBe(false); + expect(partIdSchema.safeParse('prt_storage/child').success).toBe(false); + expect(partIdSchema.safeParse('prt_storage\u0000child').success).toBe(false); + }); +}); + +describe('Kilo SDK message cursor codec', () => { + it('round-trips the existing opaque base64url wire encoding', () => { + const cursor = { id: 'msg_user_01', time: 1761000000100 }; + const encoded = encodeKiloSdkMessagesCursor(cursor); + + expect(encoded).toBe('eyJpZCI6Im1zZ191c2VyXzAxIiwidGltZSI6MTc2MTAwMDAwMDEwMH0'); + expect(decodeKiloSdkMessagesCursor(encoded)).toEqual(cursor); + expect(validateKiloSdkMessagesCursor(encoded)).toBe(true); + }); + + it('rejects malformed, non-message, and non-strict cursor payloads', () => { + const encodeUnchecked = (value: unknown) => + btoa(JSON.stringify(value)).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/g, ''); + + expect(validateKiloSdkMessagesCursor('not-valid')).toBe(false); + expect(validateKiloSdkMessagesCursor(encodeUnchecked({ id: 'other_01', time: 1 }))).toBe(false); + expect( + validateKiloSdkMessagesCursor(encodeUnchecked({ id: 'msg_parent/child', time: 1 })) + ).toBe(false); + expect( + validateKiloSdkMessagesCursor(encodeUnchecked({ id: 'msg_parent\u0000child', time: 1 })) + ).toBe(false); + expect( + validateKiloSdkMessagesCursor(encodeUnchecked({ id: 'msg_user_01', time: 1, extra: true })) + ).toBe(false); + expect(validateKiloSdkMessagesCursor(encodeUnchecked({ id: 'msg_user_01', time: -1 }))).toBe( + false + ); + expect( + validateKiloSdkMessagesCursor(encodeUnchecked({ version: 2, beforeMessageId: 'msg_user_01' })) + ).toBe(false); + expect(() => + decodeKiloSdkMessagesCursor(encodeUnchecked({ id: 'other_01', time: 1 })) + ).toThrow(); + }); +}); + +describe('SessionIngestRPC.resolveCloudAgentRootSessionForKiloSession', () => { + beforeEach(() => { + vi.resetAllMocks(); + }); + + it('returns the Cloud Agent session ID for an owned root Kilo session mapping', async () => { + const { db, select } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + const rpc = makeRpc(db); + + const result = await rpc.resolveCloudAgentRootSessionForKiloSession({ + kiloUserId: 'usr_owner', + kiloSessionId: 'ses_12345678901234567890123456', + }); + + expect(result).toEqual({ cloudAgentSessionId: 'agent_owned_root' }); + expect(db.select).toHaveBeenCalledWith({ cloudAgentSessionId: expect.anything() }); + expect(select.leftJoin).toHaveBeenCalledWith(organization_memberships, expect.anything()); + expect(or).toHaveBeenCalled(); + }); + + it('returns null when no owned Cloud Agent root mapping is found', async () => { + const { db } = makeDbFakes([]); + const rpc = makeRpc(db); + + const result = await rpc.resolveCloudAgentRootSessionForKiloSession({ + kiloUserId: 'usr_owner', + kiloSessionId: 'ses_12345678901234567890123456', + }); + + expect(result).toBeNull(); + }); + + it('returns null when the selected row has no Cloud Agent mapping', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: null }]); + const rpc = makeRpc(db); + + const result = await rpc.resolveCloudAgentRootSessionForKiloSession({ + kiloUserId: 'usr_owner', + kiloSessionId: 'ses_12345678901234567890123456', + }); + + expect(result).toBeNull(); + }); + + it('rejects invalid Kilo session IDs before querying the database', async () => { + const { db } = makeDbFakes([]); + const rpc = makeRpc(db); + + await expect( + rpc.resolveCloudAgentRootSessionForKiloSession({ + kiloUserId: 'usr_owner', + kiloSessionId: 'not-a-session', + }) + ).rejects.toThrow(); + expect(db.select).not.toHaveBeenCalled(); + }); +}); + +describe('SessionIngestRPC.getCloudAgentRootSessionSnapshot', () => { + beforeEach(() => { + vi.resetAllMocks(); + }); + + it('returns a materialized SDK snapshot only for an owned root Cloud Agent mapping', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkSessionSnapshot: vi.fn(async () => ({ + kind: 'value', + info: sdkSessionInfoFixture, + byteLength: 512, + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionSnapshot({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + snapshot: { kind: 'value', info: sdkSessionInfoFixture, byteLength: 512 }, + }); + + const { db: missingDb } = makeDbFakes([]); + const missingRpc = makeRpc(missingDb); + await expect( + missingRpc.getCloudAgentRootSessionSnapshot({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toBeNull(); + }); + + it('preserves explicit bounded and pending outcomes for an authorized root', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_pending_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkSessionSnapshot: vi + .fn() + .mockResolvedValueOnce({ kind: 'pending' }) + .mockResolvedValueOnce({ kind: 'too_large', maximumBytes: 8 * 1024 * 1024 }) + .mockResolvedValueOnce({ kind: 'retryable_failure' }), + } as never); + const rpc = makeRpc(db); + + for (const snapshot of [ + { kind: 'pending' }, + { kind: 'too_large', maximumBytes: 8 * 1024 * 1024 }, + { kind: 'retryable_failure' }, + ]) { + await expect( + rpc.getCloudAgentRootSessionSnapshot({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_pending_root', + snapshot, + }); + } + }); + + it('returns invalid_data when a persisted snapshot is outside the strict outward contract', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_invalid_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkSessionSnapshot: vi.fn(async () => ({ + kind: 'value', + info: { ...sdkSessionInfoFixture, time: { created: 'invalid', updated: 1761000001000 } }, + byteLength: 512, + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionSnapshot({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_invalid_root', + snapshot: { kind: 'invalid_data' }, + }); + }); + + it('does not convert snapshot DO failures into invalid_data', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_failed_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkSessionSnapshot: vi.fn(async () => { + throw new Error('snapshot unavailable'); + }), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionSnapshot({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).rejects.toThrow('snapshot unavailable'); + }); +}); + +describe('SessionIngestRPC.getCloudAgentRootSessionMessages', () => { + beforeEach(() => { + vi.resetAllMocks(); + }); + + it('returns null for an unavailable root and distinguishes pending from empty materialized history', async () => { + const { db: missingDb } = makeDbFakes([]); + const missingRpc = makeRpc(missingDb); + await expect( + missingRpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toBeNull(); + + const { db: ownedDb } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi + .fn() + .mockResolvedValueOnce(null) + .mockResolvedValueOnce({ messages: [], nextCursor: null }), + } as never); + const ownedRpc = makeRpc(ownedDb); + + await expect( + ownedRpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: null, + }); + await expect( + ownedRpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + }); + }); + + it('returns full materialized history for native limit zero requests', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + const readKiloSdkMessages = vi.fn(async () => ({ + messages: [sdkStoredMessageFixture], + nextCursor: null, + })); + vi.mocked(getSessionIngestDO).mockReturnValue({ readKiloSdkMessages } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + limit: 0, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { messages: [sdkStoredMessageFixture], nextCursor: null, omittedItemCount: 0 }, + }); + expect(readKiloSdkMessages).toHaveBeenCalledWith({ limit: 0, before: undefined }); + }); + + it('normalizes legacy history pages without omission metadata to zero', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ messages: [], nextCursor: null })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toMatchObject({ + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + }); + }); + + it('returns exact persisted SDK message history and forwards native paging input', async () => { + const cursor = 'eyJpZCI6Im1zZ191c2VyXzAxIiwidGltZSI6MTc2MTAwMDAwMDEwMH0'; + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + const readKiloSdkMessages = vi.fn(async () => ({ + messages: [sdkStoredMessageFixture], + nextCursor: cursor, + omittedItemCount: 3, + })); + vi.mocked(getSessionIngestDO).mockReturnValue({ readKiloSdkMessages } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + limit: 2, + before: cursor, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { messages: [sdkStoredMessageFixture], nextCursor: cursor, omittedItemCount: 3 }, + }); + expect(readKiloSdkMessages).toHaveBeenCalledWith({ limit: 2, before: cursor }); + }); + + it('omits legacy before/after summary diffs while preserving current patch diffs for public projection', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + const currentDiff = { + file: '/workspace/private/current.ts', + patch: '@@ -1 +1 @@', + additions: 1, + deletions: 1, + }; + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ + messages: [ + { + info: { + ...sdkUserMessageFixture, + summary: { + title: 'Persisted summary', + diffs: [ + { + file: '/workspace/private/historical.ts', + before: 'const value = 1;', + after: 'const value = 2;', + additions: 1, + deletions: 1, + }, + currentDiff, + ], + }, + }, + parts: [sdkTextPartFixture], + }, + ], + nextCursor: null, + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { + messages: [ + { + info: { + ...sdkUserMessageFixture, + summary: { title: 'Persisted summary', diffs: [currentDiff] }, + }, + parts: [sdkTextPartFixture], + }, + ], + nextCursor: null, + omittedItemCount: 0, + }, + }); + }); + + it('returns invalid_data for ambiguous summary diffs instead of silently discarding current patch detail', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ + messages: [ + { + info: { + ...sdkUserMessageFixture, + summary: { + diffs: [ + { + file: '/workspace/private/ambiguous.ts', + patch: '@@ -1 +1 @@', + before: 'const value = 1;', + after: 'const value = 2;', + additions: 1, + deletions: 1, + }, + ], + }, + }, + parts: [sdkTextPartFixture], + }, + ], + nextCursor: null, + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { kind: 'invalid_data' }, + }); + }); + + it('returns invalid_data for malformed historical summary diffs instead of silently dropping them', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ + messages: [ + { + info: { + ...sdkUserMessageFixture, + summary: { + diffs: [ + { + file: '/workspace/private/historical.ts', + before: 1, + after: 'const value = 2;', + additions: 1, + deletions: 1, + }, + ], + }, + }, + parts: [sdkTextPartFixture], + }, + ], + nextCursor: null, + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { kind: 'invalid_data' }, + }); + }); + + it('does not convert transcript DO failures into invalid_data', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_failed_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => { + throw new Error('transcript unavailable'); + }), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).rejects.toThrow('transcript unavailable'); + }); + + it('preserves a retryable history outcome for facade error mapping', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ + kind: 'retryable_failure', + phase: 'page_parts', + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + limit: 1, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { kind: 'retryable_failure', phase: 'page_parts' }, + }); + }); + + it('preserves a durable too-large history outcome for facade error mapping', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ + kind: 'too_large', + maximumBytes: 8 * 1024 * 1024, + phase: 'message_scan', + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + limit: 1, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { + kind: 'too_large', + maximumBytes: 8 * 1024 * 1024, + phase: 'message_scan', + }, + }); + }); + + it('rejects before without a positive limit and invalid paging input before mapping lookup', async () => { + const { db } = makeDbFakes([]); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + before: 'not-valid', + }) + ).rejects.toThrow(); + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + limit: 0, + before: 'not-valid', + }) + ).rejects.toThrow(); + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + limit: 2, + before: 'not-valid', + }) + ).rejects.toThrow(); + expect(db.select).not.toHaveBeenCalled(); + }); + + it('rejects a decodable non-message cursor before mapping lookup', async () => { + const { db } = makeDbFakes([]); + const rpc = makeRpc(db); + const cursor = btoa(JSON.stringify({ id: 'other_01', time: 1 })).replace(/=+$/g, ''); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + limit: 2, + before: cursor, + }) + ).rejects.toThrow(); + expect(db.select).not.toHaveBeenCalled(); + }); + + it('rejects positive page limits above the shared maximum before mapping lookup', async () => { + const { db } = makeDbFakes([]); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + limit: MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE + 1, + }) + ).rejects.toThrow(); + expect(db.select).not.toHaveBeenCalled(); + }); +}); + +describe('SessionIngestRPC.listCloudAgentRootSessions', () => { + beforeEach(() => { + vi.resetAllMocks(); + }); + + it('returns mapped root summaries in database order without opening session DOs', async () => { + const secondSessionId = 'ses_abcdefghijklmnopqrstuvwxyz'; + const createdAt = '2026-05-27 20:53:24.190157+00'; + const updatedAt = '2026-05-28 09:13:37.651263+00'; + const { db, select } = makeDbFakes([ + { + kiloSessionId: secondSessionId, + cloudAgentSessionId: 'agent_same_time_b', + title: 'Second', + createdAt, + updatedAt, + }, + { + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_same_time_a', + title: 'First', + createdAt, + updatedAt, + }, + ]); + const rpc = makeRpc(db); + + await expect( + rpc.listCloudAgentRootSessions({ + kiloUserId: 'usr_owner', + start: 1761000000000, + limit: 2, + }) + ).resolves.toEqual([ + { + kiloSessionId: secondSessionId, + cloudAgentSessionId: 'agent_same_time_b', + title: 'Second', + created: new Date(createdAt).getTime(), + updated: new Date(updatedAt).getTime(), + }, + { + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_same_time_a', + title: 'First', + created: new Date(createdAt).getTime(), + updated: new Date(updatedAt).getTime(), + }, + ]); + expect(getSessionIngestDO).not.toHaveBeenCalled(); + expect(select.leftJoin).toHaveBeenCalledWith(organization_memberships, expect.anything()); + expect(or).toHaveBeenCalled(); + expect(isNotNull).toHaveBeenCalledWith(cli_sessions_v2.cloud_agent_session_id); + expect(gte).toHaveBeenCalledWith( + cli_sessions_v2.updated_at, + new Date(1761000000000).toISOString() + ); + expect(desc).toHaveBeenNthCalledWith(1, cli_sessions_v2.updated_at); + expect(desc).toHaveBeenNthCalledWith(2, cli_sessions_v2.session_id); + expect(select.orderBy).toHaveBeenCalled(); + expect(select.limit).toHaveBeenCalledWith(2); + }); + + it('returns mapped roots without requiring a materialized SDK snapshot and bounds titles', async () => { + const timestamp = '2026-05-28 09:13:37.651263+00'; + const longTitle = 'x'.repeat(600); + const { db } = makeDbFakes([ + { + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_org_root', + title: longTitle, + createdAt: timestamp, + updatedAt: timestamp, + }, + ]); + const rpc = makeRpc(db); + + await expect(rpc.listCloudAgentRootSessions({ kiloUserId: 'usr_owner' })).resolves.toEqual([ + { + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_org_root', + title: longTitle.slice(0, 512), + created: new Date(timestamp).getTime(), + updated: new Date(timestamp).getTime(), + }, + ]); + expect(getSessionIngestDO).not.toHaveBeenCalled(); + }); + + it('rejects unsafe list bounds before querying root mappings', async () => { + const { db } = makeDbFakes([]); + const rpc = makeRpc(db); + + await expect( + rpc.listCloudAgentRootSessions({ kiloUserId: 'usr_owner', limit: 0 }) + ).rejects.toThrow(); + await expect( + rpc.listCloudAgentRootSessions({ kiloUserId: 'usr_owner', limit: 101 }) + ).rejects.toThrow(); + expect(db.select).not.toHaveBeenCalled(); + }); +}); diff --git a/services/session-ingest/src/session-ingest-rpc.ts b/services/session-ingest/src/session-ingest-rpc.ts index 6c51fdd617..34252535a7 100644 --- a/services/session-ingest/src/session-ingest-rpc.ts +++ b/services/session-ingest/src/session-ingest-rpc.ts @@ -1,8 +1,28 @@ import { WorkerEntrypoint } from 'cloudflare:workers'; -import { z } from 'zod'; -import { eq, and } from 'drizzle-orm'; +import { eq, and, desc, gte, isNotNull, isNull, or, sql } from 'drizzle-orm'; import { getWorkerDb } from '@kilocode/db/client'; -import { cli_sessions_v2 } from '@kilocode/db/schema'; +import { cli_sessions_v2, organization_memberships } from '@kilocode/db/schema'; +import { + createSessionForCloudAgentSchema, + deleteSessionForCloudAgentSchema, + getCloudAgentRootSessionMessagesSchema, + kiloSdkSessionSnapshotOutcomeSchema, + listCloudAgentRootSessionsSchema, + persistedKiloSdkMessageHistorySchema, + resolveCloudAgentRootSessionSchema, + type CloudAgentRootSessionSnapshot, + type CloudAgentRootSessionSummary, + type CreateSessionForCloudAgentParams, + type DeleteSessionForCloudAgentParams, + type GetCloudAgentRootSessionMessagesParams, + type GetCloudAgentRootSessionMessagesResult, + type GetCloudAgentRootSessionSnapshotParams, + type GetCloudAgentRootSessionSnapshotResult, + type ListCloudAgentRootSessionsParams, + type ResolveCloudAgentRootSessionForKiloSessionParams, + type ResolveCloudAgentRootSessionForKiloSessionResult, + type SessionIngestRpcMethods, +} from '@kilocode/session-ingest-contracts'; import type { Env } from './env'; import { getSessionIngestDO } from './dos/SessionIngestDO'; @@ -11,9 +31,28 @@ import { withDORetry } from '@kilocode/worker-utils'; import { app } from './app'; import { mapSessionEventRow, notifyUserSessionEvent } from './session-events'; -const sessionIdSchema = z.string().startsWith('ses_').length(30); +const MAX_CLOUD_AGENT_ROOT_SESSION_TITLE_CHARACTERS = 512; -export class SessionIngestRPC extends WorkerEntrypoint { +function databaseTimestampToMilliseconds(value: string): number { + const timestamp = new Date(value).getTime(); + if (!Number.isFinite(timestamp)) { + throw new Error('Invalid Cloud Agent root session timestamp'); + } + return timestamp; +} + +function organizationMembershipJoinCondition(kiloUserId: string) { + return and( + eq(organization_memberships.organization_id, cli_sessions_v2.organization_id), + eq(organization_memberships.kilo_user_id, kiloUserId) + ); +} + +function personalOrAccessibleOrganizationCondition() { + return or(isNull(cli_sessions_v2.organization_id), isNotNull(organization_memberships.id)); +} + +export class SessionIngestRPC extends WorkerEntrypoint implements SessionIngestRpcMethods { // Delegate HTTP requests to the Hono app so callers using the service // binding can `.fetch()` against this entrypoint (not just call RPC methods). fetch(request: Request): Response | Promise { @@ -27,24 +66,8 @@ export class SessionIngestRPC extends WorkerEntrypoint { * Uses ON CONFLICT DO UPDATE to set cloud_agent_session_id (and organization_id * if provided), matching the behavior previously in the backend routers. */ - async createSessionForCloudAgent(params: { - sessionId: string; - kiloUserId: string; - cloudAgentSessionId: string; - organizationId?: string; - createdOnPlatform: string; - title?: string; - }): Promise { - const parsed = z - .object({ - sessionId: sessionIdSchema, - kiloUserId: z.string().min(1), - cloudAgentSessionId: z.string().min(1), - organizationId: z.string().optional(), - createdOnPlatform: z.string().min(1), - title: z.string().optional(), - }) - .parse(params); + async createSessionForCloudAgent(params: CreateSessionForCloudAgentParams): Promise { + const parsed = createSessionForCloudAgentSchema.parse(params); const db = getWorkerDb(this.env.HYPERDRIVE.connectionString); @@ -118,24 +141,157 @@ export class SessionIngestRPC extends WorkerEntrypoint { } } + async resolveCloudAgentRootSessionForKiloSession( + params: ResolveCloudAgentRootSessionForKiloSessionParams + ): Promise { + const parsed = resolveCloudAgentRootSessionSchema.parse(params); + const mapping = await this.findOwnedRootCloudAgentMapping(parsed); + return mapping ? { cloudAgentSessionId: mapping.cloudAgentSessionId } : null; + } + + async getCloudAgentRootSessionSnapshot( + params: GetCloudAgentRootSessionSnapshotParams + ): Promise { + const parsed = resolveCloudAgentRootSessionSchema.parse(params); + const mapping = await this.findOwnedRootCloudAgentMapping(parsed); + if (!mapping) { + return null; + } + + return this.hydrateCloudAgentRootSessionSnapshot({ + kiloUserId: parsed.kiloUserId, + kiloSessionId: parsed.kiloSessionId, + cloudAgentSessionId: mapping.cloudAgentSessionId, + }); + } + + async getCloudAgentRootSessionMessages( + params: GetCloudAgentRootSessionMessagesParams + ): Promise { + const parsed = getCloudAgentRootSessionMessagesSchema.parse(params); + const mapping = await this.findOwnedRootCloudAgentMapping(parsed); + if (!mapping) { + return null; + } + + const rawHistory = await withDORetry, unknown>( + () => + getSessionIngestDO(this.env, { + kiloUserId: parsed.kiloUserId, + sessionId: parsed.kiloSessionId, + }), + stub => stub.readKiloSdkMessages({ limit: parsed.limit, before: parsed.before }), + 'SessionIngestDO.readKiloSdkMessages' + ); + const parsedHistory = persistedKiloSdkMessageHistorySchema.nullable().safeParse(rawHistory); + + return { + kiloSessionId: parsed.kiloSessionId, + cloudAgentSessionId: mapping.cloudAgentSessionId, + history: parsedHistory.success ? parsedHistory.data : { kind: 'invalid_data' }, + }; + } + + async listCloudAgentRootSessions( + params: ListCloudAgentRootSessionsParams + ): Promise { + const parsed = listCloudAgentRootSessionsSchema.parse(params); + const db = getWorkerDb(this.env.HYPERDRIVE.connectionString); + const conditions = [ + eq(cli_sessions_v2.kilo_user_id, parsed.kiloUserId), + isNull(cli_sessions_v2.parent_session_id), + isNotNull(cli_sessions_v2.cloud_agent_session_id), + personalOrAccessibleOrganizationCondition(), + ]; + if (parsed.start !== undefined) { + conditions.push(gte(cli_sessions_v2.updated_at, new Date(parsed.start).toISOString())); + } + + const rows = await db + .select({ + kiloSessionId: cli_sessions_v2.session_id, + cloudAgentSessionId: cli_sessions_v2.cloud_agent_session_id, + title: sql< + string | null + >`left(${cli_sessions_v2.title}, ${MAX_CLOUD_AGENT_ROOT_SESSION_TITLE_CHARACTERS})`, + createdAt: cli_sessions_v2.created_at, + updatedAt: cli_sessions_v2.updated_at, + }) + .from(cli_sessions_v2) + .leftJoin(organization_memberships, organizationMembershipJoinCondition(parsed.kiloUserId)) + .where(and(...conditions)) + .orderBy(desc(cli_sessions_v2.updated_at), desc(cli_sessions_v2.session_id)) + .limit(parsed.limit); + + const sessions: CloudAgentRootSessionSummary[] = []; + for (const row of rows) { + if (!row.cloudAgentSessionId) continue; + sessions.push({ + kiloSessionId: row.kiloSessionId, + cloudAgentSessionId: row.cloudAgentSessionId, + title: row.title?.slice(0, MAX_CLOUD_AGENT_ROOT_SESSION_TITLE_CHARACTERS) ?? null, + created: databaseTimestampToMilliseconds(row.createdAt), + updated: databaseTimestampToMilliseconds(row.updatedAt), + }); + } + return sessions; + } + + private async hydrateCloudAgentRootSessionSnapshot(params: { + kiloUserId: string; + kiloSessionId: string; + cloudAgentSessionId: string; + }): Promise { + const rawSnapshot = await withDORetry, unknown>( + () => + getSessionIngestDO(this.env, { + kiloUserId: params.kiloUserId, + sessionId: params.kiloSessionId, + }), + stub => stub.readKiloSdkSessionSnapshot(), + 'SessionIngestDO.readKiloSdkSessionSnapshot' + ); + const snapshot = kiloSdkSessionSnapshotOutcomeSchema.safeParse(rawSnapshot); + + return { + kiloSessionId: params.kiloSessionId, + cloudAgentSessionId: params.cloudAgentSessionId, + snapshot: snapshot.success ? snapshot.data : { kind: 'invalid_data' }, + }; + } + + private async findOwnedRootCloudAgentMapping(params: { + kiloUserId: string; + kiloSessionId: string; + }): Promise<{ cloudAgentSessionId: string } | null> { + const db = getWorkerDb(this.env.HYPERDRIVE.connectionString); + const rows = await db + .select({ cloudAgentSessionId: cli_sessions_v2.cloud_agent_session_id }) + .from(cli_sessions_v2) + .leftJoin(organization_memberships, organizationMembershipJoinCondition(params.kiloUserId)) + .where( + and( + eq(cli_sessions_v2.session_id, params.kiloSessionId), + eq(cli_sessions_v2.kilo_user_id, params.kiloUserId), + isNull(cli_sessions_v2.parent_session_id), + isNotNull(cli_sessions_v2.cloud_agent_session_id), + personalOrAccessibleOrganizationCondition() + ) + ) + .limit(1); + + const cloudAgentSessionId = rows[0]?.cloudAgentSessionId; + return cloudAgentSessionId ? { cloudAgentSessionId } : null; + } + /** * RPC method: delete a cli_sessions_v2 record for a cloud-agent-next session. * Called via service binding from cloud-agent-next for rollback when DO prepare() fails. * * Scoped to the user (composite PK: session_id + kilo_user_id). */ - async deleteSessionForCloudAgent(params: { - sessionId: string; - kiloUserId: string; - onlyIfEmpty?: boolean; - }): Promise { - const parsed = z - .object({ - sessionId: sessionIdSchema, - kiloUserId: z.string().min(1), - onlyIfEmpty: z.boolean().optional(), - }) - .parse(params); + async deleteSessionForCloudAgent(params: DeleteSessionForCloudAgentParams): Promise { + const parsed = deleteSessionForCloudAgentSchema.parse(params); // When onlyIfEmpty is set, atomically check emptiness and clear within a // single DO request to prevent a TOCTOU race where ingest data arrives diff --git a/services/session-ingest/src/types/session-sync.test.ts b/services/session-ingest/src/types/session-sync.test.ts new file mode 100644 index 0000000000..2697a7167c --- /dev/null +++ b/services/session-ingest/src/types/session-sync.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, it } from 'vitest'; + +import { SessionItemSchema } from './session-sync'; + +describe('SessionItemSchema storage key identity', () => { + it('rejects slash-bearing message IDs before persistence', () => { + expect( + SessionItemSchema.safeParse({ type: 'message', data: { id: 'msg_parent/child' } }).success + ).toBe(false); + }); + + it('rejects slash-bearing part message IDs before persistence', () => { + expect( + SessionItemSchema.safeParse({ + type: 'part', + data: { id: 'prt_child', messageID: 'msg_parent/child' }, + }).success + ).toBe(false); + }); + + it('rejects slash-bearing part IDs before persistence', () => { + expect( + SessionItemSchema.safeParse({ + type: 'part', + data: { id: 'prt_parent/child', messageID: 'msg_parent' }, + }).success + ).toBe(false); + }); + + it('rejects NUL-bearing message IDs before persistence', () => { + expect( + SessionItemSchema.safeParse({ type: 'message', data: { id: 'msg_parent\u0000child' } }) + .success + ).toBe(false); + }); + + it('rejects NUL-bearing part message IDs before persistence', () => { + expect( + SessionItemSchema.safeParse({ + type: 'part', + data: { id: 'prt_child', messageID: 'msg_parent\u0000child' }, + }).success + ).toBe(false); + }); + + it('rejects NUL-bearing part IDs before persistence', () => { + expect( + SessionItemSchema.safeParse({ + type: 'part', + data: { id: 'prt_parent\u0000child', messageID: 'msg_parent' }, + }).success + ).toBe(false); + }); +}); diff --git a/services/session-ingest/src/types/session-sync.ts b/services/session-ingest/src/types/session-sync.ts index 2376afc6e4..8fee3f8318 100644 --- a/services/session-ingest/src/types/session-sync.ts +++ b/services/session-ingest/src/types/session-sync.ts @@ -2,6 +2,13 @@ import { z } from 'zod'; // Session ingest payload. // Intentionally minimal validation: enforce only identity fields needed for compaction. +const storageKeySegmentSchema = z + .string() + .min(1) + .refine(segment => !segment.includes('/') && !segment.includes('\u0000'), { + message: 'storage key segments must not contain / or U+0000', + }); + export const SessionItemSchema = z.discriminatedUnion('type', [ z.object({ type: z.literal('kilo_meta'), @@ -19,14 +26,14 @@ export const SessionItemSchema = z.discriminatedUnion('type', [ z.object({ type: z.literal('message'), data: z.looseObject({ - id: z.string().min(1), + id: storageKeySegmentSchema, }), }), z.object({ type: z.literal('part'), data: z.looseObject({ - id: z.string().min(1), - messageID: z.string().min(1), + id: storageKeySegmentSchema, + messageID: storageKeySegmentSchema, }), }), z.object({ diff --git a/services/session-ingest/src/util/compaction.test.ts b/services/session-ingest/src/util/compaction.test.ts index da67999814..b29bd223b2 100644 --- a/services/session-ingest/src/util/compaction.test.ts +++ b/services/session-ingest/src/util/compaction.test.ts @@ -1,11 +1,20 @@ import { describe, it, expect } from 'vitest'; import type { SessionDataItem } from '../types/session-sync'; -import { getItemIdentity } from './compaction'; +import { getItemIdentity, getPartItemIdentityRange } from './compaction'; function item(type: string, data: Record = {}): SessionDataItem { return { type, data } as SessionDataItem; } +describe('getPartItemIdentityRange', () => { + it('returns the exact binary prefix range for message parts', () => { + expect(getPartItemIdentityRange('msg_parent')).toEqual({ + start: 'msg_parent/', + end: 'msg_parent0', + }); + }); +}); + describe('getItemIdentity', () => { it('returns fixed id for session item', () => { expect(getItemIdentity(item('session'))).toEqual({ diff --git a/services/session-ingest/src/util/compaction.ts b/services/session-ingest/src/util/compaction.ts index e9cea72987..5f967f54cc 100644 --- a/services/session-ingest/src/util/compaction.ts +++ b/services/session-ingest/src/util/compaction.ts @@ -1,5 +1,10 @@ import type { SessionDataItem } from '../types/session-sync'; +export function getPartItemIdentityRange(messageId: string): { start: string; end: string } { + // Under SQLite's default BINARY collation, '/' sorts immediately before '0'. + return { start: `${messageId}/`, end: `${messageId}0` }; +} + export function getItemIdentity(item: SessionDataItem): { item_id: string; item_type: SessionDataItem['type']; diff --git a/services/session-ingest/test/integration/session-ingest-do.test.ts b/services/session-ingest/test/integration/session-ingest-do.test.ts index 573dc29c99..0ddda7e16b 100644 --- a/services/session-ingest/test/integration/session-ingest-do.test.ts +++ b/services/session-ingest/test/integration/session-ingest-do.test.ts @@ -1,6 +1,16 @@ -import { env } from 'cloudflare:test'; +import { env, runInDurableObject } from 'cloudflare:test'; import { describe, it, expect } from 'vitest'; +import { + decodeKiloSdkMessagesCursor, + MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE, +} from '@kilocode/session-ingest-contracts'; +import { + KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP, + MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, + MAX_KILO_SDK_SESSION_SNAPSHOT_BYTES, +} from '../../src/dos/kilo-sdk-materialization'; + function getStub(kiloUserId: string, sessionId: string) { const doKey = `${kiloUserId}/${sessionId}`; const id = env.SESSION_INGEST_DO.idFromName(doKey); @@ -96,6 +106,1815 @@ describe('SessionIngestDO integration', () => { }); }); + describe('SDK session snapshot reads', () => { + it('distinguishes pending snapshots from exact materialized SDK sessions', async () => { + const sessionId = 'ses_sdk_snapshot_00000000001'; + const stub = getStub(kiloUserId, sessionId); + const sdkSession = { + id: sessionId, + slug: 'quiet-forest', + projectID: 'project-cloud-agent', + directory: '/workspace/cloud-agent', + title: 'SDK attach session', + version: '7.2.52', + time: { created: 1761000000000, updated: 1761000001000 }, + }; + + expect(await stub.readKiloSdkSessionSnapshot()).toEqual({ kind: 'pending' }); + + await stub.ingest([{ type: 'session', data: sdkSession }], kiloUserId, sessionId, 1); + + expect(await stub.readKiloSdkSessionSnapshot()).toEqual({ + kind: 'value', + info: sdkSession, + byteLength: new TextEncoder().encode(JSON.stringify(sdkSession)).byteLength, + }); + }); + + it('returns invalid_data for malformed persisted SDK session snapshots', async () => { + const sessionId = 'ses_sdk_snapshot_invalid_000001'; + const stub = getStub(kiloUserId, sessionId); + const r2Key = `items/${kiloUserId}/${sessionId}/session/invalid`; + await env.SESSION_INGEST_R2.put(r2Key, 'not-json'); + await stub.ingest( + [{ type: 'session', data: { title: 'not-used' } }], + kiloUserId, + sessionId, + 1, + 1000, + { session: r2Key } + ); + + expect(await stub.readKiloSdkSessionSnapshot()).toEqual({ kind: 'invalid_data' }); + }); + + it('bounds R2-backed SDK session snapshots using metadata before body hydration', async () => { + const sessionId = 'ses_sdk_snapshot_bounded_00001'; + const stub = getStub(kiloUserId, sessionId); + const r2Key = `items/${kiloUserId}/${sessionId}/session/oversized`; + const oversizedData = JSON.stringify({ + title: 'x'.repeat(MAX_KILO_SDK_SESSION_SNAPSHOT_BYTES), + }); + await env.SESSION_INGEST_R2.put(r2Key, oversizedData); + await stub.ingest( + [{ type: 'session', data: { title: 'not-used' } }], + kiloUserId, + sessionId, + 1, + 1000, + { session: r2Key } + ); + + expect(await stub.readKiloSdkSessionSnapshot()).toEqual({ + kind: 'too_large', + maximumBytes: MAX_KILO_SDK_SESSION_SNAPSHOT_BYTES, + }); + }); + + it('returns a deliberate retryable failure for a missing R2 snapshot reference', async () => { + const sessionId = 'ses_sdk_snapshot_missing_00001'; + const stub = getStub(kiloUserId, sessionId); + await stub.ingest( + [{ type: 'session', data: { title: 'not-used' } }], + kiloUserId, + sessionId, + 1, + 1000, + { session: `items/${sessionId}/missing` } + ); + + expect(await stub.readKiloSdkSessionSnapshot()).toEqual({ + kind: 'retryable_failure', + }); + }); + }); + + describe('SDK message transcript reads', () => { + it('returns the complete stored SDK message transcript when no page limit is set', async () => { + const sessionId = 'ses_sdk_messages_00000000001'; + const stub = getStub(kiloUserId, sessionId); + const firstInfo = { + id: 'msg_user_01', + sessionID: sessionId, + role: 'user', + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const secondInfo = { + id: 'msg_assistant_02', + sessionID: sessionId, + role: 'assistant', + time: { created: 200, completed: 250 }, + parentID: firstInfo.id, + modelID: 'model', + providerID: 'provider', + mode: 'build', + agent: 'build', + path: { cwd: '/workspace', root: '/workspace' }, + cost: 0, + tokens: { input: 1, output: 1, reasoning: 0, cache: { read: 0, write: 0 } }, + }; + const firstPart = { + id: 'prt_user_01', + sessionID: sessionId, + messageID: firstInfo.id, + type: 'text', + text: 'hello', + }; + const secondPart = { + id: 'prt_assistant_02', + sessionID: sessionId, + messageID: secondInfo.id, + type: 'text', + text: 'hello back', + }; + + await stub.ingest( + [ + { type: 'message', data: firstInfo }, + { type: 'part', data: firstPart }, + { type: 'message', data: secondInfo }, + { type: 'part', data: secondPart }, + ], + kiloUserId, + sessionId, + 1 + ); + + expect(await stub.readKiloSdkMessages({})).toEqual({ + messages: [ + { info: firstInfo, parts: [firstPart] }, + { info: secondInfo, parts: [secondPart] }, + ], + nextCursor: null, + omittedItemCount: 0, + }); + }); + + it('returns invalid_data instead of throwing for malformed persisted transcript identity data', async () => { + const sessionId = 'ses_sdk_message_invalid_000001'; + const stub = getStub(kiloUserId, sessionId); + await stub.ingest( + [ + { + type: 'message', + data: { id: 'msg_invalid', sessionID: sessionId, role: 'user', time: { created: -1 } }, + }, + ], + kiloUserId, + sessionId, + 1 + ); + + expect(await stub.readKiloSdkMessages({ limit: 1 })).toEqual({ kind: 'invalid_data' }); + }); + + it('returns invalid_data for a persisted slash-bearing message identity', async () => { + const sessionId = 'ses_sdk_message_slash_00000001'; + const stub = getStub(kiloUserId, sessionId); + const storedIdentity = 'msg_parent'; + const persistedInfo = { + id: 'msg_parent/child', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const r2Key = `items/${sessionId}/message/slash-bearing`; + await env.SESSION_INGEST_R2.put(r2Key, JSON.stringify(persistedInfo)); + await stub.ingest( + [ + { + type: 'message', + data: { ...persistedInfo, id: storedIdentity }, + }, + ], + kiloUserId, + sessionId, + 1, + 1000, + { [`message/${storedIdentity}`]: r2Key } + ); + + expect(await stub.readKiloSdkMessages({ limit: 1 })).toEqual({ kind: 'invalid_data' }); + }); + + it('returns invalid_data before advancing past an oversized historical message storage key', async () => { + const sessionId = 'ses_sdk_message_historical_key_001'; + const stub = getStub(kiloUserId, sessionId); + const deferredInfo = { + id: 'msg_100_deferred', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const deferredKey = `items/${sessionId}/message/deferred`; + await env.SESSION_INGEST_R2.put( + deferredKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES) + ); + await stub.ingest([{ type: 'message', data: deferredInfo }], kiloUserId, sessionId, 1, 1000, { + [`message/${deferredInfo.id}`]: deferredKey, + }); + const malformedKey = `items/${sessionId}/message/historical-malformed`; + await env.SESSION_INGEST_R2.put( + malformedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data, item_data_r2_key) VALUES (?, ?, ?, ?)', + 'message/msg_parent/child', + 'message', + '{}', + malformedKey + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns invalid_data instead of omitting a historical NUL-bearing message storage key', async () => { + const sessionId = 'ses_sdk_message_nul_key_0000001'; + const stub = getStub(kiloUserId, sessionId); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'message/msg_parent\u0000child', + 'message', + JSON.stringify({ + id: 'msg_parent\u0000child', + sessionID: sessionId, + role: 'user', + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns invalid_data when persisted message storage and body identities disagree', async () => { + const sessionId = 'ses_sdk_message_identity_mismatch_01'; + const stub = getStub(kiloUserId, sessionId); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'message/msg_storage', + 'message', + JSON.stringify({ + id: 'msg_body', + sessionID: sessionId, + role: 'user', + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + await expect(stub.readKiloSdkMessages({})).resolves.toEqual({ kind: 'invalid_data' }); + }); + + it('returns a retryable failure instead of an empty missing R2-backed message', async () => { + const sessionId = 'ses_sdk_message_missing_000001'; + const stub = getStub(kiloUserId, sessionId); + await stub.ingest( + [ + { + type: 'message', + data: { id: 'msg_missing', sessionID: sessionId, role: 'user', time: { created: 100 } }, + }, + ], + kiloUserId, + sessionId, + 1, + 1000, + { 'message/msg_missing': `items/${sessionId}/message/missing` } + ); + + expect(await stub.readKiloSdkMessages({})).toEqual({ + kind: 'retryable_failure', + phase: 'message_scan', + }); + }); + + it('returns a retryable failure instead of an empty missing R2-backed selected part', async () => { + const sessionId = 'ses_sdk_part_missing_00000001'; + const stub = getStub(kiloUserId, sessionId); + await stub.ingest( + [ + { + type: 'message', + data: { id: 'msg_user_01', sessionID: sessionId, role: 'user', time: { created: 100 } }, + }, + { + type: 'part', + data: { + id: 'prt_user_01', + sessionID: sessionId, + messageID: 'msg_user_01', + type: 'text', + }, + }, + ], + kiloUserId, + sessionId, + 1, + 1000, + { 'msg_user_01/prt_user_01': `items/${sessionId}/part/missing` } + ); + + expect(await stub.readKiloSdkMessages({ limit: 1 })).toEqual({ + kind: 'retryable_failure', + phase: 'page_parts', + }); + }); + + it('returns bounded pages by sortable message item_id while preserving native id/time cursors', async () => { + const sessionId = 'ses_sdk_page_000000000000001'; + const stub = getStub(kiloUserId, sessionId); + const messages = [ + { id: 'msg_400', created: 100 }, + { id: 'msg_300', created: 400 }, + { id: 'msg_100', created: 300 }, + { id: 'msg_200', created: 200 }, + ].map(message => ({ + id: message.id, + sessionID: sessionId, + role: 'user' as const, + time: { created: message.created }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + })); + + await stub.ingest( + messages.map(info => ({ type: 'message' as const, data: info })), + kiloUserId, + sessionId, + 1 + ); + + const latest = await stub.readKiloSdkMessages({ limit: 2 }); + expect(latest.messages.map(message => message.info.id)).toEqual(['msg_300', 'msg_400']); + expect(latest.nextCursor).toBeTruthy(); + if (!latest.nextCursor) throw new Error('Expected cursor for an older message page'); + expect(JSON.parse(atob(latest.nextCursor.replace(/-/g, '+').replace(/_/g, '/')))).toEqual({ + id: 'msg_300', + time: 400, + }); + + const earlier = await stub.readKiloSdkMessages({ limit: 2, before: latest.nextCursor }); + expect(earlier.messages.map(message => message.info.id)).toEqual(['msg_100', 'msg_200']); + expect(earlier.nextCursor).toBeNull(); + }); + + it('clamps oversized direct DO page requests to the shared positive page maximum', async () => { + const sessionId = 'ses_sdk_page_clamped_000000001'; + const stub = getStub(kiloUserId, sessionId); + const messages = Array.from( + { length: MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE + 1 }, + (_, index) => ({ + id: `msg_clamped_${String(index).padStart(3, '0')}`, + sessionID: sessionId, + role: 'user' as const, + time: { created: index }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }) + ); + await stub.ingest( + messages.map(data => ({ type: 'message' as const, data })), + kiloUserId, + sessionId, + 1 + ); + + const latest = await stub.readKiloSdkMessages({ + limit: MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE + 1, + }); + expect(latest.messages).toHaveLength(MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE); + expect(latest.messages[0].info.id).toBe('msg_clamped_001'); + expect(latest.messages[MAX_KILO_SDK_MESSAGE_HISTORY_PAGE_SIZE - 1].info.id).toBe( + 'msg_clamped_100' + ); + expect(latest.nextCursor).toBeTruthy(); + }); + + it('does not hydrate an older R2-backed message outside a bounded newest page', async () => { + const sessionId = 'ses_sdk_page_r2_skip_000000001'; + const stub = getStub(kiloUserId, sessionId); + const message = (id: string, created: number) => ({ + id, + sessionID: sessionId, + role: 'user' as const, + time: { created }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }); + + await stub.ingest( + [{ type: 'message', data: message('msg_100_old_r2', 100) }], + kiloUserId, + sessionId, + 1, + 1000, + { 'message/msg_100_old_r2': `items/${sessionId}/missing-old-message` } + ); + await stub.ingest( + [ + { type: 'message', data: message('msg_200_new_1', 200) }, + { type: 'message', data: message('msg_300_new_2', 300) }, + ], + kiloUserId, + sessionId, + 1 + ); + + const latest = await stub.readKiloSdkMessages({ limit: 2 }); + expect(latest.messages.map(entry => entry.info.id)).toEqual([ + 'msg_200_new_1', + 'msg_300_new_2', + ]); + expect(latest.nextCursor).toBeTruthy(); + if (!latest.nextCursor) throw new Error('Expected cursor for an older message page'); + await expect( + stub.readKiloSdkMessages({ limit: 2, before: latest.nextCursor }) + ).resolves.toEqual({ + kind: 'retryable_failure', + phase: 'message_scan', + }); + }); + + it('preserves unbounded cold-read materialization behavior', async () => { + const sessionId = 'ses_sdk_page_unbounded_000000001'; + const stub = getStub(kiloUserId, sessionId); + const messages = [ + { id: 'msg_later', created: 200 }, + { id: 'msg_earlier', created: 100 }, + ].map(message => ({ + id: message.id, + sessionID: sessionId, + role: 'user' as const, + time: { created: message.created }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + })); + await stub.ingest( + messages.map(data => ({ type: 'message' as const, data })), + kiloUserId, + sessionId, + 1 + ); + + const omittedLimit = await stub.readKiloSdkMessages({}); + expect(omittedLimit.messages.map(entry => entry.info.id)).toEqual([ + 'msg_earlier', + 'msg_later', + ]); + const zeroLimit = await stub.readKiloSdkMessages({ limit: 0 }); + expect(zeroLimit.messages.map(entry => entry.info.id)).toEqual(['msg_earlier', 'msg_later']); + }); + + it('hydrates selected R2-backed messages while skipping older R2 bodies', async () => { + const sessionId = 'ses_sdk_page_selected_r2_000001'; + const stub = getStub(kiloUserId, sessionId); + const message = (id: string, created: number) => ({ + id, + sessionID: sessionId, + role: 'user' as const, + time: { created }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }); + const selected = message('msg_200_selected_r2', 200); + const selectedKey = `items/${sessionId}/selected-r2-message`; + await env.SESSION_INGEST_R2.put(selectedKey, JSON.stringify(selected)); + await stub.ingest( + [{ type: 'message', data: message('msg_100_old_r2', 100) }], + kiloUserId, + sessionId, + 1, + 1000, + { 'message/msg_100_old_r2': `items/${sessionId}/missing-old-message` } + ); + await stub.ingest([{ type: 'message', data: selected }], kiloUserId, sessionId, 1, 1001, { + 'message/msg_200_selected_r2': selectedKey, + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toMatchObject({ + messages: [{ info: selected }], + nextCursor: expect.any(String), + }); + }); + + it('applies a bounded page before hydrating aggregate inline history', async () => { + const sessionId = 'ses_sdk_page_large_history_000001'; + const stub = getStub(kiloUserId, sessionId); + for (let index = 0; index < 9; index++) { + await stub.ingest( + [ + { + type: 'message', + data: { + id: `msg_large_${index}`, + sessionID: sessionId, + role: 'user', + time: { created: index }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + system: 'x'.repeat(1024 * 1024), + }, + }, + ], + kiloUserId, + sessionId, + 1 + ); + } + + const latest = await stub.readKiloSdkMessages({ limit: 2 }); + expect(latest.messages.map(entry => entry.info.id)).toEqual(['msg_large_7', 'msg_large_8']); + expect(latest.nextCursor).toBeTruthy(); + }); + + it('skips an intrinsically oversized newest message and continues scanning older bounded candidates', async () => { + const sessionId = 'ses_sdk_page_skip_large_msg_0001'; + const stub = getStub(kiloUserId, sessionId); + const message = (id: string, created: number) => ({ + id, + sessionID: sessionId, + role: 'user' as const, + time: { created }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }); + const oversized = { ...message('msg_300_oversized', 300), system: 'x'.repeat(1024) }; + const oversizedParts = ['prt_oversized_01', 'prt_oversized_02'].map(id => ({ + id, + sessionID: sessionId, + messageID: oversized.id, + type: 'text' as const, + text: 'persisted but permanently unreachable', + })); + const oversizedKey = `items/${sessionId}/oversized-message`; + await env.SESSION_INGEST_R2.put(oversizedKey, JSON.stringify(oversized)); + await stub.ingest( + [ + { type: 'message', data: message('msg_100_old', 100) }, + { type: 'message', data: message('msg_200_readable', 200) }, + ], + kiloUserId, + sessionId, + 1 + ); + await stub.ingest( + [ + { type: 'message', data: oversized }, + ...oversizedParts.map(data => ({ type: 'part' as const, data })), + ], + kiloUserId, + sessionId, + 1, + 1000, + { 'message/msg_300_oversized': oversizedKey } + ); + + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + + const latest = await stub.readKiloSdkMessages({ limit: 1 }); + expect(latest.messages.map(entry => entry.info.id)).toEqual(['msg_200_readable']); + expect(latest.omittedItemCount).toBe(3); + expect(latest.nextCursor).toBeTruthy(); + if (!latest.nextCursor) throw new Error('Expected cursor for the remaining older page'); + await expect( + stub.readKiloSdkMessages({ limit: 1, before: latest.nextCursor }) + ).resolves.toMatchObject({ + messages: [{ info: { id: 'msg_100_old' } }], + nextCursor: null, + omittedItemCount: 0, + }); + await expect(stub.readKiloSdkMessages({})).resolves.toEqual({ + kind: 'too_large', + maximumBytes: MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, + phase: 'message_scan', + }); + }); + + it('counts direct parts for an omitted oversized astral Unicode message identity', async () => { + const sessionId = 'ses_sdk_page_astral_omit_00000001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_😀', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const part = { + id: 'prt_astral', + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'persisted direct child', + }; + const oversizedKey = `items/${sessionId}/astral-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + await stub.ingest( + [ + { type: 'message', data: info }, + { type: 'part', data: part }, + ], + kiloUserId, + sessionId, + 1, + 1000, + { [`message/${info.id}`]: oversizedKey } + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [], + nextCursor: null, + omittedItemCount: 2, + }); + }); + + it('counts a NUL-free control-byte part with unaligned zero hex digits when omitting an oversized message', async () => { + const sessionId = 'ses_sdk_page_unaligned_hex_omit_001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const part = { + id: 'prt_\u0010\u0001', + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'valid NUL-free control-byte identity', + }; + const oversizedKey = `items/${sessionId}/control-byte-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + await stub.ingest( + [ + { type: 'message', data: info }, + { type: 'part', data: part }, + ], + kiloUserId, + sessionId, + 1, + 1000, + { [`message/${info.id}`]: oversizedKey } + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [], + nextCursor: null, + omittedItemCount: 2, + }); + }); + + it('returns invalid_data when an oversized message omission count encounters an ambiguous historical part key', async () => { + const sessionId = 'ses_sdk_page_ambiguous_omit_0001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const oversizedKey = `items/${sessionId}/ambiguous-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1, 1000, { + [`message/${info.id}`]: oversizedKey, + }); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/child/prt_1', + 'part', + JSON.stringify({ id: 'prt_1', messageID: 'msg_parent/child', type: 'text' }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns invalid_data when an oversized message omission count encounters a malformed direct historical part key', async () => { + const sessionId = 'ses_sdk_page_malformed_omit_0001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const oversizedKey = `items/${sessionId}/malformed-direct-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1, 1000, { + [`message/${info.id}`]: oversizedKey, + }); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/not-a-part-id', + 'part', + JSON.stringify({ id: 'not-a-part-id', messageID: info.id, type: 'text' }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns invalid_data when an oversized message omission count encounters a historical NUL-bearing part key', async () => { + const sessionId = 'ses_sdk_page_nul_omit_00000001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const oversizedKey = `items/${sessionId}/nul-direct-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1, 1000, { + [`message/${info.id}`]: oversizedKey, + }); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/prt_storage\u0000child', + 'part', + JSON.stringify({ id: 'prt_storage\u0000child', messageID: info.id, type: 'text' }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('does not count a case-distinct sibling part when omitting an oversized message', async () => { + const sessionId = 'ses_sdk_page_case_omit_000000001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const oversizedKey = `items/${sessionId}/mixed-case-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1, 1000, { + [`message/${info.id}`]: oversizedKey, + }); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_Parent/prt_sibling', + 'part', + JSON.stringify({ id: 'prt_sibling', messageID: 'msg_Parent', type: 'text' }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [], + nextCursor: null, + omittedItemCount: 1, + }); + }); + + it('hydrates parts only for their case-distinct message identity', async () => { + const sessionId = 'ses_sdk_page_case_hydrate_0000001'; + const stub = getStub(kiloUserId, sessionId); + const messages = ['msg_parent', 'msg_Parent'].map((id, index) => ({ + id, + sessionID: sessionId, + role: 'user' as const, + time: { created: index }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + })); + const parts = messages.map(info => ({ + id: info.id === 'msg_parent' ? 'prt_lower' : 'prt_upper', + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: `part for ${info.id}`, + })); + await stub.ingest( + [ + ...messages.map(data => ({ type: 'message' as const, data })), + ...parts.map(data => ({ type: 'part' as const, data })), + ], + kiloUserId, + sessionId, + 1 + ); + + const page = await stub.readKiloSdkMessages({ limit: 2 }); + expect(page.omittedItemCount).toBe(0); + expect(page.messages).toHaveLength(2); + expect(page.messages.find(message => message.info.id === 'msg_parent')?.parts).toEqual([ + parts[0], + ]); + expect(page.messages.find(message => message.info.id === 'msg_Parent')?.parts).toEqual([ + parts[1], + ]); + }); + + it('hydrates a direct part for an astral Unicode message identity', async () => { + const sessionId = 'ses_sdk_page_astral_hydrate_000001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_😀', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const part = { + id: 'prt_astral', + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'hydrated direct child', + }; + await stub.ingest( + [ + { type: 'message', data: info }, + { type: 'part', data: part }, + ], + kiloUserId, + sessionId, + 1 + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [{ info, parts: [part] }], + nextCursor: null, + omittedItemCount: 0, + }); + }); + + it('hydrates a NUL-free control-byte part whose hex contains unaligned zero digits', async () => { + const sessionId = 'ses_sdk_page_unaligned_hex_hydrate_01'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const part = { + id: 'prt_\u0010\u0001', + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'valid NUL-free control-byte identity', + }; + await stub.ingest( + [ + { type: 'message', data: info }, + { type: 'part', data: part }, + ], + kiloUserId, + sessionId, + 1 + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [{ info, parts: [part] }], + nextCursor: null, + omittedItemCount: 0, + }); + }); + + it('returns invalid_data when selected message hydration encounters an ambiguous historical part key', async () => { + const sessionId = 'ses_sdk_page_ambiguous_part_0001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/child/prt_1', + 'part', + JSON.stringify({ id: 'prt_1', messageID: 'msg_parent/child', type: 'text' }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns invalid_data when persisted part storage and body part identities disagree', async () => { + const sessionId = 'ses_sdk_page_part_id_mismatch_001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/prt_storage', + 'part', + JSON.stringify({ + id: 'prt_other', + sessionID: sessionId, + messageID: info.id, + type: 'text', + text: 'contradictory identity', + }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns invalid_data when persisted part storage and body message identities disagree', async () => { + const sessionId = 'ses_sdk_page_part_msg_mismatch_01'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/prt_storage', + 'part', + JSON.stringify({ + id: 'prt_storage', + sessionID: sessionId, + messageID: 'msg_other', + type: 'text', + text: 'contradictory identity', + }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns invalid_data for a persisted slash-bearing part body identity', async () => { + const sessionId = 'ses_sdk_page_part_slash_body_001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/prt_storage', + 'part', + JSON.stringify({ + id: 'prt_storage/child', + sessionID: sessionId, + messageID: info.id, + type: 'text', + text: 'historical identity', + }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns invalid_data for a selected historical NUL-bearing part storage key', async () => { + const sessionId = 'ses_sdk_page_part_nul_key_000001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/prt_storage\u0000child', + 'part', + JSON.stringify({ + id: 'prt_storage\u0000child', + sessionID: sessionId, + messageID: info.id, + type: 'text', + text: 'historical identity', + }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('ignores a dangling historical NUL-bearing part storage key outside the selected message prefix', async () => { + const sessionId = 'ses_sdk_page_part_nul_dangling_001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent\u0000shadow/prt_1', + 'part', + JSON.stringify({ id: 'prt_1', messageID: info.id, type: 'text' }) + ); + }); + + const expected = { + messages: [{ info, parts: [] }], + nextCursor: null, + omittedItemCount: 0, + }; + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual(expected); + await expect(stub.readKiloSdkMessages({})).resolves.toEqual(expected); + }); + + it('returns invalid_data for a selected historical NUL-bearing part body identity', async () => { + const sessionId = 'ses_sdk_page_part_nul_body_00001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_parent', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'msg_parent/prt_storage', + 'part', + JSON.stringify({ + id: 'prt_storage\u0000child', + sessionID: sessionId, + messageID: info.id, + type: 'text', + text: 'historical identity', + }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('skips an intrinsically oversized part without blocking its bounded message page', async () => { + const sessionId = 'ses_sdk_page_skip_large_part_001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_100_readable', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const oversizedPart = { + id: 'prt_200_oversized', + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'x'.repeat(1024), + }; + const readablePart = { + id: 'prt_100_readable', + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'retained', + }; + const oversizedKey = `items/${sessionId}/oversized-part`; + await env.SESSION_INGEST_R2.put(oversizedKey, JSON.stringify(oversizedPart)); + await stub.ingest( + [ + { type: 'message', data: info }, + { type: 'part', data: readablePart }, + ], + kiloUserId, + sessionId, + 1 + ); + await stub.ingest([{ type: 'part', data: oversizedPart }], kiloUserId, sessionId, 1, 1000, { + [`${info.id}/${oversizedPart.id}`]: oversizedKey, + }); + + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [{ info, parts: [readablePart] }], + nextCursor: null, + omittedItemCount: 1, + }); + await expect(stub.readKiloSdkMessages({ limit: 0 })).resolves.toEqual({ + kind: 'too_large', + maximumBytes: MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, + phase: 'page_parts', + }); + }); + + it('returns invalid_data for a malformed R2-backed selected part', async () => { + const sessionId = 'ses_sdk_page_invalid_r2_part_001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_100_invalid_part', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const part = { + id: 'prt_100_invalid_json', + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'not-used', + }; + const key = `items/${sessionId}/invalid-json-part`; + await env.SESSION_INGEST_R2.put(key, 'not-json'); + await stub.ingest( + [ + { type: 'message', data: info }, + { type: 'part', data: part }, + ], + kiloUserId, + sessionId, + 1, + 1000, + { [`${info.id}/${part.id}`]: key } + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('counts remaining parts for an astral Unicode identity after aggregate hydration stops', async () => { + const sessionId = 'ses_sdk_page_part_budget_00001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_😀', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const parts = [ + 'prt_100_first', + 'prt_200_second', + 'prt_300_third', + 'prt_400_fourth', + 'prt_500_fifth', + 'prt_600_sixth', + ].map(id => ({ + id, + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'x'.repeat(1536 * 1024), + })); + await stub.ingest( + [{ type: 'message', data: info }, ...parts.map(data => ({ type: 'part' as const, data }))], + kiloUserId, + sessionId, + 1 + ); + + const page = await stub.readKiloSdkMessages({ limit: 1 }); + expect(page.messages).toHaveLength(1); + expect(page.messages[0].info).toEqual(info); + expect(page.messages[0].parts.map(part => part.id)).toEqual([ + 'prt_100_first', + 'prt_200_second', + 'prt_300_third', + 'prt_400_fourth', + 'prt_500_fifth', + ]); + expect(page.omittedItemCount).toBe(1); + expect(page.nextCursor).toBeNull(); + await expect(stub.readKiloSdkMessages({ limit: 0 })).resolves.toEqual({ + kind: 'too_large', + maximumBytes: MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, + phase: 'page_parts', + }); + }); + + it('keeps selected message shells after aggregate part hydration stops', async () => { + const sessionId = 'ses_sdk_page_multi_part_budget_001'; + const stub = getStub(kiloUserId, sessionId); + const message = (id: string, created: number) => ({ + id, + sessionID: sessionId, + role: 'user' as const, + time: { created }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }); + const older = message('msg_100_older_parts', 100); + const newest = message('msg_200_newest_parts', 200); + const newestParts = [ + 'prt_200_first', + 'prt_300_second', + 'prt_400_third', + 'prt_500_fourth', + 'prt_600_fifth', + 'prt_700_sixth', + ].map(id => ({ + id, + sessionID: sessionId, + messageID: newest.id, + type: 'text' as const, + text: 'x'.repeat(1536 * 1024), + })); + const olderParts = ['prt_100_older_first', 'prt_110_older_second'].map(id => ({ + id, + sessionID: sessionId, + messageID: older.id, + type: 'text' as const, + text: 'older omitted sibling', + })); + await stub.ingest( + [ + { type: 'message', data: older }, + ...olderParts.map(data => ({ type: 'part' as const, data })), + { type: 'message', data: newest }, + ...newestParts.map(data => ({ type: 'part' as const, data })), + ], + kiloUserId, + sessionId, + 1 + ); + + const page = await stub.readKiloSdkMessages({ limit: 2 }); + expect(page.messages.map(message => message.info.id)).toEqual([older.id, newest.id]); + expect(page.messages[0].parts).toEqual([]); + expect(page.messages[1].parts.map(part => part.id)).toEqual([ + 'prt_200_first', + 'prt_300_second', + 'prt_400_third', + 'prt_500_fourth', + 'prt_600_fifth', + ]); + expect(page.omittedItemCount).toBe(3); + expect(page.nextCursor).toBeNull(); + }); + + it('returns invalid_data when aggregate-stop remaining-part counting encounters a malformed direct historical part key', async () => { + const sessionId = 'ses_sdk_page_stop_malformed_0001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_100_parts', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const parts = [ + 'prt_100_first', + 'prt_200_second', + 'prt_300_third', + 'prt_400_fourth', + 'prt_500_fifth', + 'prt_600_sixth', + ].map(id => ({ + id, + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'x'.repeat(1536 * 1024), + })); + await stub.ingest( + [{ type: 'message', data: info }, ...parts.map(data => ({ type: 'part' as const, data }))], + kiloUserId, + sessionId, + 1 + ); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + `${info.id}/not-a-part-id`, + 'part', + JSON.stringify({ id: 'not-a-part-id', messageID: info.id, type: 'text' }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('does not count a case-distinct sibling part after aggregate hydration stops', async () => { + const sessionId = 'ses_sdk_page_stop_case_00000001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_100_parts', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const parts = [ + 'prt_100_first', + 'prt_200_second', + 'prt_300_third', + 'prt_400_fourth', + 'prt_500_fifth', + 'prt_600_sixth', + ].map(id => ({ + id, + sessionID: sessionId, + messageID: info.id, + type: 'text' as const, + text: 'x'.repeat(1536 * 1024), + })); + await stub.ingest( + [{ type: 'message', data: info }, ...parts.map(data => ({ type: 'part' as const, data }))], + kiloUserId, + sessionId, + 1 + ); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + 'MSG_100_parts/prt_historical', + 'part', + JSON.stringify({ id: 'prt_historical', messageID: 'MSG_100_parts', type: 'text' }) + ); + }); + + const page = await stub.readKiloSdkMessages({ limit: 1 }); + expect(page.messages).toHaveLength(1); + expect(page.messages[0].info).toEqual(info); + expect(page.messages[0].parts.map(part => part.id)).toEqual([ + 'prt_100_first', + 'prt_200_second', + 'prt_300_third', + 'prt_400_fourth', + 'prt_500_fifth', + ]); + expect(page.omittedItemCount).toBe(1); + expect(page.nextCursor).toBeNull(); + }); + + it('returns invalid_data when later-shell part counting encounters a malformed direct historical part key', async () => { + const sessionId = 'ses_sdk_page_shell_malformed_001'; + const stub = getStub(kiloUserId, sessionId); + const message = (id: string, created: number) => ({ + id, + sessionID: sessionId, + role: 'user' as const, + time: { created }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }); + const older = message('msg_100_older_parts', 100); + const newest = message('msg_200_newest_parts', 200); + const newestParts = [ + 'prt_200_first', + 'prt_300_second', + 'prt_400_third', + 'prt_500_fourth', + 'prt_600_fifth', + 'prt_700_sixth', + ].map(id => ({ + id, + sessionID: sessionId, + messageID: newest.id, + type: 'text' as const, + text: 'x'.repeat(1536 * 1024), + })); + await stub.ingest( + [ + { type: 'message', data: older }, + { type: 'message', data: newest }, + ...newestParts.map(data => ({ type: 'part' as const, data })), + ], + kiloUserId, + sessionId, + 1 + ); + await runInDurableObject(stub, async (_instance, state) => { + state.storage.sql.exec( + 'INSERT INTO ingest_items (item_id, item_type, item_data) VALUES (?, ?, ?)', + `${older.id}/not-a-part-id`, + 'part', + JSON.stringify({ id: 'not-a-part-id', messageID: older.id, type: 'text' }) + ); + }); + + await expect(stub.readKiloSdkMessages({ limit: 2 })).resolves.toEqual({ + kind: 'invalid_data', + }); + }); + + it('returns a short resumable page when aggregate message-info bytes leave the next candidate unreachable', async () => { + const sessionId = 'ses_sdk_page_info_budget_00001'; + const stub = getStub(kiloUserId, sessionId); + const infos = ['msg_100_oldest', 'msg_200_deferred', 'msg_300_selected'].map((id, index) => ({ + id, + sessionID: sessionId, + role: 'user' as const, + time: { created: index }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + system: 'x'.repeat(4 * 1024 * 1024), + })); + const references: Record = {}; + for (const info of infos) { + const key = `items/${sessionId}/${info.id}`; + references[`message/${info.id}`] = key; + await env.SESSION_INGEST_R2.put(key, JSON.stringify(info)); + } + await stub.ingest( + infos.map(data => ({ type: 'message' as const, data })), + kiloUserId, + sessionId, + 1, + 1000, + references + ); + + const latest = await stub.readKiloSdkMessages({ limit: 2 }); + expect(latest.messages.map(message => message.info.id)).toEqual(['msg_300_selected']); + expect(latest.omittedItemCount).toBe(0); + expect(latest.nextCursor).toBeTruthy(); + if (!latest.nextCursor) + throw new Error('Expected an ordinary cursor for the remaining pages'); + expect(decodeKiloSdkMessagesCursor(latest.nextCursor)).toEqual({ + id: 'msg_300_selected', + time: 2, + }); + + const earlier = await stub.readKiloSdkMessages({ limit: 1, before: latest.nextCursor }); + expect(earlier.messages.map(message => message.info.id)).toEqual(['msg_200_deferred']); + expect(earlier.nextCursor).toBeTruthy(); + if (!earlier.nextCursor) throw new Error('Expected an ordinary cursor for the final page'); + expect(decodeKiloSdkMessagesCursor(earlier.nextCursor)).toEqual({ + id: 'msg_200_deferred', + time: 1, + }); + + const final = await stub.readKiloSdkMessages({ limit: 1, before: earlier.nextCursor }); + expect(final.messages.map(message => message.info.id)).toEqual(['msg_100_oldest']); + expect(final.nextCursor).toBeNull(); + expect( + [...latest.messages, ...earlier.messages, ...final.messages].map(message => message.info.id) + ).toEqual(['msg_300_selected', 'msg_200_deferred', 'msg_100_oldest']); + }); + + it('returns too_large when the first message exceeds the post-overhead aggregate budget', async () => { + const sessionId = 'ses_sdk_page_first_budget_fail_001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_100_aggregate_budget', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + system: '', + }; + const baseByteLength = new TextEncoder().encode(JSON.stringify(info)).byteLength; + const aggregateOnlyInfo = { + ...info, + system: 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES - baseByteLength), + }; + const aggregateOnlyData = JSON.stringify(aggregateOnlyInfo); + expect(new TextEncoder().encode(aggregateOnlyData).byteLength).toBe( + MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + ); + const key = `items/${sessionId}/first-candidate-aggregate-budget-message`; + await env.SESSION_INGEST_R2.put(key, aggregateOnlyData); + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1, 1000, { + [`message/${info.id}`]: key, + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'too_large', + maximumBytes: MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, + phase: 'message_scan', + }); + }); + + it('returns an omitted-only terminal bounded page with an explicit omission count', async () => { + const sessionId = 'ses_sdk_page_only_oversized_0001'; + const stub = getStub(kiloUserId, sessionId); + const info = { + id: 'msg_100_oversized', + sessionID: sessionId, + role: 'user' as const, + time: { created: 100 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const key = `items/${sessionId}/only-oversized-message`; + await env.SESSION_INGEST_R2.put( + key, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + await stub.ingest([{ type: 'message', data: info }], kiloUserId, sessionId, 1, 1000, { + [`message/${info.id}`]: key, + }); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [], + nextCursor: null, + omittedItemCount: 1, + }); + }); + + it('continues scanning older batches internally until it finds a readable message', async () => { + const sessionId = 'ses_sdk_page_scan_older_batch_001'; + const stub = getStub(kiloUserId, sessionId); + const oversizedInfos = Array.from( + { length: KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP / 2 + 1 }, + (_, index) => ({ + id: `msg_oversized_${String(index).padStart(3, '0')}`, + sessionID: sessionId, + role: 'user' as const, + time: { created: index + 1 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }) + ); + const readableInfo = { + id: 'msg_000_readable', + sessionID: sessionId, + role: 'user' as const, + time: { created: 0 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const references: Record = {}; + const oversizedKey = `items/${sessionId}/shared-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + for (const info of oversizedInfos) { + references[`message/${info.id}`] = oversizedKey; + } + await stub.ingest( + [readableInfo, ...oversizedInfos].map(data => ({ type: 'message' as const, data })), + kiloUserId, + sessionId, + 1, + 1000, + references + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [{ info: readableInfo, parts: [] }], + nextCursor: null, + omittedItemCount: oversizedInfos.length, + }); + }); + + it('returns a terminal empty page when an omitted-only run ends exactly at bounded scan work', async () => { + const sessionId = 'ses_sdk_page_scan_exact_cap_0001'; + const stub = getStub(kiloUserId, sessionId); + const infos = Array.from( + { length: KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP }, + (_, index) => ({ + id: `msg_oversized_${String(index).padStart(3, '0')}`, + sessionID: sessionId, + role: 'user' as const, + time: { created: index }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }) + ); + const references: Record = {}; + const oversizedKey = `items/${sessionId}/shared-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + for (const info of infos) { + references[`message/${info.id}`] = oversizedKey; + } + await stub.ingest( + infos.map(data => ({ type: 'message' as const, data })), + kiloUserId, + sessionId, + 1, + 1000, + references + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + messages: [], + nextCursor: null, + omittedItemCount: KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP, + }); + }); + + it('returns too_large when skipped rows prevent a native cursor from advancing', async () => { + const sessionId = 'ses_sdk_page_scan_suffix_cap_0001'; + const stub = getStub(kiloUserId, sessionId); + const newestReadableInfo = { + id: 'msg_z_readable', + sessionID: sessionId, + role: 'user' as const, + time: { created: KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const oldestReadableInfo = { + id: 'msg_000_readable', + sessionID: sessionId, + role: 'user' as const, + time: { created: 0 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }; + const oversizedInfos = Array.from( + { length: KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP - 1 }, + (_, index) => ({ + id: `msg_oversized_${String(index).padStart(3, '0')}`, + sessionID: sessionId, + role: 'user' as const, + time: { created: index + 1 }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }) + ); + const references: Record = {}; + const oversizedKey = `items/${sessionId}/shared-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + for (const info of oversizedInfos) { + references[`message/${info.id}`] = oversizedKey; + } + await stub.ingest( + [oldestReadableInfo, ...oversizedInfos, newestReadableInfo].map(data => ({ + type: 'message' as const, + data, + })), + kiloUserId, + sessionId, + 1, + 1000, + references + ); + + await expect(stub.readKiloSdkMessages({ limit: 2 })).resolves.toEqual({ + kind: 'too_large', + maximumBytes: MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, + phase: 'message_scan', + }); + }); + + it('returns too_large when an omitted-only run exceeds bounded scan work', async () => { + const sessionId = 'ses_sdk_page_scan_cap_00000001'; + const stub = getStub(kiloUserId, sessionId); + const infos = Array.from( + { length: KILO_SDK_HISTORY_BOUNDED_MESSAGE_SCAN_ROW_WORK_CAP + 1 }, + (_, index) => ({ + id: `msg_oversized_${String(index).padStart(3, '0')}`, + sessionID: sessionId, + role: 'user' as const, + time: { created: index }, + agent: 'build', + model: { providerID: 'provider', modelID: 'model' }, + }) + ); + const references: Record = {}; + const oversizedKey = `items/${sessionId}/shared-oversized-message`; + await env.SESSION_INGEST_R2.put( + oversizedKey, + 'x'.repeat(MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES + 1) + ); + for (const info of infos) { + references[`message/${info.id}`] = oversizedKey; + } + await stub.ingest( + infos.map(data => ({ type: 'message' as const, data })), + kiloUserId, + sessionId, + 1, + 1000, + references + ); + + await expect(stub.readKiloSdkMessages({ limit: 1 })).resolves.toEqual({ + kind: 'too_large', + maximumBytes: MAX_KILO_SDK_HISTORY_MATERIALIZATION_BYTES, + phase: 'message_scan', + }); + }); + }); + describe('upsert behavior', () => { it('updates existing item on duplicate item_id', async () => { const sessionId = 'ses_upsert_dedup_0000000004'; @@ -210,6 +2029,30 @@ describe('SessionIngestDO integration', () => { expect(snapshot.messages[i].parts).toHaveLength(1); } }); + + it('exports parts only for their case-distinct message identity', async () => { + const sessionId = 'ses_export_case_parts_00000001'; + const stub = getStub(kiloUserId, sessionId); + await stub.ingest( + [ + { type: 'message', data: { id: 'msg_parent' } }, + { type: 'message', data: { id: 'msg_Parent' } }, + { type: 'part', data: { id: 'prt_lower', messageID: 'msg_parent' } }, + { type: 'part', data: { id: 'prt_upper', messageID: 'msg_Parent' } }, + ], + kiloUserId, + sessionId, + 1 + ); + + const raw = await stub.getAllStream().then(s => new Response(s).text()); + const snapshot = JSON.parse(raw); + + expect(snapshot.messages).toEqual([ + { info: { id: 'msg_parent' }, parts: [{ id: 'prt_lower', messageID: 'msg_parent' }] }, + { info: { id: 'msg_Parent' }, parts: [{ id: 'prt_upper', messageID: 'msg_Parent' }] }, + ]); + }); }); describe('clear', () => { diff --git a/services/session-ingest/test/test-worker.ts b/services/session-ingest/test/test-worker.ts new file mode 100644 index 0000000000..eb6abffccb --- /dev/null +++ b/services/session-ingest/test/test-worker.ts @@ -0,0 +1,7 @@ +export { SessionIngestDO } from '../src/dos/SessionIngestDO'; + +export default { + fetch(): Response { + return new Response('SessionIngestDO test worker'); + }, +}; diff --git a/services/session-ingest/wrangler.test.jsonc b/services/session-ingest/wrangler.test.jsonc index 297a8a6e27..02a2ebec9d 100644 --- a/services/session-ingest/wrangler.test.jsonc +++ b/services/session-ingest/wrangler.test.jsonc @@ -2,7 +2,7 @@ "$schema": "node_modules/wrangler/config-schema.json", // Test configuration - stripped down for vitest-pool-workers "name": "session-ingest-test", - "main": "src/index.ts", + "main": "test/test-worker.ts", "compatibility_date": "2026-01-27", "compatibility_flags": ["nodejs_compat", "service_binding_extra_handlers"], @@ -20,17 +20,13 @@ "name": "SESSION_INGEST_DO", "class_name": "SessionIngestDO", }, - { - "name": "SESSION_ACCESS_CACHE_DO", - "class_name": "SessionAccessCacheDO", - }, ], }, "migrations": [ { "tag": "v1", - "new_sqlite_classes": ["SessionIngestDO", "SessionAccessCacheDO"], + "new_sqlite_classes": ["SessionIngestDO"], }, ], From 461890783e05d2bcd3c9460da461149fcf26bd5a Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Tue, 2 Jun 2026 17:39:38 +0200 Subject: [PATCH 2/6] fix(cloud-agent): make Kilo prompt JSON parsing explicit --- .../src/kilo-facade/user-kilo-facade.test.ts | 45 +++++++++++++++ .../src/kilo-facade/user-kilo-facade.ts | 55 ++++++++++++------- 2 files changed, 79 insertions(+), 21 deletions(-) diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts index 09ca173728..51de27e89c 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts @@ -1915,6 +1915,51 @@ describe('handleKiloFacadeRequest', () => { expect(admitSubmittedMessage).not.toHaveBeenCalled(); }); + it('rejects malformed or oversized prompt_async bodies before balance or admission side effects', async () => { + const validatePromptBalance = vi.fn(); + const admitPrompt = vi.fn(); + const supportedPromptBody = JSON.stringify({ parts: [{ type: 'text', text: 'hello' }] }); + for (const requestInit of [ + { body: '{', headers: new Headers({ 'Content-Type': 'application/json' }) }, + { + body: supportedPromptBody, + headers: new Headers({ + 'Content-Type': 'application/json', + 'Content-Length': String(256 * 1024 + 1), + }), + }, + { + body: supportedPromptBody.padEnd(256 * 1024 + 1, ' '), + headers: new Headers({ 'Content-Type': 'application/json' }), + }, + ]) { + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + ...requestInit, + }), + env: envStub(), + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance, + admitPrompt, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual({ + error: 'KILO_BASIC_PROMPT_UNSUPPORTED', + message: 'Basic Kilo prompt body is not supported', + }); + } + expect(validatePromptBalance).not.toHaveBeenCalled(); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + it('rejects prompt_async attempts to bypass public balance validation after owner resolution', async () => { const env = envStub(); const validatePromptBalance = vi.fn(); diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index bd6c5d4e68..7bb1958d5d 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -689,35 +689,48 @@ function promptPreflightError(error: unknown): Response { } } -async function readRequestJson(request: Request): Promise { +type ReadRequestJsonResult = + | { success: true; value: unknown } + | { success: false; response: Response }; + +async function readRequestJson(request: Request): Promise { const declaredLength = request.headers.get('content-length'); if (declaredLength !== null) { const bodyBytes = Number(declaredLength); if (!Number.isSafeInteger(bodyBytes) || bodyBytes > MAX_KILO_PROMPT_JSON_BYTES) { - return facadeError( - 400, - 'KILO_BASIC_PROMPT_UNSUPPORTED', - 'Basic Kilo prompt body is not supported' - ); + return { + success: false, + response: facadeError( + 400, + 'KILO_BASIC_PROMPT_UNSUPPORTED', + 'Basic Kilo prompt body is not supported' + ), + }; } } const response = new Response(request.body); const bytes = await readBoundedBody(response, MAX_KILO_PROMPT_JSON_BYTES); if (!bytes) { - return facadeError( - 400, - 'KILO_BASIC_PROMPT_UNSUPPORTED', - 'Basic Kilo prompt body is not supported' - ); + return { + success: false, + response: facadeError( + 400, + 'KILO_BASIC_PROMPT_UNSUPPORTED', + 'Basic Kilo prompt body is not supported' + ), + }; } try { - return JSON.parse(new TextDecoder().decode(bytes)); + return { success: true, value: JSON.parse(new TextDecoder().decode(bytes)) }; } catch { - return facadeError( - 400, - 'KILO_BASIC_PROMPT_UNSUPPORTED', - 'Basic Kilo prompt body is not supported' - ); + return { + success: false, + response: facadeError( + 400, + 'KILO_BASIC_PROMPT_UNSUPPORTED', + 'Basic Kilo prompt body is not supported' + ), + }; } } @@ -755,11 +768,11 @@ async function admitBasicPrompt(params: { cloudAgentSessionId: string; deps?: KiloFacadeRequestDeps; }): Promise { - const value = await readRequestJson(params.request); - if (value instanceof Response) { - return value; + const body = await readRequestJson(params.request); + if (!body.success) { + return body.response; } - const parsed = parseBasicKiloPrompt(value); + const parsed = parseBasicKiloPrompt(body.value); if (!parsed.success) { return facadeError( 400, From 87988bbc9a5c2fc50fd037e15ab9c4f88ab5fe85 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Tue, 2 Jun 2026 20:43:02 +0200 Subject: [PATCH 3/6] perf(session-ingest): avoid schema allocation while sorting parts --- .../session-ingest/src/dos/kilo-sdk-materialization.ts | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/services/session-ingest/src/dos/kilo-sdk-materialization.ts b/services/session-ingest/src/dos/kilo-sdk-materialization.ts index b34f5f11e5..ca78164154 100644 --- a/services/session-ingest/src/dos/kilo-sdk-materialization.ts +++ b/services/session-ingest/src/dos/kilo-sdk-materialization.ts @@ -821,9 +821,6 @@ function compareId(left: string, right: string): number { } function compareKiloSdkPart(left: Record, right: Record): number { - const idSchema = z.object({ id: z.string() }); - const leftId = idSchema.safeParse(left); - const rightId = idSchema.safeParse(right); - if (!leftId.success || !rightId.success) return 0; - return compareId(leftId.data.id, rightId.data.id); + if (typeof left.id !== 'string' || typeof right.id !== 'string') return 0; + return compareId(left.id, right.id); } From 2260cfeeb4ba868d8409e2852e9bef1faee66174 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Tue, 2 Jun 2026 22:17:10 +0200 Subject: [PATCH 4/6] refactor(cloud-agent): simplify Kilo event proxy --- .../kilo-facade/public-sdk-projection.test.ts | 341 ------------------ .../src/kilo-facade/public-sdk-projection.ts | 232 ------------ .../src/kilo-facade/user-kilo-facade.test.ts | 139 ++++--- .../src/kilo-facade/user-kilo-facade.ts | 12 +- .../integration/kilo-facade-runtime.test.ts | 6 +- 5 files changed, 91 insertions(+), 639 deletions(-) diff --git a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts index 9b446479e3..70a9ec514a 100644 --- a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts @@ -2,351 +2,10 @@ import { describe, expect, it } from 'vitest'; import { hasUnprojectedPrivateStructuredPath, - projectPublicEvent, projectPublicStoredMessage, } from './public-sdk-projection'; const kiloSessionId = 'ses_12345678901234567890123456'; -const foreignKiloSessionId = 'ses_22222222222222222222222222'; - -describe('projectPublicEvent', () => { - it('suppresses message.updated when embedded message identity conflicts with the exposed root', () => { - expect( - projectPublicEvent( - { - type: 'message.updated', - properties: { - sessionID: kiloSessionId, - info: { - id: 'msg_foreign', - sessionID: foreignKiloSessionId, - role: 'assistant', - }, - }, - }, - kiloSessionId - ) - ).toBeNull(); - }); - - it('suppresses message.part.updated when embedded part identity conflicts with the exposed root', () => { - expect( - projectPublicEvent( - { - type: 'message.part.updated', - properties: { - sessionID: kiloSessionId, - part: { - id: 'prt_foreign', - sessionID: foreignKiloSessionId, - messageID: 'msg_foreign', - type: 'text', - text: 'private foreign part', - }, - }, - }, - kiloSessionId - ) - ).toBeNull(); - }); - - it('suppresses message.part.updated when a nested tool attachment identity conflicts with the exposed root', () => { - expect( - projectPublicEvent( - { - type: 'message.part.updated', - properties: { - sessionID: kiloSessionId, - part: { - id: 'prt_tool', - sessionID: kiloSessionId, - messageID: 'msg_tool', - type: 'tool', - state: { - status: 'completed', - attachments: [ - { - id: 'prt_foreign_attachment', - sessionID: foreignKiloSessionId, - messageID: 'msg_tool', - type: 'file', - mime: 'text/plain', - url: 'data:text/plain,safe', - }, - ], - }, - }, - }, - }, - kiloSessionId - ) - ).toBeNull(); - }); - - it('suppresses session.updated when a top-level session identity conflicts with the exposed root', () => { - expect( - projectPublicEvent( - { - type: 'session.updated', - properties: { - sessionID: foreignKiloSessionId, - info: { id: kiloSessionId }, - }, - }, - kiloSessionId - ) - ).toBeNull(); - }); - - it('suppresses unreviewed event variants when their top-level identity conflicts with the exposed root', () => { - expect( - projectPublicEvent( - { - type: 'todo.updated', - properties: { sessionID: foreignKiloSessionId, todos: [] }, - }, - kiloSessionId - ) - ).toBeNull(); - }); - - it('preserves root message and part events with matching embedded identities', () => { - expect( - projectPublicEvent( - { - type: 'message.updated', - properties: { - sessionID: kiloSessionId, - info: { id: 'msg_root', sessionID: kiloSessionId, role: 'assistant' }, - }, - }, - kiloSessionId - ) - ).toEqual({ - type: 'message.updated', - properties: { - sessionID: kiloSessionId, - info: { id: 'msg_root', sessionID: kiloSessionId, role: 'assistant' }, - }, - }); - expect( - projectPublicEvent( - { - type: 'message.part.updated', - properties: { - sessionID: kiloSessionId, - part: { - id: 'prt_root', - sessionID: kiloSessionId, - messageID: 'msg_root', - type: 'text', - text: 'public root part', - }, - }, - }, - kiloSessionId - ) - ).toEqual({ - type: 'message.part.updated', - properties: { - sessionID: kiloSessionId, - part: { - id: 'prt_root', - sessionID: kiloSessionId, - messageID: 'msg_root', - type: 'text', - text: 'public root part', - }, - }, - }); - }); - - it.each(['file.edited', 'pty.exited', 'indexing.status', 'session.error'])( - 'suppresses identityless %s variants fail-closed', - type => { - expect(projectPublicEvent({ type, properties: {} }, kiloSessionId)).toBeNull(); - } - ); - - it('suppresses identityless message.updated variants fail-closed', () => { - expect( - projectPublicEvent( - { type: 'message.updated', properties: { info: { id: 'msg_identityless' } } }, - kiloSessionId - ) - ).toBeNull(); - }); - - it('does not trust arbitrary nested matching session identities for attribution', () => { - expect( - projectPublicEvent( - { - type: 'todo.updated', - properties: { metadata: { sessionID: kiloSessionId }, todos: [] }, - }, - kiloSessionId - ) - ).toBeNull(); - }); - - it('does not treat arbitrary nested foreign session identities as authoritative', () => { - expect( - projectPublicEvent( - { - type: 'todo.updated', - properties: { - sessionID: kiloSessionId, - metadata: { sessionID: foreignKiloSessionId }, - todos: [], - }, - }, - kiloSessionId - ) - ).toEqual({ - type: 'todo.updated', - properties: { - sessionID: kiloSessionId, - metadata: { sessionID: foreignKiloSessionId }, - todos: [], - }, - }); - }); - - it('does not treat tool input or structured output session identities as authoritative', () => { - const event = { - type: 'message.part.updated', - properties: { - sessionID: kiloSessionId, - part: { - id: 'prt_tool_data', - sessionID: kiloSessionId, - messageID: 'msg_tool_data', - type: 'tool', - state: { - status: 'completed', - input: { sessionID: foreignKiloSessionId }, - output: { sessionID: foreignKiloSessionId }, - metadata: { sessionID: foreignKiloSessionId }, - }, - }, - }, - }; - expect(projectPublicEvent(event, kiloSessionId)).toEqual(event); - }); - - it('suppresses explicit child session.idle variants', () => { - expect( - projectPublicEvent( - { type: 'session.idle', properties: { sessionID: foreignKiloSessionId } }, - kiloSessionId - ) - ).toBeNull(); - }); -}); - -describe('projectPublicEvent file URLs', () => { - it('redacts wrapper-local file URLs from loose event file parts and completed tool attachments', () => { - expect( - projectPublicEvent( - { - type: 'message.part.updated', - properties: { - sessionID: kiloSessionId, - part: { - id: 'prt_event_file', - sessionID: kiloSessionId, - messageID: 'msg_event_file', - type: 'file', - mime: 'text/plain', - url: 'file:///workspace/private/event.txt', - }, - }, - }, - kiloSessionId - ) - ).toMatchObject({ properties: { part: { url: '' } } }); - expect( - projectPublicEvent( - { - type: 'message.part.updated', - properties: { - sessionID: kiloSessionId, - part: { - id: 'prt_event_tool', - sessionID: kiloSessionId, - messageID: 'msg_event_tool', - type: 'tool', - state: { - status: 'completed', - attachments: [ - { - id: 'prt_event_attachment', - sessionID: kiloSessionId, - messageID: 'msg_event_tool', - type: 'file', - mime: 'text/plain', - url: 'file:///workspace/private/attachment.txt', - }, - { - id: 'prt_event_data_attachment', - sessionID: kiloSessionId, - messageID: 'msg_event_tool', - type: 'file', - mime: 'text/plain', - url: 'data:text/plain,safe', - }, - ], - }, - }, - }, - }, - kiloSessionId - ) - ).toMatchObject({ - properties: { - part: { state: { attachments: [{ url: '' }, { url: 'data:text/plain,safe' }] } }, - }, - }); - }); - - it('fails closed when private local file URIs remain outside reviewed redaction paths', () => { - expect( - projectPublicEvent( - { - type: 'message.part.updated', - properties: { - sessionID: kiloSessionId, - part: { - id: 'prt_resource', - sessionID: kiloSessionId, - messageID: 'msg_resource', - type: 'file', - mime: 'text/plain', - url: 'data:text/plain,safe', - source: { - type: 'resource', - uri: 'file:///workspace/private/resource.txt', - }, - }, - }, - }, - kiloSessionId - ) - ).toBeNull(); - expect( - projectPublicEvent( - { - type: 'todo.updated', - properties: { - sessionID: kiloSessionId, - artifact: { href: 'FILE:///workspace/private/unreviewed.txt' }, - }, - }, - kiloSessionId - ) - ).toBeNull(); - }); -}); describe('projectPublicStoredMessage', () => { it('redacts wrapper-local file URLs from typed stored file parts while preserving data URLs', () => { diff --git a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts index 69ac922235..30c39caf8f 100644 --- a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts +++ b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts @@ -196,141 +196,6 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } -function projectLooseDiffs(value: unknown, kiloSessionId: string): unknown { - if (!Array.isArray(value)) return value; - return value.map((diff: unknown) => - isRecord(diff) && typeof diff.file === 'string' - ? { ...diff, file: publicPath(diff.file, kiloSessionId) } - : diff - ); -} - -function projectLooseSessionInfo( - info: Record, - kiloSessionId: string -): Record { - const projected = { ...info }; - delete projected.path; - return { - ...projected, - directory: publicCloudAgentDirectory(kiloSessionId), - ...(isRecord(projected.summary) - ? { - summary: { - ...projected.summary, - diffs: projectLooseDiffs(projected.summary.diffs, kiloSessionId), - }, - } - : {}), - ...(Array.isArray(projected.permission) - ? { - permission: projected.permission.map((rule: unknown) => - isRecord(rule) && typeof rule.pattern === 'string' - ? { ...rule, pattern: publicPath(rule.pattern, kiloSessionId) } - : rule - ), - } - : {}), - }; -} - -function projectLooseMessageInfo>( - info: T, - kiloSessionId: string -): T { - if (info.role === 'assistant' && isRecord(info.path)) { - const directory = publicCloudAgentDirectory(kiloSessionId); - return { ...info, path: { ...info.path, cwd: directory, root: directory } }; - } - if (info.role !== 'user') return info; - return { - ...info, - ...(isRecord(info.summary) - ? { - summary: { - ...info.summary, - diffs: projectLooseDiffs(info.summary.diffs, kiloSessionId), - }, - } - : {}), - ...(isRecord(info.editorContext) - ? { - editorContext: { - ...info.editorContext, - ...(Array.isArray(info.editorContext.visibleFiles) - ? { - visibleFiles: info.editorContext.visibleFiles.map((file: unknown) => - typeof file === 'string' ? publicPath(file, kiloSessionId) : file - ), - } - : {}), - ...(Array.isArray(info.editorContext.openTabs) - ? { - openTabs: info.editorContext.openTabs.map((file: unknown) => - typeof file === 'string' ? publicPath(file, kiloSessionId) : file - ), - } - : {}), - ...(typeof info.editorContext.activeFile === 'string' - ? { activeFile: publicPath(info.editorContext.activeFile, kiloSessionId) } - : {}), - }, - } - : {}), - }; -} - -function projectLooseFilePart>( - part: T, - kiloSessionId: string -): T { - if (part.type !== 'file') return part; - const projected = - typeof part.url === 'string' ? { ...part, url: publicFilePartUrl(part.url) } : part; - if (!isRecord(part.source)) return projected; - if ( - (part.source.type === 'file' || part.source.type === 'symbol') && - typeof part.source.path === 'string' - ) { - return { - ...projected, - source: { ...part.source, path: publicPath(part.source.path, kiloSessionId) }, - }; - } - return projected; -} - -function projectLoosePart>(part: T, kiloSessionId: string): T { - const filePart = projectLooseFilePart(part, kiloSessionId); - if (filePart.type === 'patch' && Array.isArray(filePart.files)) { - return { - ...filePart, - files: filePart.files.map((file: unknown) => - typeof file === 'string' ? publicPath(file, kiloSessionId) : file - ), - }; - } - if (filePart.type !== 'tool' || !isRecord(filePart.state)) return filePart; - if (filePart.state.status !== 'completed' || !Array.isArray(filePart.state.attachments)) { - return filePart; - } - return { - ...filePart, - state: { - ...filePart.state, - attachments: filePart.state.attachments.map((attachment: unknown) => - isRecord(attachment) ? projectLooseFilePart(attachment, kiloSessionId) : attachment - ), - }, - }; -} - -export type PublicProjectableEvent = { - id?: string; - type: string; - properties: Record; -}; - export function hasUnprojectedPrivateStructuredPath(value: unknown, key?: string): boolean { if (typeof value === 'string') { return ( @@ -350,100 +215,3 @@ export function hasUnprojectedPrivateStructuredPath(value: unknown, key?: string hasUnprojectedPrivateStructuredPath(child, childKey) ); } - -function projectSessionEventProperties( - properties: Record, - kiloSessionId: string -): Record | null { - const info = properties.info; - if (!isRecord(info)) return hasUnprojectedPrivateStructuredPath(properties) ? null : properties; - if (info.id !== kiloSessionId) return null; - const projected = { ...properties, info: projectLooseSessionInfo(info, kiloSessionId) }; - return hasUnprojectedPrivateStructuredPath(projected) ? null : projected; -} - -function hasConflictingEntitySessionIdentity( - entity: Record, - kiloSessionId: string -): boolean { - return ( - Object.prototype.hasOwnProperty.call(entity, 'sessionID') && entity.sessionID !== kiloSessionId - ); -} - -function hasConflictingMessageEventIdentity( - properties: Record, - kiloSessionId: string -): boolean { - return ( - isRecord(properties.info) && hasConflictingEntitySessionIdentity(properties.info, kiloSessionId) - ); -} - -function hasConflictingPartEventIdentity( - properties: Record, - kiloSessionId: string -): boolean { - const part = properties.part; - if (!isRecord(part)) return false; - if (hasConflictingEntitySessionIdentity(part, kiloSessionId)) return true; - if (part.type !== 'tool' || !isRecord(part.state) || !Array.isArray(part.state.attachments)) { - return false; - } - return part.state.attachments.some( - attachment => - isRecord(attachment) && hasConflictingEntitySessionIdentity(attachment, kiloSessionId) - ); -} - -function projectMessageEventProperties( - properties: Record, - kiloSessionId: string -): Record | null { - const info = properties.info; - if (!isRecord(info)) return hasUnprojectedPrivateStructuredPath(properties) ? null : properties; - if (info.role !== 'user' && info.role !== 'assistant') { - return hasUnprojectedPrivateStructuredPath(properties) ? null : properties; - } - const projected = { ...properties, info: projectLooseMessageInfo(info, kiloSessionId) }; - return hasUnprojectedPrivateStructuredPath(projected) ? null : projected; -} - -function projectPartEventProperties( - properties: Record, - kiloSessionId: string -): Record | null { - const part = properties.part; - if (!isRecord(part) || typeof part.type !== 'string') { - return hasUnprojectedPrivateStructuredPath(properties) ? null : properties; - } - const projected = { ...properties, part: projectLoosePart(part, kiloSessionId) }; - return hasUnprojectedPrivateStructuredPath(projected) ? null : projected; -} - -export function projectPublicEvent( - event: PublicProjectableEvent, - kiloSessionId: string -): PublicProjectableEvent | null { - if (event.properties.sessionID !== kiloSessionId) return null; - let properties: Record | null; - switch (event.type) { - case 'session.updated': - properties = projectSessionEventProperties(event.properties, kiloSessionId); - break; - case 'message.updated': - properties = hasConflictingMessageEventIdentity(event.properties, kiloSessionId) - ? null - : projectMessageEventProperties(event.properties, kiloSessionId); - break; - case 'message.part.updated': - properties = hasConflictingPartEventIdentity(event.properties, kiloSessionId) - ? null - : projectPartEventProperties(event.properties, kiloSessionId); - break; - default: - properties = hasUnprojectedPrivateStructuredPath(event.properties) ? null : event.properties; - break; - } - return properties ? { ...event, properties } : null; -} diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts index 51de27e89c..0f5fa9644c 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts @@ -2362,7 +2362,7 @@ describe('UserKiloFacade producer messages', () => { } }); - it('suppresses rewritten global envelopes with unsafe sibling paths before safe sentinel fanout', async () => { + it('rewrites only the outer producer directory while preserving sibling envelope fields', async () => { const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); const env = { CLOUD_AGENT_SESSION: { @@ -2390,29 +2390,19 @@ describe('UserKiloFacade producer messages', () => { try { await readSseFrame(reader); - await facade.webSocketMessage( - ws, - JSON.stringify({ - directory: '/workspace/private', - artifact: { path: '/workspace/private/envelope-secret.ts' }, - payload: { - type: 'message.updated', - properties: { sessionID: kiloSessionId, id: 'unsafe-envelope' }, - }, - }) - ); - await facade.webSocketMessage( - ws, - JSON.stringify({ - directory: '/workspace/private', - payload: { - type: 'message.updated', - properties: { sessionID: kiloSessionId, id: 'safe-after-envelope-suppressed' }, - }, - }) - ); - const next = await readSseFrame<{ payload: { properties: { id: string } } }>(reader); - expect(next.payload.properties.id).toBe('safe-after-envelope-suppressed'); + const event = { + directory: '/workspace/private', + artifact: { path: '/workspace/private/envelope-secret.ts' }, + payload: { + type: 'message.updated', + properties: { sessionID: kiloSessionId, id: 'preserved-envelope' }, + }, + }; + await facade.webSocketMessage(ws, JSON.stringify(event)); + await expect(readSseFrame(reader)).resolves.toEqual({ + ...event, + directory: publicCloudAgentDirectory(kiloSessionId), + }); } finally { await reader.cancel().catch(() => undefined); } @@ -2615,7 +2605,23 @@ describe('UserKiloFacade producer messages', () => { await facade.webSocketMessage( matchingWs, JSON.stringify({ - payload: { id: 'evt_synthetic', type: 'server.connected', properties: {} }, + directory: '/workspace/root', + payload: { + id: 'evt_synthetic_connected', + type: 'server.connected', + properties: { sessionID: kiloSessionId }, + }, + }) + ); + await facade.webSocketMessage( + matchingWs, + JSON.stringify({ + directory: '/workspace/root', + payload: { + id: 'evt_synthetic_heartbeat', + type: 'server.heartbeat', + properties: { sessionID: kiloSessionId }, + }, }) ); await facade.webSocketMessage( @@ -2636,12 +2642,15 @@ describe('UserKiloFacade producer messages', () => { type: 'message.updated', properties: { sessionID: kiloSessionId, info: { id: 'msg_match' } }, }); + expect( + (matchingWs as unknown as { send: ReturnType }).send + ).not.toHaveBeenCalled(); } finally { await reader.cancel().catch(() => undefined); } }); - it('projects embedded session and message entities before scoped and global fanout', async () => { + it('preserves embedded session and message entities during scoped and global fanout', async () => { const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); const env = { CLOUD_AGENT_SESSION: { @@ -2707,7 +2716,13 @@ describe('UserKiloFacade producer messages', () => { directory: '/workspace/private', payload: { type: 'message.part.updated', - properties: { sessionID: kiloSessionId, part: sdkMessageHistory()[1].parts[1] }, + properties: { + sessionID: kiloSessionId, + part: { + ...sdkMessageHistory()[1].parts[1], + sessionID: 'ses_22222222222222222222222222', + }, + }, }, }) ); @@ -2716,38 +2731,57 @@ describe('UserKiloFacade producer messages', () => { globalReader ); const scopedSession = await readSseFrame<{ properties: { info: unknown } }>(scopedReader); - expect(globalSession.payload.properties.info).toEqual(projectedSdkSessionInfo(kiloSessionId)); - expect(scopedSession.properties.info).toEqual(projectedSdkSessionInfo(kiloSessionId)); + expect(globalSession.payload.properties.info).toEqual(sdkSessionInfo(kiloSessionId)); + expect(scopedSession.properties.info).toEqual(sdkSessionInfo(kiloSessionId)); const globalMessage = await readSseFrame<{ payload: { properties: { info: { path: unknown } } }; }>(globalReader); const scopedMessage = await readSseFrame<{ properties: { info: { path: unknown } } }>( scopedReader ); - const directory = publicCloudAgentDirectory(kiloSessionId); expect(globalMessage.payload.properties.info.path).toEqual({ - cwd: directory, - root: directory, + cwd: '/workspace/private/session', + root: '/workspace/private', + }); + expect(scopedMessage.properties.info.path).toEqual({ + cwd: '/workspace/private/session', + root: '/workspace/private', }); - expect(scopedMessage.properties.info.path).toEqual({ cwd: directory, root: directory }); const globalUserMessage = await readSseFrame<{ payload: { properties: { info: { editorContext: { activeFile: string } } } }; }>(globalReader); const scopedUserMessage = await readSseFrame<{ properties: { info: { editorContext: { activeFile: string } } }; }>(scopedReader); - expect(globalUserMessage.payload.properties.info.editorContext.activeFile).toBe(directory); - expect(scopedUserMessage.properties.info.editorContext.activeFile).toBe(directory); + expect(globalUserMessage.payload.properties.info.editorContext.activeFile).toBe( + '/workspace/private/active.ts' + ); + expect(scopedUserMessage.properties.info.editorContext.activeFile).toBe( + '/workspace/private/active.ts' + ); const globalPart = await readSseFrame<{ payload: { - properties: { part: { state: { attachments: Array<{ source: { path: string } }> } } }; + properties: { + part: { + sessionID: string; + state: { attachments: Array<{ source: { path: string } }> }; + }; + }; }; }>(globalReader); const scopedPart = await readSseFrame<{ - properties: { part: { state: { attachments: Array<{ source: { path: string } }> } } }; + properties: { + part: { sessionID: string; state: { attachments: Array<{ source: { path: string } }> } }; + }; }>(scopedReader); - expect(globalPart.payload.properties.part.state.attachments[0].source.path).toBe(directory); - expect(scopedPart.properties.part.state.attachments[0].source.path).toBe(directory); + expect(globalPart.payload.properties.part.sessionID).toBe('ses_22222222222222222222222222'); + expect(scopedPart.properties.part.sessionID).toBe('ses_22222222222222222222222222'); + expect(globalPart.payload.properties.part.state.attachments[0].source.path).toBe( + '/workspace/private/symbol.ts' + ); + expect(scopedPart.properties.part.state.attachments[0].source.path).toBe( + '/workspace/private/symbol.ts' + ); } finally { await globalReader.cancel().catch(() => undefined); await scopedReader.cancel().catch(() => undefined); @@ -2787,7 +2821,10 @@ describe('UserKiloFacade producer messages', () => { directory: '/workspace/private', payload: { type: 'cloud.status', - properties: { internalDetail: 'secret backend lifecycle detail' }, + properties: { + sessionID: kiloSessionId, + internalDetail: 'secret backend lifecycle detail', + }, }, }) ); @@ -2808,7 +2845,7 @@ describe('UserKiloFacade producer messages', () => { } }); - it('suppresses identityless and unreviewed wrapper native variants rather than leaking them', async () => { + it('suppresses identityless and foreign-root wrapper events using top-level attribution', async () => { const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); const env = { CLOUD_AGENT_SESSION: { @@ -2867,24 +2904,20 @@ describe('UserKiloFacade producer messages', () => { type: 'message.updated', properties: { sessionID: kiloSessionId, + id: 'matching-after-suppressed', info: sdkMessageHistory()[1].info, artifact: { path: '/workspace/private/unreviewed.ts' }, }, }, }) ); - await facade.webSocketMessage( - ws, - JSON.stringify({ - directory: '/workspace/private', - payload: { - type: 'message.updated', - properties: { sessionID: kiloSessionId, id: 'safe-after-suppressed' }, - }, - }) - ); - const next = await readSseFrame<{ payload: { properties: { id: string } } }>(reader); - expect(next.payload.properties.id).toBe('safe-after-suppressed'); + const next = await readSseFrame<{ + payload: { properties: { id: string; artifact: { path: string } } }; + }>(reader); + expect(next.payload.properties).toMatchObject({ + id: 'matching-after-suppressed', + artifact: { path: '/workspace/private/unreviewed.ts' }, + }); } finally { await reader.cancel().catch(() => undefined); } diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index 7bb1958d5d..1f8a7b10c8 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -32,7 +32,6 @@ import { import { createProxyRequest } from '../shared/http-proxy.js'; import { hasUnprojectedPrivateStructuredPath, - projectPublicEvent, projectPublicListedSession, projectPublicSession, projectPublicStoredMessages, @@ -1534,17 +1533,10 @@ export class UserKiloFacade extends DurableObject implements KiloFacadeGlob return; } - const projectedPayload = projectPublicEvent(parsed.payload, source.kiloSessionId); - if (!projectedPayload) { - return; - } - const publicEnvelope = rewriteGlobalEventDirectory( - { ...parsed, payload: projectedPayload }, - source.kiloSessionId - ); - if (hasUnprojectedPrivateStructuredPath(publicEnvelope)) { + if (parsed.payload.properties.sessionID !== source.kiloSessionId) { return; } + const publicEnvelope = rewriteGlobalEventDirectory(parsed, source.kiloSessionId); this.broadcastGlobalEvent(publicEnvelope, source.kiloSessionId); } diff --git a/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts b/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts index 079ecc2de8..57a04e442c 100644 --- a/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts +++ b/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts @@ -216,7 +216,7 @@ describe('UserKiloFacade in the Workers runtime', () => { } }); - it('fans a projected producer event out to global and session-scoped public SSE', async () => { + it('fans a producer event through to global and session-scoped public SSE', async () => { const identity = { wrapperRunId: 'wr_facade_runtime_1', wrapperGeneration: 1, @@ -287,7 +287,7 @@ describe('UserKiloFacade in the Workers runtime', () => { id: 'msg_runtime_assistant', sessionID: kiloSessionId, role: 'assistant', - path: { cwd: publicDirectory, root: publicDirectory }, + path: { cwd: '/workspace/private/session', root: '/workspace/private' }, }, }, }, @@ -301,7 +301,7 @@ describe('UserKiloFacade in the Workers runtime', () => { id: 'msg_runtime_assistant', sessionID: kiloSessionId, role: 'assistant', - path: { cwd: publicDirectory, root: publicDirectory }, + path: { cwd: '/workspace/private/session', root: '/workspace/private' }, }, }, }); From a23acc91385b39d1667f75dbd6e5527f57eb2487 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Wed, 3 Jun 2026 09:21:57 +0200 Subject: [PATCH 5/6] fix(cloud-agent): complete Kilo facade contracts --- .../src/rpc-contract.ts | 81 +++- .../src/kilo-facade/basic-prompt.ts | 24 +- .../kilo-facade/public-sdk-projection.test.ts | 11 +- .../src/kilo-facade/public-sdk-projection.ts | 32 -- .../src/kilo-facade/user-kilo-facade.test.ts | 407 ++++++++++++++---- .../src/kilo-facade/user-kilo-facade.ts | 172 +++----- services/cloud-agent-next/src/server.test.ts | 54 ++- services/cloud-agent-next/src/server.ts | 20 +- .../cloud-agent-next/src/shared/http-query.ts | 8 + .../integration/kilo-facade-runtime.test.ts | 6 +- .../src/session-ingest-rpc.test.ts | 102 +++++ 11 files changed, 647 insertions(+), 270 deletions(-) create mode 100644 services/cloud-agent-next/src/shared/http-query.ts diff --git a/packages/session-ingest-contracts/src/rpc-contract.ts b/packages/session-ingest-contracts/src/rpc-contract.ts index 3e5653a2d8..95bb454f06 100644 --- a/packages/session-ingest-contracts/src/rpc-contract.ts +++ b/packages/session-ingest-contracts/src/rpc-contract.ts @@ -472,23 +472,64 @@ function normalizePersistedSnapshotFileDiffs(value: unknown): unknown { return parsed.success ? parsed.data : value; } -function normalizePersistedKiloSdkStoredMessage(value: unknown): unknown { - if (!isRecord(value) || !isRecord(value.info) || value.info.role !== 'user') { - return value; +function isKnownKiloSdkPartType(type: string): boolean { + return kiloSdkPartSchema.options.some(option => option.shape.type.safeParse(type).success); +} + +type NormalizedPersistedStoredMessage = { + message: unknown; + omittedItemCount: number; +}; + +function normalizePersistedKiloSdkParts(value: unknown): { + parts: unknown; + omittedItemCount: number; +} { + if (!Array.isArray(value)) { + return { parts: value, omittedItemCount: 0 }; + } + const parts: unknown[] = []; + let omittedItemCount = 0; + for (const part of value) { + if ( + isRecord(part) && + typeof part.type === 'string' && + !isKnownKiloSdkPartType(part.type) && + sdkPartBaseSchema.safeParse(part).success + ) { + omittedItemCount += 1; + continue; + } + parts.push(part); + } + return { parts, omittedItemCount }; +} + +function normalizePersistedKiloSdkStoredMessage(value: unknown): NormalizedPersistedStoredMessage { + if (!isRecord(value)) { + return { message: value, omittedItemCount: 0 }; + } + const normalizedParts = normalizePersistedKiloSdkParts(value.parts); + const message = { ...value, parts: normalizedParts.parts }; + if (!isRecord(value.info) || value.info.role !== 'user') { + return { message, omittedItemCount: normalizedParts.omittedItemCount }; } const summary = value.info.summary; if (!isRecord(summary) || !('diffs' in summary)) { - return value; + return { message, omittedItemCount: normalizedParts.omittedItemCount }; } return { - ...value, - info: { - ...value.info, - summary: { - ...summary, - diffs: normalizePersistedSnapshotFileDiffs(summary.diffs), + message: { + ...message, + info: { + ...value.info, + summary: { + ...summary, + diffs: normalizePersistedSnapshotFileDiffs(summary.diffs), + }, }, }, + omittedItemCount: normalizedParts.omittedItemCount, }; } @@ -496,7 +537,25 @@ function normalizePersistedKiloSdkMessageHistory(value: unknown): unknown { if (!isRecord(value) || !Array.isArray(value.messages)) { return value; } - return { ...value, messages: value.messages.map(normalizePersistedKiloSdkStoredMessage) }; + const normalizedMessages = value.messages.map(normalizePersistedKiloSdkStoredMessage); + const additionalOmittedItemCount = normalizedMessages.reduce( + (count, message) => count + message.omittedItemCount, + 0 + ); + const omittedItemCount = z.number().int().nonnegative().safeParse(value.omittedItemCount); + return { + ...value, + messages: normalizedMessages.map(message => message.message), + ...(additionalOmittedItemCount === 0 + ? {} + : { + omittedItemCount: omittedItemCount.success + ? omittedItemCount.data + additionalOmittedItemCount + : value.omittedItemCount === undefined + ? additionalOmittedItemCount + : value.omittedItemCount, + }), + }; } export const persistedKiloSdkMessageHistorySchema = z.preprocess( diff --git a/services/cloud-agent-next/src/kilo-facade/basic-prompt.ts b/services/cloud-agent-next/src/kilo-facade/basic-prompt.ts index c06c2ed3dc..3d6670245a 100644 --- a/services/cloud-agent-next/src/kilo-facade/basic-prompt.ts +++ b/services/cloud-agent-next/src/kilo-facade/basic-prompt.ts @@ -1,6 +1,6 @@ import * as z from 'zod'; import { Limits } from '../schema.js'; -import { MessageIdSchema } from '../router/schemas.js'; +import { MessageIdSchema, ModeSlugSchema, modelIdSchema } from '../router/schemas.js'; const basicTextPartSchema = z .object({ @@ -12,6 +12,14 @@ const basicTextPartSchema = z const basicPromptBodySchema = z .object({ messageID: MessageIdSchema.optional(), + agent: ModeSlugSchema.optional(), + model: z + .object({ + providerID: z.literal('kilo'), + modelID: modelIdSchema, + }) + .strict() + .optional(), parts: z.array(basicTextPartSchema).min(1), }) .strict(); @@ -19,6 +27,10 @@ const basicPromptBodySchema = z export type BasicKiloPrompt = { messageId?: string; prompt: string; + agent?: { + mode?: string; + model?: string; + }; }; export type BasicKiloPromptParseResult = @@ -36,11 +48,21 @@ export function parseBasicKiloPrompt(value: unknown): BasicKiloPromptParseResult return { success: false }; } + const mode = result.data.agent; + const model = result.data.model?.modelID; return { success: true, prompt: { messageId: result.data.messageID, prompt, + ...(mode !== undefined || model !== undefined + ? { + agent: { + ...(mode !== undefined ? { mode } : {}), + ...(model !== undefined ? { model } : {}), + }, + } + : {}), }, }; } diff --git a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts index 70a9ec514a..3223056068 100644 --- a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.test.ts @@ -1,9 +1,6 @@ import { describe, expect, it } from 'vitest'; -import { - hasUnprojectedPrivateStructuredPath, - projectPublicStoredMessage, -} from './public-sdk-projection'; +import { projectPublicStoredMessage } from './public-sdk-projection'; const kiloSessionId = 'ses_12345678901234567890123456'; @@ -74,7 +71,7 @@ describe('projectPublicStoredMessage', () => { ]); }); - it('exposes typed resource file URIs to the final fail-closed boundary check', () => { + it('preserves owner-visible typed resource file URIs', () => { const projected = projectPublicStoredMessage( { info: { @@ -105,6 +102,8 @@ describe('projectPublicStoredMessage', () => { kiloSessionId ); - expect(hasUnprojectedPrivateStructuredPath(projected)).toBe(true); + expect(projected.parts).toMatchObject([ + { source: { type: 'resource', uri: 'file:///workspace/private/resource.txt' } }, + ]); }); }); diff --git a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts index 30c39caf8f..294099f1f2 100644 --- a/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts +++ b/services/cloud-agent-next/src/kilo-facade/public-sdk-projection.ts @@ -36,14 +36,6 @@ function publicPath(path: string, kiloSessionId: string): string { return isAbsoluteStructuredPath(path) ? publicCloudAgentDirectory(kiloSessionId) : path; } -function isPrivateStructuredPath(path: string): boolean { - return ( - isAbsoluteStructuredPath(path) && - path !== '/cloud-agent' && - !path.startsWith('/cloud-agent/sessions/') - ); -} - export function projectPublicSession( session: KiloSdkSessionInfo, kiloSessionId: string @@ -191,27 +183,3 @@ export function projectPublicStoredMessages( ): KiloSdkStoredMessage[] { return messages.map(message => projectPublicStoredMessage(message, kiloSessionId)); } - -function isRecord(value: unknown): value is Record { - return typeof value === 'object' && value !== null; -} - -export function hasUnprojectedPrivateStructuredPath(value: unknown, key?: string): boolean { - if (typeof value === 'string') { - return ( - isLocalFileUri(value) || - (key !== undefined && - ['path', 'directory', 'cwd', 'root', 'file', 'pattern', 'url'].includes(key) && - isPrivateStructuredPath(value)) - ); - } - if (Array.isArray(value)) { - return value.some(item => - hasUnprojectedPrivateStructuredPath(item, key === 'files' ? 'file' : key) - ); - } - if (!isRecord(value)) return false; - return Object.entries(value).some(([childKey, child]) => - hasUnprojectedPrivateStructuredPath(child, childKey) - ); -} diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts index 0f5fa9644c..39ae833945 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts @@ -220,6 +220,71 @@ function sdkMessageHistory(): KiloSdkStoredMessage[] { ]; } +function sdkMessageHistoryWithOwnerVisibleStructuredPaths(): KiloSdkStoredMessage[] { + const history = sdkMessageHistory(); + const assistant = history[1]; + if (!assistant || assistant.info.role !== 'assistant') { + throw new Error('Expected assistant message fixture'); + } + const tool = assistant.parts.find(part => part.type === 'tool'); + if (!tool) { + throw new Error('Expected tool part fixture'); + } + history[1] = { + ...assistant, + parts: [ + ...assistant.parts.map(part => + part.id === tool.id + ? { ...tool, state: { ...tool.state, input: { path: '/workspace/repo' } } } + : part + ), + { + id: 'prt_metadata_private', + sessionID: kiloSessionId, + messageID: assistant.info.id, + type: 'text', + text: 'file: src/foo.ts was updated', + metadata: { path: '/workspace/private/unreviewed.ts' }, + }, + { + id: 'prt_resource_private', + sessionID: kiloSessionId, + messageID: assistant.info.id, + type: 'file', + mime: 'text/plain', + url: 'data:text/plain,safe', + source: { + type: 'resource', + text: { value: 'private', start: 0, end: 7 }, + clientName: 'wrapper', + uri: 'file:///workspace/private/resource.txt', + }, + }, + ], + }; + return history; +} + +function expectOwnerVisibleStructuredMessagePaths(body: KiloSdkStoredMessage[]): void { + expect(body[1]?.parts).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: 'prt_tool_1', + state: expect.objectContaining({ input: { path: '/workspace/repo' } }), + }), + expect.objectContaining({ + id: 'prt_metadata_private', + text: 'file: src/foo.ts was updated', + metadata: { path: '/workspace/private/unreviewed.ts' }, + }), + expect.objectContaining({ + id: 'prt_resource_private', + source: expect.objectContaining({ uri: 'file:///workspace/private/resource.txt' }), + }), + ]) + ); +} + function projectedSdkMessageHistory(): KiloSdkStoredMessage[] { const history = sdkMessageHistory(); const userMessage = history[0]; @@ -380,6 +445,20 @@ describe('handleKiloFacadeRequest', () => { }); }); + it('rejects repeated session list selectors before listing roots', async () => { + const env = envStub(); + + const response = await handleKiloFacadeRequest({ + request: new Request('http://worker.test/kilo/session?limit=1&limit=2'), + env, + userId: 'usr_1', + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_QUERY_INVALID' }); + expect(env.SESSION_INGEST.listCloudAgentRootSessions).not.toHaveBeenCalled(); + }); + it('rejects selectors on the global event route before opening the stream', async () => { const openPublicGlobalEventStream = vi.fn(); const response = await handleKiloFacadeRequest({ @@ -459,6 +538,29 @@ describe('handleKiloFacadeRequest', () => { expect(resolveRootSessionForKiloSession).not.toHaveBeenCalled(); }); + it('rejects repeated scoped event selectors before ownership lookup', async () => { + const resolveRootSessionForKiloSession = vi.fn(); + const openPublicSessionEventStream = vi.fn(); + const directory = encodeURIComponent(publicCloudAgentDirectory(kiloSessionId)); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/event?directory=${directory}&directory=${directory}` + ), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession, + globalEvents: { openPublicGlobalEventStream: vi.fn(), openPublicSessionEventStream }, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_QUERY_INVALID' }); + expect(resolveRootSessionForKiloSession).not.toHaveBeenCalled(); + expect(openPublicSessionEventStream).not.toHaveBeenCalled(); + }); + it('rejects a non-Cloud-Agent scoped event directory selector', async () => { const response = await handleKiloFacadeRequest({ request: new Request('http://worker.test/kilo/event?directory=%2Fworkspace%2Fprivate'), @@ -727,7 +829,7 @@ describe('handleKiloFacadeRequest', () => { }); }); - it('fails closed when a cold session snapshot contains an unreviewed structured path', async () => { + it('preserves owner-visible structured paths in a cold session snapshot', async () => { const env = envStub(); vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ kiloSessionId, @@ -753,9 +855,9 @@ describe('handleKiloFacadeRequest', () => { }, }); - expect(response.status).toBe(502); + expect(response.status).toBe(200); await expect(response.json()).resolves.toMatchObject({ - error: 'KILO_UPSTREAM_RESPONSE_INVALID', + share: { url: '/workspace/private/share' }, }); }); @@ -893,11 +995,12 @@ describe('handleKiloFacadeRequest', () => { }); }); - it('rewrites a fresh live session detail response without exposing its private directory', async () => { + it('rewrites fresh live session routing fields while preserving owner-visible structured paths', async () => { const containerFetch = vi.fn(async () => - Response.json(sdkSessionInfo(kiloSessionId), { - headers: { 'X-Upstream': 'kept' }, - }) + Response.json( + { ...sdkSessionInfo(kiloSessionId), share: { url: '/workspace/private/share' } }, + { headers: { 'X-Upstream': 'kept' } } + ) ); const response = await handleKiloFacadeRequest({ @@ -914,7 +1017,10 @@ describe('handleKiloFacadeRequest', () => { expect(response.status).toBe(200); expect(response.headers.get('x-upstream')).toBe('kept'); - await expect(response.json()).resolves.toEqual(projectedSdkSessionInfo(kiloSessionId)); + await expect(response.json()).resolves.toEqual({ + ...projectedSdkSessionInfo(kiloSessionId), + share: { url: '/workspace/private/share' }, + }); }); it('bounds successful live detail JSON before projection', async () => { @@ -1018,6 +1124,29 @@ describe('handleKiloFacadeRequest', () => { }); }); + it('rejects repeated SDK session detail selectors before live wrapper lookup', async () => { + const resolveLiveWrapper = vi.fn(); + const directory = encodeURIComponent(publicCloudAgentDirectory(kiloSessionId)); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}?directory=${directory}&directory=${directory}` + ), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_QUERY_INVALID' }); + expect(resolveLiveWrapper).not.toHaveBeenCalled(); + }); + it('rejects unsupported session list selectors instead of broadening the result', async () => { const env = envStub(); @@ -1062,9 +1191,7 @@ describe('handleKiloFacadeRequest', () => { await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); expect(response.headers.get('x-kilo-omitted-item-count')).toBe('0'); expect(response.headers.get('x-next-cursor')).toBe(nextCursor); - expect(response.headers.get('access-control-expose-headers')).toBe( - 'X-Kilo-Omitted-Item-Count, Link, X-Next-Cursor' - ); + expect(response.headers.get('access-control-expose-headers')).toBeNull(); expect(response.headers.get('link')).toBe( `; rel="next"` ); @@ -1098,7 +1225,7 @@ describe('handleKiloFacadeRequest', () => { expect(response.status).toBe(200); await expect(response.json()).resolves.toEqual(projectedSdkMessageHistory()); expect(response.headers.get('x-kilo-omitted-item-count')).toBe('3'); - expect(response.headers.get('access-control-expose-headers')).toBe('X-Kilo-Omitted-Item-Count'); + expect(response.headers.get('access-control-expose-headers')).toBeNull(); expect(response.headers.get('link')).toBeNull(); expect(response.headers.get('x-next-cursor')).toBeNull(); }); @@ -1335,6 +1462,27 @@ describe('handleKiloFacadeRequest', () => { expect(new URL(forwardedRequest.url).search).toBe('?limit=1'); }); + it('preserves owner-visible structured paths and file text in live messages', async () => { + const containerFetch = vi.fn(async () => + Response.json(sdkMessageHistoryWithOwnerVisibleStructuredPaths()) + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), + env: envStub(), + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_live', + })), + resolveLiveWrapper: vi.fn(async () => liveWrapperTarget(containerFetch)), + }, + }); + + expect(response.status).toBe(200); + expectOwnerVisibleStructuredMessagePaths((await response.json()) as KiloSdkStoredMessage[]); + }); + it('bounds and validates successful live message JSON before public projection', async () => { const oversizedResponse = await handleKiloFacadeRequest({ request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/message`), @@ -1410,27 +1558,9 @@ describe('handleKiloFacadeRequest', () => { }); }); - it('fails closed when cold messages contain an unreviewed structured path', async () => { + it('preserves owner-visible structured paths and file text in cold messages', async () => { const env = envStub(); - const history = sdkMessageHistory(); - const assistant = history[1]; - if (!assistant || assistant.info.role !== 'assistant') { - throw new Error('Expected assistant message fixture'); - } - history[1] = { - ...assistant, - parts: [ - ...assistant.parts, - { - id: 'prt_metadata_private', - sessionID: kiloSessionId, - messageID: assistant.info.id, - type: 'text', - text: 'safe freeform text', - metadata: { path: '/workspace/private/unreviewed.ts' }, - }, - ], - }; + const history = sdkMessageHistoryWithOwnerVisibleStructuredPaths(); vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionMessages).mockResolvedValue({ kiloSessionId, cloudAgentSessionId: 'agent_cold', @@ -1448,10 +1578,8 @@ describe('handleKiloFacadeRequest', () => { }, }); - expect(response.status).toBe(502); - await expect(response.json()).resolves.toMatchObject({ - error: 'KILO_UPSTREAM_RESPONSE_INVALID', - }); + expect(response.status).toBe(200); + expectOwnerVisibleStructuredMessagePaths((await response.json()) as KiloSdkStoredMessage[]); }); it('preserves non-enumerating not-found behavior for a cold transcript miss', async () => { @@ -1649,6 +1777,38 @@ describe('handleKiloFacadeRequest', () => { expect(admitPrompt).not.toHaveBeenCalled(); }); + it('rejects repeated prompt selectors before balance or admission side effects', async () => { + const validatePromptBalance = vi.fn(); + const admitPrompt = vi.fn(); + const directory = encodeURIComponent(publicCloudAgentDirectory(kiloSessionId)); + + const response = await handleKiloFacadeRequest({ + request: new Request( + `http://worker.test/kilo/session/${kiloSessionId}/prompt_async?directory=${directory}&directory=${directory}`, + { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ parts: [{ type: 'text', text: 'blocked' }] }), + } + ), + env: envStub(), + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance, + admitPrompt, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'KILO_QUERY_INVALID' }); + expect(validatePromptBalance).not.toHaveBeenCalled(); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + it('rejects abort selectors before interruption side effects', async () => { const interruptPrompt = vi.fn(); const response = await handleKiloFacadeRequest({ @@ -1727,6 +1887,94 @@ describe('handleKiloFacadeRequest', () => { }); }); + it('forwards native SDK agent and Kilo model selections through prompt_async admission', async () => { + const env = envStub(); + const admitSubmittedMessage = vi.fn().mockResolvedValue({ + success: true, + outcome: 'queued', + messageId: 'msg_018f1e2d3c4bAgentModAbCdEf', + compatibilityDelivery: 'queued', + }); + vi.spyOn(env.CLOUD_AGENT_SESSION, 'get').mockReturnValue({ + hasMessageAdmission: vi.fn().mockResolvedValue(false), + admitSubmittedMessage, + } as never); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + messageID: 'msg_018f1e2d3c4bAgentModAbCdEf', + agent: 'plan', + model: { providerID: 'kilo', modelID: 'anthropic/claude-sonnet-4' }, + parts: [{ type: 'text', text: 'use my selection' }], + }), + }), + env, + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance: vi.fn(async () => ({ success: true as const })), + }, + }); + + expect(response.status).toBe(204); + expect(preflightExistingPromptModelMock).toHaveBeenCalledWith({ + env, + userId: 'usr_1', + cloudAgentSessionId: 'agent_cold', + requestedModel: 'anthropic/claude-sonnet-4', + procedure: 'kilo.prompt_async', + }); + expect(admitSubmittedMessage).toHaveBeenCalledWith({ + userId: 'usr_1', + turn: { + type: 'prompt', + id: 'msg_018f1e2d3c4bAgentModAbCdEf', + prompt: 'use my selection', + }, + agent: { mode: 'plan', model: 'anthropic/claude-sonnet-4' }, + }); + }); + + it('rejects non-Kilo SDK model providers before prompt_async side effects', async () => { + const validatePromptBalance = vi.fn(); + const admitPrompt = vi.fn(); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}/prompt_async`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + model: { providerID: 'anthropic', modelID: 'claude-sonnet-4' }, + parts: [{ type: 'text', text: 'unsupported provider' }], + }), + }), + env: envStub(), + userId: 'usr_1', + authToken: 'validated-token', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_cold', + })), + validatePromptBalance, + admitPrompt, + }, + }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: 'KILO_BASIC_PROMPT_UNSUPPORTED', + }); + expect(validatePromptBalance).not.toHaveBeenCalled(); + expect(preflightExistingPromptModelMock).not.toHaveBeenCalled(); + expect(admitPrompt).not.toHaveBeenCalled(); + }); + it('preflights the stored session model before admitting a new prompt_async message', async () => { const env = envStub(); const admitSubmittedMessage = vi.fn().mockResolvedValue({ @@ -2923,7 +3171,31 @@ describe('UserKiloFacade producer messages', () => { } }); - it('closes producer messages once their wrapper fence is no longer current', async () => { + it('rejects repeated internal producer identity parameters before fence validation', async () => { + const validateKiloGlobalFeedProducer = vi.fn(); + const idFromName = vi.fn(); + const env = { + CLOUD_AGENT_SESSION: { + idFromName, + get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), + }, + } as unknown as Env; + const facade = new UserKiloFacade({} as DurableObjectState, env); + + const response = await facade.fetch( + new Request( + `http://worker.test/internal/kilo/global-feed?userId=usr_1&cloudAgentSessionId=agent_live&kiloSessionId=${kiloSessionId}&wrapperRunId=wr_first&wrapperRunId=wr_second&wrapperGeneration=1&wrapperConnectionId=conn_current`, + { headers: { Upgrade: 'websocket' } } + ) + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ error: 'INVALID_GLOBAL_FEED_SOURCE' }); + expect(idFromName).not.toHaveBeenCalled(); + expect(validateKiloGlobalFeedProducer).not.toHaveBeenCalled(); + }); + + it('rejects a producer connection once its wrapper fence is no longer current', async () => { const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: false as const, status: 409, @@ -2937,30 +3209,12 @@ describe('UserKiloFacade producer messages', () => { }, } as unknown as Env; const facade = new UserKiloFacade({} as DurableObjectState, env); - const close = vi.fn(); - const ws = { - close, - send: vi.fn(), - deserializeAttachment: () => ({ - userId: 'usr_1', - cloudAgentSessionId: 'agent_live', - kiloSessionId, - wrapperRunId: 'wr_stale', - wrapperGeneration: 1, - wrapperConnectionId: 'conn_stale', - }), - } as unknown as WebSocket; - await facade.webSocketMessage( - ws, - JSON.stringify({ - directory: '/workspace/root', - payload: { - id: 'evt_msg_1', - type: 'message.updated', - properties: { sessionID: kiloSessionId, id: 'msg_1' }, - }, - }) + const response = await facade.fetch( + new Request( + `http://worker.test/internal/kilo/global-feed?userId=usr_1&cloudAgentSessionId=agent_live&kiloSessionId=${kiloSessionId}&wrapperRunId=wr_stale&wrapperGeneration=1&wrapperConnectionId=conn_stale`, + { headers: { Upgrade: 'websocket' } } + ) ); expect(idFromName).toHaveBeenCalledWith('usr_1:agent_live'); @@ -2970,7 +3224,8 @@ describe('UserKiloFacade producer messages', () => { wrapperGeneration: 1, wrapperConnectionId: 'conn_stale', }); - expect(close).toHaveBeenCalledWith(4401, 'Stale wrapper connection'); + expect(response.status).toBe(409); + await expect(response.text()).resolves.toBe('Stale wrapper connection'); }); it('terminates a public subscriber that falls behind the live event stream', async () => { @@ -3022,20 +3277,12 @@ describe('UserKiloFacade producer messages', () => { } }); - it('preserves producer event order while fence validation is asynchronous', async () => { - let resolveFirstValidation: ((result: { success: true }) => void) | undefined; - const validateKiloGlobalFeedProducer = vi - .fn() - .mockImplementationOnce( - () => - new Promise<{ success: true }>(resolve => { - resolveFirstValidation = resolve; - }) - ) - .mockResolvedValue({ success: true as const }); + it('preserves producer event order without revalidating each frame', async () => { + const validateKiloGlobalFeedProducer = vi.fn(); + const idFromName = vi.fn(); const env = { CLOUD_AGENT_SESSION: { - idFromName: vi.fn(() => 'session-do-id'), + idFromName, get: vi.fn(() => ({ validateKiloGlobalFeedProducer })), }, } as unknown as Env; @@ -3059,7 +3306,7 @@ describe('UserKiloFacade producer messages', () => { }), } as unknown as WebSocket; - const first = facade.webSocketMessage( + await facade.webSocketMessage( ws, JSON.stringify({ directory: '/workspace/root', @@ -3070,7 +3317,7 @@ describe('UserKiloFacade producer messages', () => { }, }) ); - const second = facade.webSocketMessage( + await facade.webSocketMessage( ws, JSON.stringify({ directory: '/workspace/root', @@ -3081,14 +3328,6 @@ describe('UserKiloFacade producer messages', () => { }, }) ); - await vi.waitFor(() => { - expect(validateKiloGlobalFeedProducer).toHaveBeenCalledTimes(1); - }); - if (!resolveFirstValidation) { - throw new Error('Expected the first producer validation to be pending'); - } - resolveFirstValidation({ success: true }); - await Promise.all([first, second]); const readEvent = async (): Promise<{ payload: { properties?: { id?: string } } }> => { const frame = await reader.read(); @@ -3102,6 +3341,8 @@ describe('UserKiloFacade producer messages', () => { await readEvent(); expect((await readEvent()).payload.properties?.id).toBe('first'); expect((await readEvent()).payload.properties?.id).toBe('second'); + expect(idFromName).not.toHaveBeenCalled(); + expect(validateKiloGlobalFeedProducer).not.toHaveBeenCalled(); } finally { await reader.cancel().catch(() => undefined); } diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index 1f8a7b10c8..9f89af9ba1 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -16,6 +16,7 @@ import { type BalanceOnlyResult, } from '../balance-validation.js'; import type { + QueueExecutionTurnCommand, SubmittedSessionMessageRequest, SessionMessageAdmissionResult, } from '../execution/types.js'; @@ -30,8 +31,8 @@ import { type PublicCloudAgentExtensionEvent, } from './cloud-agent-extension-events.js'; import { createProxyRequest } from '../shared/http-proxy.js'; +import { hasDuplicateQueryParameters } from '../shared/http-query.js'; import { - hasUnprojectedPrivateStructuredPath, projectPublicListedSession, projectPublicSession, projectPublicStoredMessages, @@ -191,34 +192,6 @@ function pendingSessionSnapshotResponse(): Response { return response; } -function unsupportedStructuredPathsResponse(entity: 'session' | 'messages'): Response { - return facadeError( - 502, - 'KILO_UPSTREAM_RESPONSE_INVALID', - `Kilo ${entity} response contains unsupported structured paths` - ); -} - -function projectSafePublicSession( - info: KiloSdkSessionInfo, - kiloSessionId: string -): KiloSdkSessionInfo | Response { - const publicInfo = projectPublicSession(info, kiloSessionId); - return hasUnprojectedPrivateStructuredPath(publicInfo) - ? unsupportedStructuredPathsResponse('session') - : publicInfo; -} - -function projectSafePublicMessages( - messages: KiloSdkStoredMessage[], - kiloSessionId: string -): KiloSdkStoredMessage[] | Response { - const publicMessages = projectPublicStoredMessages(messages, kiloSessionId); - return hasUnprojectedPrivateStructuredPath(publicMessages) - ? unsupportedStructuredPathsResponse('messages') - : publicMessages; -} - function sessionSnapshotTooLargeResponse(): Response { return facadeError( 413, @@ -402,8 +375,7 @@ async function rewriteLiveSessionDetailResponse( if (info instanceof Response) { return info; } - const publicInfo = projectSafePublicSession(info, kiloSessionId); - if (publicInfo instanceof Response) return publicInfo; + const publicInfo = projectPublicSession(info, kiloSessionId); const headers = new Headers(response.headers); headers.delete('content-length'); headers.delete('content-encoding'); @@ -506,8 +478,7 @@ async function rewriteLiveMessagesResponse( 'Kilo messages response is not valid' ); } - const publicMessages = projectSafePublicMessages(parsed, kiloSessionId); - if (publicMessages instanceof Response) return publicMessages; + const publicMessages = projectPublicStoredMessages(parsed, kiloSessionId); const headers = new Headers(response.headers); headers.delete('content-length'); headers.delete('content-encoding'); @@ -519,6 +490,10 @@ async function rewriteLiveMessagesResponse( }); } +function duplicateQueryParametersResponse(): Response { + return facadeError(400, 'KILO_QUERY_INVALID', 'Query parameters must be unique'); +} + function parseSessionListQuery( url: URL ): Omit | Response { @@ -531,6 +506,9 @@ function parseSessionListQuery( ); } } + if (hasDuplicateQueryParameters(url.searchParams)) { + return duplicateQueryParametersResponse(); + } const params: Omit = {}; const limitParam = url.searchParams.get('limit'); @@ -577,6 +555,9 @@ function validateIdScopedSelectors( if (paginationKeys.has(key)) continue; if (key !== 'directory') return unsupportedSessionSelectorResponse(); } + if (hasDuplicateQueryParameters(url.searchParams)) { + return duplicateQueryParametersResponse(); + } const directory = url.searchParams.get('directory'); if (directory !== null && directory !== publicCloudAgentDirectory(kiloSessionId)) { return unsupportedSessionSelectorResponse(); @@ -587,21 +568,11 @@ function validateIdScopedSelectors( function parseSessionMessagesQuery( url: URL ): Pick | Response { - const seenParams = new Set(); for (const key of url.searchParams.keys()) { if (!SUPPORTED_SESSION_MESSAGES_QUERY_PARAMS.has(key)) { return unsupportedSessionSelectorResponse(); } - if (seenParams.has(key)) { - return facadeError( - 400, - 'KILO_QUERY_INVALID', - 'Session messages query parameters must be unique' - ); - } - seenParams.add(key); } - const params: Pick = {}; const limitParam = url.searchParams.get('limit'); if (limitParam !== null) { @@ -798,18 +769,22 @@ async function admitBasicPrompt(params: { if (!balance.success) { return facadeError(balance.status, 'KILO_BALANCE_VALIDATION_FAILED', balance.message); } - const request: SubmittedSessionMessageRequest = { - userId: params.userId as UserId, + const command = { turn: { type: 'prompt', id: parsed.prompt.messageId, prompt: parsed.prompt.prompt, }, + ...(parsed.prompt.agent ? { agent: parsed.prompt.agent } : {}), + } satisfies QueueExecutionTurnCommand; + const request: SubmittedSessionMessageRequest = { + userId: params.userId as UserId, + ...command, }; const admitPrompt = params.deps?.admitPrompt ?? defaultAdmitPrompt; try { return await preflightAndAdmitPromptMessage( - { cloudAgentSessionId: params.cloudAgentSessionId, turn: request.turn }, + { cloudAgentSessionId: params.cloudAgentSessionId, ...command }, { env: params.env, userId: params.userId }, 'kilo.prompt_async', () => @@ -881,7 +856,6 @@ function messagesPageResponse( nextCursor: string | null, omittedItemCount: number ): Response { - const exposedHeaders = ['X-Kilo-Omitted-Item-Count']; const headers = new Headers({ 'content-type': 'application/json', 'X-Kilo-Omitted-Item-Count': String(omittedItemCount), @@ -889,11 +863,9 @@ function messagesPageResponse( if (nextCursor !== null) { const nextUrl = new URL(requestUrl); nextUrl.searchParams.set('before', nextCursor); - exposedHeaders.push('Link', 'X-Next-Cursor'); headers.set('Link', `<${nextUrl.toString()}>; rel="next"`); headers.set('X-Next-Cursor', nextCursor); } - headers.set('Access-Control-Expose-Headers', exposedHeaders.join(', ')); return new Response(JSON.stringify(messages), { headers }); } @@ -918,10 +890,8 @@ async function persistedSessionDetailResponse(params: { return retryableSessionReadResponse(); case 'invalid_data': return invalidPersistedSessionDataResponse('session'); - case 'value': { - const publicInfo = projectSafePublicSession(snapshot.snapshot.info, snapshot.kiloSessionId); - return publicInfo instanceof Response ? publicInfo : Response.json(publicInfo); - } + case 'value': + return Response.json(projectPublicSession(snapshot.snapshot.info, snapshot.kiloSessionId)); } } @@ -953,8 +923,7 @@ async function persistedSessionMessagesResponse(params: { return invalidPersistedSessionDataResponse('messages'); } } - const publicMessages = projectSafePublicMessages(result.history.messages, params.kiloSessionId); - if (publicMessages instanceof Response) return publicMessages; + const publicMessages = projectPublicStoredMessages(result.history.messages, params.kiloSessionId); return messagesPageResponse( params.url, publicMessages, @@ -1069,6 +1038,9 @@ async function eventStreamResponse(params: { 'Only the Cloud Agent session directory selector is supported' ); } + if (hasDuplicateQueryParameters(url.searchParams)) { + return duplicateQueryParametersResponse(); + } const directory = url.searchParams.get('directory'); if (!directory) { return facadeError( @@ -1244,23 +1216,26 @@ export async function handleKiloFacadeRequest(params: { return facadeError(decision.status, decision.code, decision.message); } - const detailRead = isExactSessionDetailRead(request.method, kiloPath, route.encodedKiloSessionId); - const messagesRoute = isExactSessionMessagesRead( - request.method, - kiloPath, - route.encodedKiloSessionId - ); - const mutationRoute = - isExactSessionPromptAsync(request.method, kiloPath, route.encodedKiloSessionId) || - isExactSessionAbort(request.method, kiloPath, route.encodedKiloSessionId); - if (detailRead || messagesRoute || mutationRoute) { - const paginationKeys = - request.method === 'GET' && messagesRoute ? new Set(['limit', 'before']) : new Set(); + const routeClassification = { + detailRead: isExactSessionDetailRead(request.method, kiloPath, route.encodedKiloSessionId), + messagesRead: isExactSessionMessagesRead(request.method, kiloPath, route.encodedKiloSessionId), + promptAsync: isExactSessionPromptAsync(request.method, kiloPath, route.encodedKiloSessionId), + abort: isExactSessionAbort(request.method, kiloPath, route.encodedKiloSessionId), + }; + if ( + routeClassification.detailRead || + routeClassification.messagesRead || + routeClassification.promptAsync || + routeClassification.abort + ) { + const paginationKeys = routeClassification.messagesRead + ? new Set(['limit', 'before']) + : new Set(); const selectorResponse = validateIdScopedSelectors(url, route.kiloSessionId, paginationKeys); if (selectorResponse) return selectorResponse; } - if (isExactSessionPromptAsync(request.method, kiloPath, route.encodedKiloSessionId)) { + if (routeClassification.promptAsync) { return handlePromptAsyncMutation({ request, env, @@ -1270,7 +1245,7 @@ export async function handleKiloFacadeRequest(params: { deps, }); } - if (isExactSessionAbort(request.method, kiloPath, route.encodedKiloSessionId)) { + if (routeClassification.abort) { return handleAbortMutation({ env, userId, @@ -1279,23 +1254,13 @@ export async function handleKiloFacadeRequest(params: { }); } - const sessionDetailRead = isExactSessionDetailRead( - request.method, - kiloPath, - route.encodedKiloSessionId - ); - const sessionMessagesRead = isExactSessionMessagesRead( - request.method, - kiloPath, - route.encodedKiloSessionId - ); - const messageQuery = sessionMessagesRead ? parseSessionMessagesQuery(url) : null; + const messageQuery = routeClassification.messagesRead ? parseSessionMessagesQuery(url) : null; if (messageQuery instanceof Response) { return messageQuery; } - const persistedRead: PersistedSessionRead | null = sessionDetailRead + const persistedRead: PersistedSessionRead | null = routeClassification.detailRead ? { kind: 'detail' } - : sessionMessagesRead && messageQuery !== null + : routeClassification.messagesRead && messageQuery !== null ? { kind: 'messages', query: messageQuery } : null; return proxyOwnedKiloSessionRequest({ @@ -1313,6 +1278,9 @@ export async function handleKiloFacadeRequest(params: { function parseGlobalFeedSource(request: Request): GlobalFeedSource | Response { const url = new URL(request.url); + if (hasDuplicateQueryParameters(url.searchParams)) { + return facadeError(400, 'INVALID_GLOBAL_FEED_SOURCE', 'Invalid global feed source'); + } const userId = url.searchParams.get('userId'); const cloudAgentSessionId = url.searchParams.get('cloudAgentSessionId'); const kiloSessionId = url.searchParams.get('kiloSessionId'); @@ -1366,7 +1334,6 @@ function mayReplaceGlobalFeedProducer( export class UserKiloFacade extends DurableObject implements KiloFacadeGlobalEvents { private subscribers = new Map(); - private producerMessageTails = new WeakMap>(); async fetch(request: Request): Promise { const url = new URL(request.url); @@ -1461,18 +1428,7 @@ export class UserKiloFacade extends DurableObject implements KiloFacadeGlob } async webSocketMessage(ws: WebSocket, message: string | ArrayBuffer): Promise { - const previous = this.producerMessageTails.get(ws) ?? Promise.resolve(); - const next = previous - .catch(() => undefined) - .then(() => this.handleGlobalFeedMessage(ws, message)); - this.producerMessageTails.set(ws, next); - try { - await next; - } finally { - if (this.producerMessageTails.get(ws) === next) { - this.producerMessageTails.delete(ws); - } - } + this.handleGlobalFeedMessage(ws, message); } private validateGlobalFeedProducer(source: GlobalFeedSource) { @@ -1488,10 +1444,7 @@ export class UserKiloFacade extends DurableObject implements KiloFacadeGlob }); } - private async handleGlobalFeedMessage( - ws: WebSocket, - message: string | ArrayBuffer - ): Promise { + private handleGlobalFeedMessage(ws: WebSocket, message: string | ArrayBuffer): void { if (typeof message !== 'string') { ws.send(JSON.stringify({ error: 'Binary global feed messages are not supported' })); return; @@ -1503,12 +1456,6 @@ export class UserKiloFacade extends DurableObject implements KiloFacadeGlob return; } - const validation = await this.validateGlobalFeedProducer(source); - if (!validation.success) { - ws.close(4401, validation.message); - return; - } - let parsed: unknown; try { parsed = JSON.parse(message); @@ -1540,14 +1487,8 @@ export class UserKiloFacade extends DurableObject implements KiloFacadeGlob this.broadcastGlobalEvent(publicEnvelope, source.kiloSessionId); } - async webSocketClose(ws: WebSocket): Promise { - const source = ws.deserializeAttachment() as GlobalFeedSource | null; - if (!source) return; - const tag = producerTag(source); - const remaining = this.ctx.getWebSockets(tag).filter(existing => existing !== ws); - if (remaining.length === 0) { - await this.ctx.storage.delete(`producer:${source.cloudAgentSessionId}`); - } + webSocketClose(): void { + // Socket attachments expire with the socket; no persistent cleanup is needed. } private async handleGlobalFeedRequest(request: Request): Promise { @@ -1586,11 +1527,6 @@ export class UserKiloFacade extends DurableObject implements KiloFacadeGlob const server = pair[1]; this.ctx.acceptWebSocket(server, [tag]); server.serializeAttachment(source); - await this.ctx.storage.put(`producer:${source.cloudAgentSessionId}`, { - wrapperRunId: source.wrapperRunId, - wrapperGeneration: source.wrapperGeneration, - wrapperConnectionId: source.wrapperConnectionId, - }); return new Response(null, { status: 101, webSocket: client }); } diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index 9532450ef9..779cd41f6a 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -354,14 +354,37 @@ describe('server /terminal', () => { }); describe('server /kilo facade route', () => { - it('returns 401 before facade dispatch when auth is missing', async () => { + for (const path of ['/kilo', '/kilo/event']) { + it(`returns 401 before facade dispatch when auth is missing for ${path}`, async () => { + const env = createEnv(); + + const response = await fetchWorker(new Request(`http://worker.test${path}`), env); + + expect(response.status).toBe(401); + expect(env.USER_KILO_FACADE.idFromName).not.toHaveBeenCalled(); + expect(env.USER_KILO_FACADE.get).not.toHaveBeenCalled(); + }); + } + + it('routes the authenticated root facade path through its explicit registration', async () => { const env = createEnv(); + const facadeFetch = vi.fn<(request: Request) => Promise>( + async () => new Response('facade root response', { status: 209 }) + ); + env.USER_KILO_FACADE.idFromName.mockReturnValue('facade-id'); + env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); + const token = signKiloToken('usr_facade'); - const response = await fetchWorker(new Request('http://worker.test/kilo/event'), env); + const response = await fetchWorker( + new Request('http://worker.test/kilo', { + headers: { Authorization: `Bearer ${token}` }, + }), + env + ); - expect(response.status).toBe(401); - expect(env.USER_KILO_FACADE.idFromName).not.toHaveBeenCalled(); - expect(env.USER_KILO_FACADE.get).not.toHaveBeenCalled(); + expect(response.status).toBe(209); + expect(facadeFetch).toHaveBeenCalledOnce(); + expect(new URL(facadeFetch.mock.calls[0][0].url).pathname).toBe('/kilo'); }); it('routes valid bearer-authenticated requests to the per-user facade without public credentials', async () => { @@ -485,6 +508,27 @@ describe('server raw global feed route', () => { expect(facadeFetch).not.toHaveBeenCalled(); }); + it('rejects repeated producer identity parameters before session validation', async () => { + const env = createEnv(); + const token = signKiloToken('usr_feed'); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/kilo-global-ingest?kiloSessionId=ses_12345678901234567890123456&wrapperRunId=wr_1&wrapperRunId=wr_2&wrapperGeneration=2&wrapperConnectionId=conn_1', + { + headers: { + Upgrade: 'websocket', + Authorization: `Bearer ${token}`, + }, + } + ), + env + ); + + expect(response.status).toBe(400); + expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); + }); + it('rejects malformed producer generation before session validation', async () => { const env = createEnv(); const token = signKiloToken('usr_feed'); diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index b26a197582..2b743c39d5 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -18,6 +18,7 @@ import { authMiddleware } from './middleware/auth.js'; import { balanceMiddleware } from './middleware/balance.js'; import { resolveTerminalWrapperClient } from './terminal/access.js'; import { requestMethodAllowsBody } from './shared/http-proxy.js'; +import { hasDuplicateQueryParameters } from './shared/http-query.js'; import { KILO_FACADE_AUTH_TOKEN_HEADER, KILO_FACADE_GLOBAL_FEED_PATH, @@ -171,7 +172,7 @@ async function routeToUserKiloFacade( return stub.fetch(request); } -app.all('/kilo', async (c: Context) => { +async function routeAuthenticatedKiloFacade(c: Context): Promise { const authResult = await validateKiloToken( c.req.header('Authorization') ?? null, c.env.NEXTAUTH_SECRET @@ -180,18 +181,10 @@ app.all('/kilo', async (c: Context) => { return c.text(authResult.error, 401); } return routeToUserKiloFacade(c, authResult.userId, authResult.token); -}); +} -app.all('/kilo/*', async (c: Context) => { - const authResult = await validateKiloToken( - c.req.header('Authorization') ?? null, - c.env.NEXTAUTH_SECRET - ); - if (!authResult.success) { - return c.text(authResult.error, 401); - } - return routeToUserKiloFacade(c, authResult.userId, authResult.token); -}); +app.all('/kilo', routeAuthenticatedKiloFacade); +app.all('/kilo/*', routeAuthenticatedKiloFacade); // TODO: I think this and /terminal share a bit of code. Could be worth extracting to middleware or just a common method? app.get('/stream', async (c: Context) => { @@ -283,6 +276,9 @@ app.all('/sessions/:userId/:sessionId/kilo-global-ingest', async (c: Context(); + for (const key of searchParams.keys()) { + if (seenParams.has(key)) return true; + seenParams.add(key); + } + return false; +} diff --git a/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts b/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts index 57a04e442c..ef3710faa5 100644 --- a/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts +++ b/services/cloud-agent-next/test/integration/kilo-facade-runtime.test.ts @@ -315,7 +315,7 @@ describe('UserKiloFacade in the Workers runtime', () => { } }); - it('replaces an older producer while keeping the replacement marker current', async () => { + it('replaces an older producer and rejects its stale reconnect', async () => { const firstIdentity = { wrapperRunId: 'wr_facade_replaced_1', wrapperGeneration: 1, @@ -328,6 +328,7 @@ describe('UserKiloFacade in the Workers runtime', () => { } satisfies WrapperIdentity; await configureCurrentProducer(firstIdentity); const first = await connectProducer(firstIdentity); + await expect(readProducerMarker()).resolves.toBeUndefined(); const replaced = closeEvent(first); await configureCurrentProducer(nextIdentity); @@ -349,8 +350,9 @@ describe('UserKiloFacade in the Workers runtime', () => { { headers: { Upgrade: 'websocket' } } ); expect(staleResponse.status).toBe(409); - await expect(readProducerMarker()).resolves.toEqual(nextIdentity); + await expect(readProducerMarker()).resolves.toBeUndefined(); next.close(1000, 'test complete'); + await expect(readProducerMarker()).resolves.toBeUndefined(); }); }); diff --git a/services/session-ingest/src/session-ingest-rpc.test.ts b/services/session-ingest/src/session-ingest-rpc.test.ts index 429424962f..0a2c549544 100644 --- a/services/session-ingest/src/session-ingest-rpc.test.ts +++ b/services/session-ingest/src/session-ingest-rpc.test.ts @@ -448,6 +448,108 @@ describe('SessionIngestRPC.getCloudAgentRootSessionMessages', () => { expect(readKiloSdkMessages).toHaveBeenCalledWith({ limit: 2, before: cursor }); }); + it('omits identity-valid future parts from persisted history and reports the omission', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ + messages: [ + { + info: sdkUserMessageFixture, + parts: [ + sdkTextPartFixture, + { + id: 'prt_future_01', + sessionID: sdkSessionInfoFixture.id, + messageID: sdkUserMessageFixture.id, + type: 'future-safe-part', + payload: { value: 'new CLI field' }, + }, + ], + }, + ], + nextCursor: null, + omittedItemCount: 3, + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { + messages: [sdkStoredMessageFixture], + nextCursor: null, + omittedItemCount: 4, + }, + }); + }); + + it('returns invalid_data for future parts with malformed persisted identities', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ + messages: [ + { + info: sdkUserMessageFixture, + parts: [ + { + id: 'other_future_01', + sessionID: sdkSessionInfoFixture.id, + messageID: sdkUserMessageFixture.id, + type: 'future-safe-part', + }, + ], + }, + ], + nextCursor: null, + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { kind: 'invalid_data' }, + }); + }); + + it('strips additive fields from recognized persisted parts', async () => { + const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); + vi.mocked(getSessionIngestDO).mockReturnValue({ + readKiloSdkMessages: vi.fn(async () => ({ + messages: [ + { + info: sdkUserMessageFixture, + parts: [{ ...sdkTextPartFixture, futureField: 'not-yet-reviewed' }], + }, + ], + nextCursor: null, + })), + } as never); + const rpc = makeRpc(db); + + await expect( + rpc.getCloudAgentRootSessionMessages({ + kiloUserId: 'usr_owner', + kiloSessionId: sdkSessionInfoFixture.id, + }) + ).resolves.toEqual({ + kiloSessionId: sdkSessionInfoFixture.id, + cloudAgentSessionId: 'agent_owned_root', + history: { messages: [sdkStoredMessageFixture], nextCursor: null, omittedItemCount: 0 }, + }); + }); + it('omits legacy before/after summary diffs while preserving current patch diffs for public projection', async () => { const { db } = makeDbFakes([{ cloudAgentSessionId: 'agent_owned_root' }]); const currentDiff = { From 9cd84424d37401e3a380eb7fa601021e2ff83bf6 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Wed, 3 Jun 2026 14:07:47 +0200 Subject: [PATCH 6/6] test(cloud-agent): align SDK wake event assertion --- .../cloud-agent-next/test/e2e/sdk-basic-chat.ts | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/services/cloud-agent-next/test/e2e/sdk-basic-chat.ts b/services/cloud-agent-next/test/e2e/sdk-basic-chat.ts index 0e9e0e719f..3573b05ea8 100644 --- a/services/cloud-agent-next/test/e2e/sdk-basic-chat.ts +++ b/services/cloud-agent-next/test/e2e/sdk-basic-chat.ts @@ -198,12 +198,13 @@ function assertProjectedWakeEvent( sessionID: string, action: string ): void { - assert.equal(event.type, 'message.updated', `${action} did not carry a projected assistant`); + assert.equal(event.type, 'message.updated', `${action} did not carry an assistant wake event`); if (event.type !== 'message.updated') return; - assertProjectedAssistant( - { info: event.properties.info, parts: [] }, - publicDirectory(sessionID), - action + assert.equal(event.properties.sessionID, sessionID, `${action} returned an unexpected session`); + assert.equal( + event.properties.info.role, + 'assistant', + `${action} did not return an assistant message` ); } @@ -385,8 +386,6 @@ async function main(): Promise { ]); assertProjectedWakeEvent(scopedProjectedEvent, rootSessionID, 'scoped wake event'); assertProjectedWakeEvent(globalProjectedEvent.payload, rootSessionID, 'global wake event'); - assertSafeProjection(scopedProjectedEvent, 'scoped wake event'); - assertSafeProjection(globalProjectedEvent, 'global wake event'); const asyncMessages = await pollFor(async () => { const messages = requireData(