From 7c35a20745b1a49ad0dc360aeac60a2bbbdc9486 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Fri, 17 Jul 2026 21:48:51 +0200 Subject: [PATCH 1/3] fix(cli): repair cloud session imports --- .changeset/cloud-session-import.md | 6 + packages/kilo-gateway/src/cloud-sessions.ts | 286 +++++++++++++----- packages/kilo-gateway/src/index.ts | 8 +- packages/kilo-gateway/src/server/routes.ts | 8 +- .../kilo-gateway/test/cloud-sessions.test.ts | 236 ++++++++++++++- .../server/httpapi/groups/kilo-gateway.ts | 7 +- .../server/httpapi/handlers/kilo-gateway.ts | 184 +++++++---- .../server/cloud-session-import.test.ts | 273 +++++++++++++++++ .../server/kilo-gateway-statuses.test.ts | 10 +- packages/sdk/js/src/v2/gen/types.gen.ts | 8 + packages/sdk/openapi.json | 20 ++ script/check-opencode-promise-facades.ts | 1 + 12 files changed, 910 insertions(+), 137 deletions(-) create mode 100644 .changeset/cloud-session-import.md create mode 100644 packages/opencode/test/kilocode/server/cloud-session-import.test.ts diff --git a/.changeset/cloud-session-import.md b/.changeset/cloud-session-import.md new file mode 100644 index 00000000000..cc8935ae0c2 --- /dev/null +++ b/.changeset/cloud-session-import.md @@ -0,0 +1,6 @@ +--- +"@kilocode/cli": patch +"@kilocode/kilo-gateway": patch +--- + +Allow forking Cloud Agent sessions into the local CLI. diff --git a/packages/kilo-gateway/src/cloud-sessions.ts b/packages/kilo-gateway/src/cloud-sessions.ts index 2752073db19..444634254cf 100644 --- a/packages/kilo-gateway/src/cloud-sessions.ts +++ b/packages/kilo-gateway/src/cloud-sessions.ts @@ -6,6 +6,33 @@ export interface DrizzleDb { insert(table: object): { values(data: object): { onConflictDoNothing(): { run(): void } } } } +type Export = { + info: Record & { + time?: { + created?: number + updated?: number + compacting?: number + archived?: number + } + } + messages?: unknown +} + +export interface PrepareDeps { + Instance: { + readonly directory: string + readonly project: { readonly id: string } + } + readonly workspaceID?: string + readonly path?: string + Identifier: { + ascending(prefix: "session" | "message" | "part", given?: string): string + descending(prefix: "session" | "message" | "part", given?: string): string + } +} + +export class SessionImportValidationError extends Error {} + const INGEST_BASE = process.env.KILO_SESSION_INGEST_URL ?? "https://ingest.kilosessions.ai" const TIMEOUT = 30_000 @@ -56,43 +83,118 @@ export async function fetchCloudSessionForImport(token: string, sessionId: strin return { ok: true, data } } -export interface ImportDeps { +export interface ImportDeps extends PrepareDeps { Database: { transaction(callback: (db: DrizzleDb) => T): T effect(fn: () => void | Promise): void } - Instance: { - readonly directory: string - readonly project: { readonly id: string } - } SessionTable: object MessageTable: object PartTable: object SessionToRow: (info: any) => Record Bus: { publish(event: { type: string; properties: unknown }, payload: unknown): void | Promise } SessionCreatedEvent: { type: string; properties: unknown } - Identifier: { - ascending(prefix: "session" | "message" | "part", given?: string): string - descending(prefix: "session" | "message" | "part", given?: string): string - } } -export function importSessionToDb(data: any, deps: ImportDeps) { - const { - Database, - Instance, - SessionTable, - MessageTable, - PartTable, - SessionToRow, - Bus, - SessionCreatedEvent, - Identifier, - } = deps - - const localSessionID = Identifier.descending("session") - const msgMap = new Map() - const projectID = Instance.project.id +function record(input: unknown): input is Record { + return typeof input === "object" && input !== null && !Array.isArray(input) +} + +export function prepareSessionImport(data: Export, deps: PrepareDeps) { + if (!record(data) || !record(data.info) || typeof data.info.id !== "string") + throw new SessionImportValidationError("Invalid session info") + if (!Array.isArray(data.messages)) throw new SessionImportValidationError("Invalid session messages") + + const sessionID = deps.Identifier.descending("session") + const ids = new Map() + const pids = new Map() + const items: Array<{ + id: string + created: number + info: Record + parent?: string + parts: Array<{ + id: string + data: Record + tail?: string + attachments?: Array<{ id: string; data: Record }> + }> + }> = [] + + for (const msg of data.messages) { + if (!record(msg) || !record(msg.info) || !Array.isArray(msg.parts)) + throw new SessionImportValidationError("Invalid message") + const id = msg.info.id + const role = msg.info.role + const time = msg.info.time + if ( + typeof id !== "string" || + msg.info.sessionID !== data.info.id || + (role !== "user" && role !== "assistant") || + !record(time) || + typeof time.created !== "number" || + !Number.isFinite(time.created) + ) + throw new SessionImportValidationError("Invalid message info") + const parent = msg.info.parentID + if (parent !== undefined && typeof parent !== "string") + throw new SessionImportValidationError("Invalid message parent") + if (ids.has(id)) throw new SessionImportValidationError("Duplicate message ID") + ids.set(id, deps.Identifier.ascending("message")) + + const parts: Array<{ + id: string + data: Record + tail?: string + attachments?: Array<{ id: string; data: Record }> + }> = [] + for (const part of msg.parts) { + if ( + !record(part) || + typeof part.id !== "string" || + part.sessionID !== data.info.id || + part.messageID !== id || + typeof part.type !== "string" + ) + throw new SessionImportValidationError("Invalid message part") + const tail = part.type === "compaction" ? part.tail_start_id : undefined + if (tail !== undefined && typeof tail !== "string") + throw new SessionImportValidationError("Invalid compaction tail") + if (pids.has(part.id)) throw new SessionImportValidationError("Duplicate part ID") + pids.set(part.id, deps.Identifier.ascending("part")) + + const state = + part.type === "tool" && record(part.state) && part.state.status === "completed" ? part.state : undefined + const attachments = (() => { + if (!state || state.attachments === undefined) return + if (!Array.isArray(state.attachments)) throw new SessionImportValidationError("Invalid tool attachments") + const result: Array<{ id: string; data: Record }> = [] + for (const attachment of state.attachments) { + if ( + !record(attachment) || + typeof attachment.id !== "string" || + attachment.sessionID !== data.info.id || + attachment.messageID !== id || + attachment.type !== "file" || + typeof attachment.mime !== "string" || + typeof attachment.url !== "string" + ) + throw new SessionImportValidationError("Invalid tool attachment") + if (pids.has(attachment.id)) throw new SessionImportValidationError("Duplicate part ID") + pids.set(attachment.id, deps.Identifier.ascending("part")) + result.push({ id: attachment.id, data: attachment }) + } + return result + })() + parts.push({ + id: part.id, + data: part, + ...(tail !== undefined ? { tail } : {}), + ...(attachments !== undefined ? { attachments } : {}), + }) + } + items.push({ id, created: time.created, info: msg.info, parts, ...(parent !== undefined ? { parent } : {}) }) + } const now = Date.now() const time = { @@ -102,60 +204,104 @@ export function importSessionToDb(data: any, deps: ImportDeps) { ...(data.info.time?.archived !== undefined && { archived: data.info.time.archived }), } - const info = { + const info: Record & { + id: string + projectID: string + directory: string + time: typeof time + } = { ...data.info, - id: localSessionID, - projectID, + id: sessionID, + projectID: deps.Instance.project.id, slug: data.info.slug, - directory: Instance.directory, + directory: deps.Instance.directory, version: data.info.version, time, } + delete info.workspaceID + delete info.path + if (deps.workspaceID !== undefined) info.workspaceID = deps.workspaceID + if (deps.path !== undefined) info.path = deps.path + delete info.parentID + delete info.share + delete info.revert - Database.transaction((db) => { - db.insert(SessionTable) - .values(SessionToRow(info as Record)) - .onConflictDoNothing() - .run() - - const messages = Array.isArray(data.messages) ? data.messages : [] - for (const msg of messages.filter((m: any) => m.info)) { - const msgID = Identifier.ascending("message") - msgMap.set(msg.info.id, msgID) - msg.info.id = msgID - msg.info.sessionID = localSessionID - if (msg.info.parentID) msg.info.parentID = msgMap.get(msg.info.parentID) ?? msg.info.parentID - - db.insert(MessageTable) - .values({ - id: msgID, - session_id: localSessionID, - time_created: msg.info.time?.created ?? Date.now(), - data: msg.info, - }) - .onConflictDoNothing() - .run() - - for (const part of msg.parts ?? []) { - const partID = Identifier.ascending("part") - part.id = partID - part.messageID = msgID - part.sessionID = localSessionID - - db.insert(PartTable) - .values({ - id: partID, - message_id: msgID, - session_id: localSessionID, - data: part, - }) - .onConflictDoNothing() - .run() + const messages: Array<{ + id: string + session_id: string + time_created: number + data: Record + }> = [] + const parts: Array<{ + id: string + message_id: string + session_id: string + data: Record + }> = [] + for (const item of items) { + const id = ids.get(item.id) + if (!id) throw new SessionImportValidationError("Missing message ID") + const parentID = item.parent === undefined ? undefined : ids.get(item.parent) + if (item.parent !== undefined && !parentID) throw new SessionImportValidationError("Dangling message parent") + const next = { + ...item.info, + id, + sessionID, + ...(parentID ? { parentID } : {}), + } + messages.push({ id, session_id: sessionID, time_created: item.created, data: next }) + + for (const part of item.parts) { + const partID = pids.get(part.id) + if (!partID) throw new SessionImportValidationError("Missing part ID") + const tail = part.tail === undefined ? undefined : ids.get(part.tail) + if (part.tail !== undefined && !tail) throw new SessionImportValidationError("Dangling compaction tail") + const data: Record = { + ...part.data, + id: partID, + messageID: id, + sessionID, + ...(tail ? { tail_start_id: tail } : {}), + } + if (part.attachments) { + const state = part.data.state + if (!record(state)) throw new SessionImportValidationError("Invalid tool state") + data.state = { + ...state, + attachments: part.attachments.map((attachment) => { + const attachmentID = pids.get(attachment.id) + if (!attachmentID) throw new SessionImportValidationError("Missing attachment ID") + return { ...attachment.data, id: attachmentID, messageID: id, sessionID } + }), + } } + parts.push({ + id: partID, + message_id: id, + session_id: sessionID, + data, + }) + } + } + + return { info, messages, parts } +} + +export function importSessionToDb(data: Export, deps: ImportDeps) { + const prepared = prepareSessionImport(data, deps) + + deps.Database.transaction((db) => { + db.insert(deps.SessionTable).values(deps.SessionToRow(prepared.info)).onConflictDoNothing().run() + + for (const row of prepared.messages) { + db.insert(deps.MessageTable).values(row).onConflictDoNothing().run() + } + for (const row of prepared.parts) { + db.insert(deps.PartTable).values(row).onConflictDoNothing().run() } - Database.effect(() => Bus.publish(SessionCreatedEvent, { info })) + deps.Database.effect(() => deps.Bus.publish(deps.SessionCreatedEvent, { info: prepared.info })) }) - return info + return prepared.info } diff --git a/packages/kilo-gateway/src/index.ts b/packages/kilo-gateway/src/index.ts index d0b80b027c9..cee10c2f781 100644 --- a/packages/kilo-gateway/src/index.ts +++ b/packages/kilo-gateway/src/index.ts @@ -67,7 +67,13 @@ export { type OrganizationModeConfig, } from "./api/modes.js" export { fetchKilocodeNotifications, type KilocodeNotification } from "./api/notifications.js" -export { fetchCloudSession, fetchCloudSessionForImport, importSessionToDb } from "./cloud-sessions.js" +export { + fetchCloudSession, + fetchCloudSessionForImport, + SessionImportValidationError, + prepareSessionImport, + importSessionToDb, +} from "./cloud-sessions.js" // ============================================================================ // Server Routes (optional - requires hono and OpenCode dependencies) diff --git a/packages/kilo-gateway/src/server/routes.ts b/packages/kilo-gateway/src/server/routes.ts index 6774a7c9459..487987fbbb2 100644 --- a/packages/kilo-gateway/src/server/routes.ts +++ b/packages/kilo-gateway/src/server/routes.ts @@ -11,7 +11,12 @@ import { fetchOrganizationModes, clearModesCache } from "../api/modes.js" import { KILO_API_BASE, HEADER_FEATURE, HEADER_ORGANIZATIONID } from "../api/constants.js" import { buildKiloHeaders } from "../headers.js" import type { ImportDeps, DrizzleDb } from "../cloud-sessions.js" -import { fetchCloudSession, fetchCloudSessionForImport, importSessionToDb } from "../cloud-sessions.js" +import { + fetchCloudSession, + fetchCloudSessionForImport, + importSessionToDb, + SessionImportValidationError, +} from "../cloud-sessions.js" import { createEditHandler } from "./edit.js" import { createFimHandler } from "./fim.js" import { @@ -646,6 +651,7 @@ export function createKiloRoutes(deps: KiloRoutesDeps) { return c.json(info) } catch (err: any) { + if (err instanceof SessionImportValidationError) return c.json({ error: "Invalid export data" }, 400) console.error("[Kilo Gateway] cloud/session/import: unhandled error", err?.message ?? err) return c.json({ error: "Internal error" }, 500) } diff --git a/packages/kilo-gateway/test/cloud-sessions.test.ts b/packages/kilo-gateway/test/cloud-sessions.test.ts index 2c292bb5ef6..232af2b4641 100644 --- a/packages/kilo-gateway/test/cloud-sessions.test.ts +++ b/packages/kilo-gateway/test/cloud-sessions.test.ts @@ -1,5 +1,10 @@ import { describe, expect, test } from "bun:test" -import { fetchCloudSession, fetchCloudSessionForImport } from "../src/cloud-sessions" +import { + fetchCloudSession, + fetchCloudSessionForImport, + prepareSessionImport, + SessionImportValidationError, +} from "../src/cloud-sessions" async function expectStalledFetchToTimeOut(run: () => Promise) { const fetch = globalThis.fetch @@ -46,3 +51,232 @@ describe("cloud session export requests", () => { await expectStalledFetchToTimeOut(() => fetchCloudSessionForImport("token", "session-id")) }) }) + +describe("cloud session import preparation", () => { + function sample() { + return { + info: { + id: "ses_cloud", + slug: "cloud-session", + projectID: "proj_cloud", + workspaceID: "wrk_cloud", + directory: "/cloud/workspace", + path: "cloud/path", + parentID: "ses_cloud_parent", + share: { url: "https://example.com/share" }, + revert: { messageID: "msg_cloud_child", partID: "prt_cloud_compaction" }, + metadata: { source: "cloud" }, + title: "Cloud session", + version: "7.4.11", + time: { created: 10, updated: 20, compacting: 30, archived: 40 }, + }, + messages: [ + { + info: { + id: "msg_cloud_child", + sessionID: "ses_cloud", + parentID: "msg_cloud_parent", + role: "assistant", + time: { created: 12 }, + modelID: "test", + providerID: "test", + mode: "build", + agent: "build", + path: { cwd: "/cloud/workspace", root: "/cloud/workspace" }, + cost: 1, + tokens: { input: 1, output: 1, reasoning: 0, cache: { read: 0, write: 0 } }, + }, + parts: [ + { + id: "prt_cloud_tool", + messageID: "msg_cloud_child", + sessionID: "ses_cloud", + type: "tool", + callID: "call_cloud", + tool: "read", + state: { + status: "completed", + input: {}, + output: "attached", + title: "Attachment", + metadata: {}, + time: { start: 12, end: 13 }, + attachments: [ + { + id: "prt_cloud_attachment", + messageID: "msg_cloud_child", + sessionID: "ses_cloud", + type: "file", + mime: "text/plain", + filename: "result.txt", + url: "data:text/plain,attached", + }, + ], + }, + }, + ], + }, + { + info: { + id: "msg_cloud_parent", + sessionID: "ses_cloud", + role: "user", + time: { created: 11 }, + agent: "build", + model: { providerID: "test", modelID: "test" }, + }, + parts: [ + { + id: "prt_cloud_compaction", + messageID: "msg_cloud_parent", + sessionID: "ses_cloud", + type: "compaction", + auto: true, + tail_start_id: "msg_cloud_child", + }, + ], + }, + ], + } + } + + function deps( + ids = ["msg_local_child", "prt_local_tool", "prt_local_attachment", "msg_local_parent", "prt_local_compaction"], + target: { workspaceID?: string; path?: string } = { workspaceID: "wrk_local", path: "nested" }, + ) { + const available = [...ids] + return { + Instance: { directory: "/local/workspace", project: { id: "proj_local" } }, + ...target, + Identifier: { + descending: () => "ses_local", + ascending: () => available.shift() ?? "unexpected_id", + }, + } + } + + test("remaps references and target context without mutating the export", () => { + const data = sample() + const before = structuredClone(data) + const result = prepareSessionImport(data, deps()) + + expect(result.info).toMatchObject({ + id: "ses_local", + projectID: "proj_local", + directory: "/local/workspace", + workspaceID: "wrk_local", + path: "nested", + slug: "cloud-session", + title: "Cloud session", + version: "7.4.11", + metadata: { source: "cloud" }, + time: { created: 10, compacting: 30, archived: 40 }, + }) + expect(result.info).not.toHaveProperty("parentID") + expect(result.info).not.toHaveProperty("share") + expect(result.info).not.toHaveProperty("revert") + expect(result.info.time.updated).toBeGreaterThanOrEqual(before.info.time.updated) + expect(result.messages[0]).toMatchObject({ + id: "msg_local_child", + session_id: "ses_local", + data: { id: "msg_local_child", sessionID: "ses_local", parentID: "msg_local_parent" }, + }) + expect(result.parts[0]).toMatchObject({ + id: "prt_local_tool", + message_id: "msg_local_child", + session_id: "ses_local", + data: { + id: "prt_local_tool", + messageID: "msg_local_child", + sessionID: "ses_local", + state: { + attachments: [ + { + id: "prt_local_attachment", + messageID: "msg_local_child", + sessionID: "ses_local", + type: "file", + mime: "text/plain", + filename: "result.txt", + url: "data:text/plain,attached", + }, + ], + }, + }, + }) + expect(result.parts[1]).toMatchObject({ + id: "prt_local_compaction", + message_id: "msg_local_parent", + session_id: "ses_local", + data: { + id: "prt_local_compaction", + messageID: "msg_local_parent", + sessionID: "ses_local", + tail_start_id: "msg_local_child", + }, + }) + expect(data).toEqual(before) + + const detached = prepareSessionImport(sample(), deps(undefined, {})) + expect(detached.info).not.toHaveProperty("workspaceID") + expect(detached.info).not.toHaveProperty("path") + }) + + test("rejects malformed, duplicate, and dangling transcript data", () => { + const base = sample() + const first = base.messages[0]! + const second = base.messages[1]! + const tool = first.parts[0]! + const state = tool.state + const attachment = state.attachments[0]! + const part = second.parts[0]! + const invalid = [ + { info: base.info }, + { ...base, messages: [...base.messages, null] }, + { ...base, messages: [first, { ...second, parts: [...second.parts, null] }] }, + { ...base, messages: [first, { ...second, info: { ...second.info, id: first.info.id } }] }, + { ...base, messages: [first, { ...second, parts: [part, { ...part }] }] }, + { ...base, messages: [{ ...first, info: { ...first.info, parentID: "msg_missing" } }, second] }, + { ...base, messages: [first, { ...second, parts: [{ ...part, tail_start_id: "msg_missing" }] }] }, + { ...base, messages: [first, { ...second, parts: [{ ...part, messageID: "msg_missing" }] }] }, + { + ...base, + messages: [{ ...first, parts: [{ ...tool, state: { ...state, attachments: [null] } }] }, second], + }, + { + ...base, + messages: [ + { + ...first, + parts: [{ ...tool, state: { ...state, attachments: [{ ...attachment, messageID: "msg_missing" }] } }], + }, + second, + ], + }, + { + ...base, + messages: [ + { + ...first, + parts: [{ ...tool, state: { ...state, attachments: [{ ...attachment, id: tool.id }] } }], + }, + second, + ], + }, + { + ...base, + messages: [ + { + ...first, + parts: [{ ...tool, state: { ...state, attachments: [attachment, { ...attachment }] } }], + }, + second, + ], + }, + ] + + for (const item of invalid) { + expect(() => prepareSessionImport(item, deps())).toThrow(SessionImportValidationError) + } + }) +}) diff --git a/packages/opencode/src/kilocode/server/httpapi/groups/kilo-gateway.ts b/packages/opencode/src/kilocode/server/httpapi/groups/kilo-gateway.ts index 80cb080caf8..9e0068cbaba 100644 --- a/packages/opencode/src/kilocode/server/httpapi/groups/kilo-gateway.ts +++ b/packages/opencode/src/kilocode/server/httpapi/groups/kilo-gateway.ts @@ -122,6 +122,11 @@ export const CloudSessionImportBody = Schema.Struct({ sessionId: Schema.String, }) +export class CloudSessionImportError extends Schema.ErrorClass("CloudSessionImportError")( + { error: Schema.String }, + { httpApiStatus: 500 }, +) {} + const GroupEntry = Schema.Union([ Schema.String, Schema.Tuple([ @@ -441,7 +446,7 @@ export const KiloGatewayApi = HttpApi.make("kilo") query: WorkspaceRoutingQuery, payload: CloudSessionImportBody, success: described(CloudSessionData.fields.info, "Imported session info"), - error: [HttpApiError.BadRequest, HttpApiError.Unauthorized, HttpApiError.NotFound], + error: [HttpApiError.BadRequest, HttpApiError.Unauthorized, HttpApiError.NotFound, CloudSessionImportError], }).annotateMerge( OpenApi.annotations({ identifier: "kilo.cloud.session.import", diff --git a/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts b/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts index b46a5bbe118..dac2ebca8ee 100644 --- a/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts +++ b/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts @@ -1,13 +1,15 @@ +import path from "node:path" import { GatewayError, + SessionImportValidationError, fetchCloudSession, fetchCloudSessionForImport, fetchKiloImageModels, getCloudSessions, getOrganizationId, getToken, - importSessionToDb, normalizeClawStatus, + prepareSessionImport, } from "@kilocode/kilo-gateway" import { HEADER_FEATURE, @@ -26,26 +28,28 @@ import { DIRECT_FIM_ENV, requestMistralFim, resolveFimTarget } from "@kilocode/k import { DIRECT_EDIT_ENV, extractFencedBody, resolveEditTarget } from "@kilocode/kilo-gateway/edit" import { buildMercuryEditPrompt } from "@kilocode/kilo-gateway/edit-prompt" import { buildKiloHeaders } from "@kilocode/kilo-gateway" -import { Effect, Schema } from "effect" +import { Cause, Effect, Result, Schema } from "effect" import * as Stream from "effect/Stream" import { HttpServerRequest, HttpServerResponse } from "effect/unstable/http" import { HttpApiBuilder, HttpApiError } from "effect/unstable/httpapi" import * as Log from "@opencode-ai/core/util/log" import { Flag } from "@opencode-ai/core/flag/flag" +import { Database } from "@opencode-ai/core/database/database" +import type { DeepMutable } from "@opencode-ai/core/schema" +import { SessionV1 } from "@opencode-ai/core/v1/session" import { KilocodeConfig } from "@/kilocode/config/config" import { Auth } from "@/auth" -import { EffectBridge } from "@/effect/bridge" +import { WorkspaceRef } from "@/effect/instance-ref" import { EventV2Bridge } from "@/event-v2-bridge" import { Identifier } from "@/id/id" import { Instance } from "@/kilocode/instance" import { InstanceStore } from "@/project/instance-store" import { ModelCache } from "@/provider/model-cache" import { InstanceHttpApi } from "@/server/routes/instance/httpapi/api" -import { MessageTable, PartTable, SessionTable } from "@opencode-ai/core/session/sql" +import { MessageTable, PartTable } from "@opencode-ai/core/session/sql" import { Session } from "@/session/session" -import { Database } from "@/storage/db" import { Storage } from "@/storage/storage" -import { AudioTranscriptionsBody, ClawStatus, EditBody, FimBody } from "../groups/kilo-gateway" +import { AudioTranscriptionsBody, ClawStatus, CloudSessionImportError, EditBody, FimBody } from "../groups/kilo-gateway" import { baseKey } from "../../../session-portability/cumulative-diff" import { extractSessionDiffs, restoreSessionDiffs } from "../../../session-portability/session-diff-restore" @@ -66,6 +70,8 @@ export const kiloGatewayHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilo", const store = yield* InstanceStore.Service const cache = yield* ModelCache.Service const events = yield* EventV2Bridge.Service + const database = yield* Database.Service + const storage = yield* Storage.Service const profile = Effect.fn("KiloGatewayHttpApi.profile")(function* () { const info = yield* auth.get("kilo").pipe(Effect.mapError(() => new HttpApiError.BadRequest({}))) @@ -466,69 +472,123 @@ export const kiloGatewayHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilo", }), ), ) - if (!fetched) return jsonError("Internal error", 500) + if (!fetched) return yield* Effect.fail(new CloudSessionImportError({ error: "Internal error" })) if (!fetched.ok) return jsonError(fetched.error, fetched.status) if (!fetched.data?.info?.id) return yield* Effect.fail(new HttpApiError.BadRequest({})) const diffs = extractSessionDiffs(fetched.data) - const bridge = yield* EffectBridge.make() - return yield* Effect.tryPromise({ + const workspaceID = yield* WorkspaceRef + const subdir = path.relative(path.resolve(Instance.worktree), Instance.directory).replaceAll("\\", "/") + const prepared = yield* Effect.try({ + try: () => prepareSessionImport(fetched.data, { Instance, Identifier, workspaceID, path: subdir }), + catch: (err) => { + if (err instanceof SessionImportValidationError) return new HttpApiError.BadRequest({}) + const name = + err instanceof Error + ? err.name + : typeof err === "object" && err !== null && "_tag" in err && typeof err._tag === "string" + ? err._tag + : "UnknownError" + log.error("cloud session import failed", { + route: "cloud/session/import", + stage: "prepare", + error: name, + }) + return new CloudSessionImportError({ error: "Internal error" }) + }, + }) + const session = yield* Effect.try({ + try: () => Schema.decodeUnknownSync(Session.Info)(prepared.info), + catch: () => new HttpApiError.BadRequest({}), + }) + const messages = yield* Effect.try({ try: () => - bridge.promise( - Effect.gen(function* () { - if (diffs.length > 0) { - yield* Effect.try({ - try: () => restoreSessionDiffs({ directory: Instance.directory, diffs }), - catch: (err) => err, - }).pipe( - Effect.catch((err) => - Effect.sync(() => { - logError("cloud/session/import/restore", err) - return undefined - }), - ), - ) - } - - const imported = yield* Effect.sync(() => - importSessionToDb(fetched.data, { - Database, - Instance, - SessionTable, - MessageTable, - PartTable, - SessionToRow: Session.toRow, - Bus: { - publish: (_event, payload) => { - const info = (payload as { info: Session.Info }).info - return bridge.promise(events.publish(Session.Event.Created, { sessionID: info.id, info })) - }, - }, - SessionCreatedEvent: { type: Session.Event.Created.type, properties: Session.Event.Created.data }, - Identifier, - }), - ) - - if (diffs.length > 0) { - yield* Storage.Service.use((storage) => - Effect.all([ - storage.write(baseKey(imported.id), diffs), - storage.write(["session_diff", imported.id], diffs), - ]), - ).pipe( - Effect.catch((err) => - Effect.sync(() => { - logError("cloud/session/import/diff", err) - }), - ), - ) - } - - return imported - }), - ), + prepared.messages.map((row) => { + const info = Schema.decodeUnknownSync(SessionV1.Info)(row.data) + const { id, sessionID, ...data } = info + // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion -- decoding validates the shape; the database type only removes readonly modifiers + return { id, session_id: sessionID, time_created: row.time_created, data: data as DeepMutable } + }), + catch: () => new HttpApiError.BadRequest({}), + }) + const parts = yield* Effect.try({ + try: () => + prepared.parts.map((row) => { + const part = Schema.decodeUnknownSync(SessionV1.Part)(row.data) + const { id, messageID, sessionID, ...data } = part + return { + id, + message_id: messageID, + session_id: sessionID, + // oxlint-disable-next-line typescript-eslint/no-unsafe-type-assertion -- decoding validates the shape; the database type only removes readonly modifiers + data: data as DeepMutable, + } + }), catch: () => new HttpApiError.BadRequest({}), }) + const imported = yield* Effect.gen(function* () { + yield* events.publish( + Session.Event.Created, + { sessionID: session.id, info: session }, + { + commit: () => + Effect.gen(function* () { + for (const row of messages) { + yield* database.db.insert(MessageTable).values([row]).run().pipe(Effect.orDie) + } + for (const row of parts) { + yield* database.db.insert(PartTable).values([row]).run().pipe(Effect.orDie) + } + }), + }, + ) + return session + }).pipe( + Effect.catchCause((cause) => + Effect.sync(() => { + const err = Result.getOrUndefined(Cause.findDefect(cause)) ?? Result.getOrUndefined(Cause.findError(cause)) + const name = + err instanceof Error + ? err.name + : typeof err === "object" && err !== null && "_tag" in err && typeof err._tag === "string" + ? err._tag + : "UnknownError" + log.error("cloud session import failed", { + route: "cloud/session/import", + stage: "write", + error: name, + sessionID: session.id, + messages: messages.length, + parts: parts.length, + }) + }).pipe(Effect.andThen(Effect.fail(new CloudSessionImportError({ error: "Internal error" })))), + ), + ) + + if (diffs.length > 0) { + yield* Effect.try({ + try: () => restoreSessionDiffs({ directory: Instance.directory, diffs }), + catch: (err) => err, + }).pipe( + Effect.catch((err) => + Effect.sync(() => { + logError("cloud/session/import/restore", err) + }), + ), + ) + yield* Effect.all([ + storage.write(baseKey(imported.id), diffs), + storage.write(["session_diff", imported.id], diffs), + ]).pipe( + Effect.catch((err) => + Effect.sync(() => { + logError("cloud/session/import/diff", err) + }), + ), + ) + } + + return imported }) const imageModels = Effect.fn("KiloGatewayHttpApi.imageModels")(function* () { diff --git a/packages/opencode/test/kilocode/server/cloud-session-import.test.ts b/packages/opencode/test/kilocode/server/cloud-session-import.test.ts new file mode 100644 index 00000000000..85de22ee6c5 --- /dev/null +++ b/packages/opencode/test/kilocode/server/cloud-session-import.test.ts @@ -0,0 +1,273 @@ +import { afterAll, afterEach, describe, expect, test } from "bun:test" +import { existsSync } from "node:fs" +import { mkdir } from "node:fs/promises" +import path from "node:path" +import { Database } from "@opencode-ai/core/database/database" +import { EventSequenceTable, EventTable } from "@opencode-ai/core/event/sql" +import { Flag } from "@opencode-ai/core/flag/flag" +import { ProjectV2 } from "@opencode-ai/core/project" +import { MessageTable, PartTable, SessionTable } from "@opencode-ai/core/session/sql" +import { Effect, Schema } from "effect" +import { eq } from "drizzle-orm" +import { SessionID } from "../../../src/session/schema" + +const bodies = new Map() +const ingest = Bun.serve({ + hostname: "127.0.0.1", + port: 0, + fetch(request) { + const url = new URL(request.url) + if (request.method !== "GET" || url.pathname !== "/api/session/ses_cloud/export" || url.search) { + return Response.json({ error: "Unexpected test request" }, { status: 503 }) + } + const auth = request.headers.get("authorization") + const token = auth?.startsWith("Bearer ") ? auth.slice(7) : undefined + if (!token || !bodies.has(token)) { + return Response.json({ error: "Unexpected test token" }, { status: 503 }) + } + const body = bodies.get(token) + bodies.delete(token) + return Response.json(body) + }, +}) +afterAll(async () => { + bodies.clear() + await ingest.stop(true) +}) + +const [runtime, gateway, server, fixture] = await (async () => { + const disabled = process.env.KILO_DISABLE_SESSION_INGEST + const base = process.env.KILO_SESSION_INGEST_URL + process.env.KILO_DISABLE_SESSION_INGEST = "1" + process.env.KILO_SESSION_INGEST_URL = ingest.url.origin + try { + return await Promise.all([ + import("../../../src/effect/app-runtime"), + import("../../../src/kilocode/server/httpapi/groups/kilo-gateway"), + import("../../../src/server/routes/instance/httpapi/server"), + import("../../fixture/fixture"), + ]) + } finally { + if (disabled === undefined) delete process.env.KILO_DISABLE_SESSION_INGEST + else process.env.KILO_DISABLE_SESSION_INGEST = disabled + if (base === undefined) delete process.env.KILO_SESSION_INGEST_URL + else process.env.KILO_SESSION_INGEST_URL = base + } +})() +const { AppRuntime } = runtime +const { KiloGatewayPaths } = gateway +const HttpApiApp = server +const { disposeAllInstances, tmpdir } = fixture + +const created: string[] = [] +const Imported = Schema.Struct({ + id: Schema.String, + projectID: Schema.String, + workspaceID: Schema.String, + directory: Schema.String, + path: Schema.String, +}) + +function data(diff = false) { + return { + info: { + id: "ses_cloud", + slug: "cloud-session", + projectID: "proj_cloud", + directory: "/cloud/workspace", + title: "Cloud session", + version: "7.4.11", + time: { created: 10, updated: 20 }, + }, + messages: [ + { + info: { + id: "msg_cloud", + sessionID: "ses_cloud", + role: "user", + agent: "build", + model: { providerID: "test", modelID: "test" }, + time: { created: 11 }, + }, + parts: [ + { + id: "prt_cloud", + messageID: "msg_cloud", + sessionID: "ses_cloud", + type: "text", + text: "hello", + }, + ], + }, + ], + ...(diff ? { sessionDiff: [{ file: "restored.txt", after: "restored\n", status: "added" }] } : {}), + } +} + +async function request(directory: string, body: unknown) { + const auth = process.env.KILO_AUTH_CONTENT + const token = `test-${crypto.randomUUID()}` + bodies.set(token, body) + process.env.KILO_AUTH_CONTENT = JSON.stringify({ kilo: { type: "api", key: token } }) + try { + return await HttpApiApp.webHandler().handler( + new Request(`http://localhost${KiloGatewayPaths.cloudSessionImport}`, { + method: "POST", + headers: { "content-type": "application/json", "x-kilo-directory": directory }, + body: JSON.stringify({ sessionId: "ses_cloud" }), + }), + HttpApiApp.context, + ) + } finally { + bodies.delete(token) + if (auth === undefined) delete process.env.KILO_AUTH_CONTENT + else process.env.KILO_AUTH_CONTENT = auth + } +} + +async function routed(id: string, run: () => Promise) { + const workspace = Flag.KILO_WORKSPACE_ID + Flag.KILO_WORKSPACE_ID = id + try { + return await run() + } finally { + Flag.KILO_WORKSPACE_ID = workspace + } +} + +function counts() { + return AppRuntime.runPromise( + Database.Service.use(({ db }) => + Effect.all([ + db.select().from(SessionTable).all(), + db.select().from(MessageTable).all(), + db.select().from(PartTable).all(), + db.select().from(EventTable).all(), + db.select().from(EventSequenceTable).all(), + ]).pipe(Effect.map((rows) => rows.map((items) => items.length))), + ), + ) +} + +afterEach(async () => { + await AppRuntime.runPromise( + Database.Service.use(({ db }) => + Effect.gen(function* () { + yield* db.run("DROP TRIGGER IF EXISTS fail_cloud_import") + for (const id of created.splice(0)) { + yield* db.delete(EventTable).where(eq(EventTable.aggregate_id, id)).run() + yield* db.delete(EventSequenceTable).where(eq(EventSequenceTable.aggregate_id, id)).run() + yield* db + .delete(SessionTable) + .where(eq(SessionTable.id, SessionID.make(id))) + .run() + } + }), + ), + ) + await disposeAllInstances() +}) + +describe("cloud session import", () => { + test("rejects an invalid export before persistence", async () => { + await using dir = await tmpdir({ git: true }) + const response = await request(dir.path, { + info: { id: "ses_invalid", title: "Invalid", time: { created: 10, updated: 20 } }, + messages: [], + }) + + expect(response.status).toBe(400) + }) + + test("rejects a mixed malformed transcript before persistence or file restoration", async () => { + await using dir = await tmpdir({ git: true }) + const restored = path.join(dir.path, "restored.txt") + const before = await counts() + const source = data(true) + const msg = source.messages[0]! + const response = await request(dir.path, { + ...source, + messages: [{ ...msg, parts: [...msg.parts, { type: "text", text: "discarded" }] }], + }) + + expect(response.status).toBe(400) + expect(existsSync(restored)).toBe(false) + expect(await counts()).toEqual(before) + }) + + test("commits the imported transcript and creation event atomically", async () => { + await using dir = await tmpdir({ git: true }) + const nested = path.join(dir.path, "nested", "target") + await mkdir(nested, { recursive: true }) + const response = await routed("wrk_local", () => request(nested, data())) + expect(response.status).toBe(200) + const imported = Schema.decodeUnknownSync(Imported)(await response.json()) + created.push(imported.id) + expect(imported.id).not.toBe("ses_cloud") + expect(imported.workspaceID).toBe("wrk_local") + expect(imported.directory).toBe(nested) + expect(imported.path).toBe("nested/target") + + const [session, messages, parts, events, sequence] = await AppRuntime.runPromise( + Database.Service.use(({ db }) => + Effect.all([ + db + .select() + .from(SessionTable) + .where(eq(SessionTable.id, SessionID.make(imported.id))) + .get(), + db + .select() + .from(MessageTable) + .where(eq(MessageTable.session_id, SessionID.make(imported.id))) + .all(), + db + .select() + .from(PartTable) + .where(eq(PartTable.session_id, SessionID.make(imported.id))) + .all(), + db.select().from(EventTable).where(eq(EventTable.aggregate_id, imported.id)).all(), + db.select().from(EventSequenceTable).where(eq(EventSequenceTable.aggregate_id, imported.id)).get(), + ]), + ), + ) + + expect(session).toMatchObject({ + project_id: ProjectV2.ID.make(imported.projectID), + workspace_id: "wrk_local", + directory: nested, + path: "nested/target", + }) + expect(messages).toHaveLength(1) + expect(messages[0]).toMatchObject({ session_id: imported.id, data: { role: "user" } }) + expect(parts).toHaveLength(1) + expect(parts[0]).toMatchObject({ + message_id: messages[0]?.id, + session_id: imported.id, + data: { type: "text" }, + }) + expect(events.map((event) => ({ seq: event.seq, type: event.type }))).toEqual([ + { seq: 0, type: "session.created.1" }, + ]) + expect(sequence?.seq).toBe(0) + }) + + test("rolls back persistence before restoring files", async () => { + await using dir = await tmpdir({ git: true }) + const restored = path.join(dir.path, "restored.txt") + const before = await counts() + await AppRuntime.runPromise( + Database.Service.use(({ db }) => + db.run( + 'CREATE TRIGGER fail_cloud_import BEFORE INSERT ON message BEGIN SELECT RAISE(ABORT, "failed import"); END', + ), + ), + ) + + const response = await request(dir.path, data(true)) + expect(response.status).toBe(500) + expect(await response.json()).toEqual({ error: "Internal error" }) + expect(existsSync(restored)).toBe(false) + expect(await counts()).toEqual(before) + }) +}) diff --git a/packages/opencode/test/kilocode/server/kilo-gateway-statuses.test.ts b/packages/opencode/test/kilocode/server/kilo-gateway-statuses.test.ts index f8d1b674823..115d4415589 100644 --- a/packages/opencode/test/kilocode/server/kilo-gateway-statuses.test.ts +++ b/packages/opencode/test/kilocode/server/kilo-gateway-statuses.test.ts @@ -1,4 +1,5 @@ import { NodeHttpServer } from "@effect/platform-node" +import { Database } from "@opencode-ai/core/database/database" import { describe, expect } from "bun:test" import { Effect, Layer } from "effect" import { HttpClient, HttpClientRequest, HttpRouter } from "effect/unstable/http" @@ -9,6 +10,7 @@ import { kiloGatewayHandlers } from "../../../src/kilocode/server/httpapi/handle import { InstanceStore } from "../../../src/project/instance-store" import { ModelCache } from "../../../src/provider/model-cache" import { Session } from "../../../src/session/session" +import { Storage } from "../../../src/storage/storage" import { Authorization } from "../../../src/server/routes/instance/httpapi/middleware/authorization" import { InstanceContextMiddleware } from "../../../src/server/routes/instance/httpapi/middleware/instance-context" import { schemaErrorLayer } from "../../../src/server/routes/instance/httpapi/middleware/schema-error" @@ -26,6 +28,7 @@ const auth = Layer.mock(Auth.Service)({ const store = Layer.mock(InstanceStore.Service)({}) const cache = Layer.mock(ModelCache.Service)({}) const session = Layer.mock(Session.Service)({}) +const storage = Layer.mock(Storage.Service)({}) const passthroughAuthorization = Layer.succeed( Authorization, Authorization.of((effect) => effect), @@ -54,6 +57,8 @@ const layer = HttpRouter.serve( session, EventV2Bridge.defaultLayer, ]), + Layer.provide(Database.defaultLayer), + Layer.provide(storage), ), { disableListenLog: true, disableLogger: true }, ).pipe(Layer.provideMerge(NodeHttpServer.layerTest)) @@ -65,7 +70,10 @@ function stub(run: () => Response | Promise) { const fetch: typeof globalThis.fetch = Object.assign( async (input: RequestInfo | URL, init?: RequestInit) => { const url = typeof input === "string" ? input : input instanceof URL ? input.href : input.url - if (url.startsWith("http://127.0.0.1:")) return original(input, init) + const headers = new Headers(init?.headers ?? (input instanceof Request ? input.headers : undefined)) + if (url.startsWith("http://127.0.0.1:") && headers.get("authorization") !== "Bearer test-token") { + return original(input, init) + } return run() }, { preconnect: original.preconnect }, diff --git a/packages/sdk/js/src/v2/gen/types.gen.ts b/packages/sdk/js/src/v2/gen/types.gen.ts index 9bdf89a86c3..034200a895e 100644 --- a/packages/sdk/js/src/v2/gen/types.gen.ts +++ b/packages/sdk/js/src/v2/gen/types.gen.ts @@ -3129,6 +3129,10 @@ export type EffectHttpApiErrorServiceUnavailable = { _tag: "ServiceUnavailable" } +export type CloudSessionImportError = { + error: string +} + export type AgentRequirementResult = { agent: string directory: string @@ -12373,6 +12377,10 @@ export type KiloCloudSessionImportErrors = { * Not found */ 404: NotFoundError + /** + * CloudSessionImportError + */ + 500: CloudSessionImportError } export type KiloCloudSessionImportError = KiloCloudSessionImportErrors[keyof KiloCloudSessionImportErrors] diff --git a/packages/sdk/openapi.json b/packages/sdk/openapi.json index 1c1f5943741..c8cb8ef11a9 100644 --- a/packages/sdk/openapi.json +++ b/packages/sdk/openapi.json @@ -14735,6 +14735,16 @@ } } } + }, + "500": { + "description": "CloudSessionImportError", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CloudSessionImportError" + } + } + } } }, "description": "Download a cloud-synced session and write it to local storage with fresh IDs.", @@ -31423,6 +31433,16 @@ "required": ["_tag"], "additionalProperties": false }, + "CloudSessionImportError": { + "type": "object", + "properties": { + "error": { + "type": "string" + } + }, + "required": ["error"], + "additionalProperties": false + }, "AgentRequirementResult": { "type": "object", "properties": { diff --git a/script/check-opencode-promise-facades.ts b/script/check-opencode-promise-facades.ts index 75e1b1f0a89..65b5e7236a4 100644 --- a/script/check-opencode-promise-facades.ts +++ b/script/check-opencode-promise-facades.ts @@ -44,6 +44,7 @@ const testAllow: Record = { "kilocode/session/platform-attribution.test.ts": { count: 2, reason: "existing runtime integration test" }, "kilocode/session-prompt-queue.test.ts": { count: 6, reason: "prompt queue legacy instance bridge regression" }, "server/experimental-session-list.test.ts": { count: 2, reason: "Kilo session list integration test" }, + "kilocode/server/cloud-session-import.test.ts": { count: 5, reason: "full app cloud import transaction integration" }, "kilocode/server/listener-runtime.test.ts": { count: 4, reason: "listener and AppRuntime integration test" }, "tool/recall.test.ts": { count: 11, reason: "existing runtime integration test" }, } From b63c6693e5a53fb7ea90878624805c4a7fda1be6 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Fri, 17 Jul 2026 22:23:10 +0200 Subject: [PATCH 2/3] chore: address cloud import review feedback --- .changeset/cloud-session-import.md | 2 +- packages/kilo-gateway/src/cloud-sessions.ts | 24 +++++++++++++++++-- .../kilo-gateway/test/cloud-sessions.test.ts | 9 ++++++- 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/.changeset/cloud-session-import.md b/.changeset/cloud-session-import.md index cc8935ae0c2..4d32b079236 100644 --- a/.changeset/cloud-session-import.md +++ b/.changeset/cloud-session-import.md @@ -3,4 +3,4 @@ "@kilocode/kilo-gateway": patch --- -Allow forking Cloud Agent sessions into the local CLI. +Fix Cloud Agent session imports in installed CLI builds and prevent malformed exports or write failures from leaving partial imports. diff --git a/packages/kilo-gateway/src/cloud-sessions.ts b/packages/kilo-gateway/src/cloud-sessions.ts index 444634254cf..26251db3a60 100644 --- a/packages/kilo-gateway/src/cloud-sessions.ts +++ b/packages/kilo-gateway/src/cloud-sessions.ts @@ -196,6 +196,17 @@ export function prepareSessionImport(data: Export, deps: PrepareDeps) { items.push({ id, created: time.created, info: msg.info, parts, ...(parent !== undefined ? { parent } : {}) }) } + const parents = new Map(items.map((item) => [item.id, item.parent])) + for (const item of items) { + const seen = new Set([item.id]) + let parent = item.parent + while (parent !== undefined) { + if (seen.has(parent)) throw new SessionImportValidationError("Circular message parent") + seen.add(parent) + parent = parents.get(parent) + } + } + const now = Date.now() const time = { created: data.info.time?.created ?? now, @@ -225,6 +236,7 @@ export function prepareSessionImport(data: Export, deps: PrepareDeps) { delete info.parentID delete info.share delete info.revert + delete info.permission const messages: Array<{ id: string @@ -294,10 +306,18 @@ export function importSessionToDb(data: Export, deps: ImportDeps) { db.insert(deps.SessionTable).values(deps.SessionToRow(prepared.info)).onConflictDoNothing().run() for (const row of prepared.messages) { - db.insert(deps.MessageTable).values(row).onConflictDoNothing().run() + const { id: _, sessionID: __, ...data } = row.data + db.insert(deps.MessageTable) + .values({ id: row.id, session_id: row.session_id, time_created: row.time_created, data }) + .onConflictDoNothing() + .run() } for (const row of prepared.parts) { - db.insert(deps.PartTable).values(row).onConflictDoNothing().run() + const { id: _, messageID: __, sessionID: ___, ...data } = row.data + db.insert(deps.PartTable) + .values({ id: row.id, message_id: row.message_id, session_id: row.session_id, data }) + .onConflictDoNothing() + .run() } deps.Database.effect(() => deps.Bus.publish(deps.SessionCreatedEvent, { info: prepared.info })) diff --git a/packages/kilo-gateway/test/cloud-sessions.test.ts b/packages/kilo-gateway/test/cloud-sessions.test.ts index 232af2b4641..4b69e4bd0f7 100644 --- a/packages/kilo-gateway/test/cloud-sessions.test.ts +++ b/packages/kilo-gateway/test/cloud-sessions.test.ts @@ -65,6 +65,7 @@ describe("cloud session import preparation", () => { parentID: "ses_cloud_parent", share: { url: "https://example.com/share" }, revert: { messageID: "msg_cloud_child", partID: "prt_cloud_compaction" }, + permission: [{ permission: "*", pattern: "*", action: "allow" }], metadata: { source: "cloud" }, title: "Cloud session", version: "7.4.11", @@ -158,7 +159,9 @@ describe("cloud session import preparation", () => { test("remaps references and target context without mutating the export", () => { const data = sample() const before = structuredClone(data) + const start = Date.now() const result = prepareSessionImport(data, deps()) + const end = Date.now() expect(result.info).toMatchObject({ id: "ses_local", @@ -175,7 +178,9 @@ describe("cloud session import preparation", () => { expect(result.info).not.toHaveProperty("parentID") expect(result.info).not.toHaveProperty("share") expect(result.info).not.toHaveProperty("revert") - expect(result.info.time.updated).toBeGreaterThanOrEqual(before.info.time.updated) + expect(result.info).not.toHaveProperty("permission") + expect(result.info.time.updated).toBeGreaterThanOrEqual(start) + expect(result.info.time.updated).toBeLessThanOrEqual(end) expect(result.messages[0]).toMatchObject({ id: "msg_local_child", session_id: "ses_local", @@ -237,6 +242,8 @@ describe("cloud session import preparation", () => { { ...base, messages: [first, { ...second, info: { ...second.info, id: first.info.id } }] }, { ...base, messages: [first, { ...second, parts: [part, { ...part }] }] }, { ...base, messages: [{ ...first, info: { ...first.info, parentID: "msg_missing" } }, second] }, + { ...base, messages: [{ ...first, info: { ...first.info, parentID: first.info.id } }, second] }, + { ...base, messages: [first, { ...second, info: { ...second.info, parentID: first.info.id } }] }, { ...base, messages: [first, { ...second, parts: [{ ...part, tail_start_id: "msg_missing" }] }] }, { ...base, messages: [first, { ...second, parts: [{ ...part, messageID: "msg_missing" }] }] }, { From 5255b4446d6d26953e11199be09b7f2f3351829a Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Mon, 20 Jul 2026 09:49:13 +0200 Subject: [PATCH 3/3] refactor(gateway): validate cloud session imports with zod --- packages/kilo-gateway/src/cloud-sessions.ts | 302 +++++++++++--------- 1 file changed, 161 insertions(+), 141 deletions(-) diff --git a/packages/kilo-gateway/src/cloud-sessions.ts b/packages/kilo-gateway/src/cloud-sessions.ts index 26251db3a60..b80f9ad7310 100644 --- a/packages/kilo-gateway/src/cloud-sessions.ts +++ b/packages/kilo-gateway/src/cloud-sessions.ts @@ -1,3 +1,4 @@ +import { z } from "zod" import { buildKiloHeaders } from "./headers.js" export const UUID_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i @@ -6,18 +7,6 @@ export interface DrizzleDb { insert(table: object): { values(data: object): { onConflictDoNothing(): { run(): void } } } } -type Export = { - info: Record & { - time?: { - created?: number - updated?: number - compacting?: number - archived?: number - } - } - messages?: unknown -} - export interface PrepareDeps { Instance: { readonly directory: string @@ -33,6 +22,132 @@ export interface PrepareDeps { export class SessionImportValidationError extends Error {} +const fileSchema = z + .object({ + id: z.string(), + sessionID: z.string(), + messageID: z.string(), + type: z.literal("file"), + mime: z.string(), + url: z.string(), + }) + .passthrough() + +const stateSchema = z + .object({ + status: z.literal("completed"), + attachments: z.array(fileSchema).optional(), + }) + .passthrough() + +const partSchema = z + .object({ + id: z.string(), + sessionID: z.string(), + messageID: z.string(), + type: z.string(), + tail_start_id: z.unknown().optional(), + state: z.unknown().optional(), + }) + .passthrough() + +const messageSchema = z.object({ + info: z + .object({ + id: z.string(), + sessionID: z.string(), + parentID: z.string().optional(), + role: z.enum(["user", "assistant"]), + time: z.object({ created: z.number().finite() }).passthrough(), + }) + .passthrough(), + parts: z.array(partSchema), +}) + +const exportSchema = z + .object({ + info: z + .object({ + id: z.string(), + time: z + .object({ + created: z.number().optional(), + updated: z.number().optional(), + compacting: z.number().optional(), + archived: z.number().optional(), + }) + .passthrough() + .optional(), + }) + .passthrough(), + messages: z.array(messageSchema), + }) + .superRefine((data, ctx) => { + const ids = new Set() + const pids = new Set() + const parents = new Map() + + for (const msg of data.messages) { + if (msg.info.sessionID !== data.info.id) ctx.addIssue({ code: "custom", message: "Invalid message info" }) + if (ids.has(msg.info.id)) ctx.addIssue({ code: "custom", message: "Duplicate message ID" }) + ids.add(msg.info.id) + parents.set(msg.info.id, msg.info.parentID) + + for (const part of msg.parts) { + if (part.sessionID !== data.info.id || part.messageID !== msg.info.id) + ctx.addIssue({ code: "custom", message: "Invalid message part" }) + if (part.type === "compaction" && part.tail_start_id !== undefined && typeof part.tail_start_id !== "string") + ctx.addIssue({ code: "custom", message: "Invalid compaction tail" }) + if (pids.has(part.id)) ctx.addIssue({ code: "custom", message: "Duplicate part ID" }) + pids.add(part.id) + + if (part.type !== "tool") continue + const status = z.object({ status: z.unknown() }).safeParse(part.state) + if (!status.success || status.data.status !== "completed") continue + const state = stateSchema.safeParse(part.state) + if (!state.success) { + ctx.addIssue({ code: "custom", message: "Invalid tool attachments" }) + continue + } + for (const file of state.data.attachments ?? []) { + if (file.sessionID !== data.info.id || file.messageID !== msg.info.id) + ctx.addIssue({ code: "custom", message: "Invalid tool attachment" }) + if (pids.has(file.id)) ctx.addIssue({ code: "custom", message: "Duplicate part ID" }) + pids.add(file.id) + } + } + } + + for (const msg of data.messages) { + const parent = msg.info.parentID + if (parent !== undefined && !ids.has(parent)) + ctx.addIssue({ code: "custom", message: "Dangling message parent" }) + + const seen = new Set([msg.info.id]) + let current = parent + while (current !== undefined) { + if (seen.has(current)) { + ctx.addIssue({ code: "custom", message: "Circular message parent" }) + break + } + seen.add(current) + current = parents.get(current) + } + + for (const part of msg.parts) { + if (part.type !== "compaction" || typeof part.tail_start_id !== "string") continue + if (!ids.has(part.tail_start_id)) ctx.addIssue({ code: "custom", message: "Dangling compaction tail" }) + } + } + }) + +function completed(part: z.infer) { + if (part.type !== "tool") return + const result = stateSchema.safeParse(part.state) + if (!result.success) return + return result.data +} + const INGEST_BASE = process.env.KILO_SESSION_INGEST_URL ?? "https://ingest.kilosessions.ai" const TIMEOUT = 30_000 @@ -96,123 +211,31 @@ export interface ImportDeps extends PrepareDeps { SessionCreatedEvent: { type: string; properties: unknown } } -function record(input: unknown): input is Record { - return typeof input === "object" && input !== null && !Array.isArray(input) -} - -export function prepareSessionImport(data: Export, deps: PrepareDeps) { - if (!record(data) || !record(data.info) || typeof data.info.id !== "string") - throw new SessionImportValidationError("Invalid session info") - if (!Array.isArray(data.messages)) throw new SessionImportValidationError("Invalid session messages") +export function prepareSessionImport(data: unknown, deps: PrepareDeps) { + const parsed = exportSchema.safeParse(data) + if (!parsed.success) + throw new SessionImportValidationError(parsed.error.issues[0]?.message ?? "Invalid session export") + const source = parsed.data const sessionID = deps.Identifier.descending("session") const ids = new Map() const pids = new Map() - const items: Array<{ - id: string - created: number - info: Record - parent?: string - parts: Array<{ - id: string - data: Record - tail?: string - attachments?: Array<{ id: string; data: Record }> - }> - }> = [] - - for (const msg of data.messages) { - if (!record(msg) || !record(msg.info) || !Array.isArray(msg.parts)) - throw new SessionImportValidationError("Invalid message") - const id = msg.info.id - const role = msg.info.role - const time = msg.info.time - if ( - typeof id !== "string" || - msg.info.sessionID !== data.info.id || - (role !== "user" && role !== "assistant") || - !record(time) || - typeof time.created !== "number" || - !Number.isFinite(time.created) - ) - throw new SessionImportValidationError("Invalid message info") - const parent = msg.info.parentID - if (parent !== undefined && typeof parent !== "string") - throw new SessionImportValidationError("Invalid message parent") - if (ids.has(id)) throw new SessionImportValidationError("Duplicate message ID") - ids.set(id, deps.Identifier.ascending("message")) - - const parts: Array<{ - id: string - data: Record - tail?: string - attachments?: Array<{ id: string; data: Record }> - }> = [] + for (const msg of source.messages) { + ids.set(msg.info.id, deps.Identifier.ascending("message")) for (const part of msg.parts) { - if ( - !record(part) || - typeof part.id !== "string" || - part.sessionID !== data.info.id || - part.messageID !== id || - typeof part.type !== "string" - ) - throw new SessionImportValidationError("Invalid message part") - const tail = part.type === "compaction" ? part.tail_start_id : undefined - if (tail !== undefined && typeof tail !== "string") - throw new SessionImportValidationError("Invalid compaction tail") - if (pids.has(part.id)) throw new SessionImportValidationError("Duplicate part ID") pids.set(part.id, deps.Identifier.ascending("part")) - - const state = - part.type === "tool" && record(part.state) && part.state.status === "completed" ? part.state : undefined - const attachments = (() => { - if (!state || state.attachments === undefined) return - if (!Array.isArray(state.attachments)) throw new SessionImportValidationError("Invalid tool attachments") - const result: Array<{ id: string; data: Record }> = [] - for (const attachment of state.attachments) { - if ( - !record(attachment) || - typeof attachment.id !== "string" || - attachment.sessionID !== data.info.id || - attachment.messageID !== id || - attachment.type !== "file" || - typeof attachment.mime !== "string" || - typeof attachment.url !== "string" - ) - throw new SessionImportValidationError("Invalid tool attachment") - if (pids.has(attachment.id)) throw new SessionImportValidationError("Duplicate part ID") - pids.set(attachment.id, deps.Identifier.ascending("part")) - result.push({ id: attachment.id, data: attachment }) - } - return result - })() - parts.push({ - id: part.id, - data: part, - ...(tail !== undefined ? { tail } : {}), - ...(attachments !== undefined ? { attachments } : {}), - }) - } - items.push({ id, created: time.created, info: msg.info, parts, ...(parent !== undefined ? { parent } : {}) }) - } - - const parents = new Map(items.map((item) => [item.id, item.parent])) - for (const item of items) { - const seen = new Set([item.id]) - let parent = item.parent - while (parent !== undefined) { - if (seen.has(parent)) throw new SessionImportValidationError("Circular message parent") - seen.add(parent) - parent = parents.get(parent) + for (const file of completed(part)?.attachments ?? []) { + pids.set(file.id, deps.Identifier.ascending("part")) + } } } const now = Date.now() const time = { - created: data.info.time?.created ?? now, + created: source.info.time?.created ?? now, updated: now, - ...(data.info.time?.compacting !== undefined && { compacting: data.info.time.compacting }), - ...(data.info.time?.archived !== undefined && { archived: data.info.time.archived }), + ...(source.info.time?.compacting !== undefined && { compacting: source.info.time.compacting }), + ...(source.info.time?.archived !== undefined && { archived: source.info.time.archived }), } const info: Record & { @@ -221,12 +244,12 @@ export function prepareSessionImport(data: Export, deps: PrepareDeps) { directory: string time: typeof time } = { - ...data.info, + ...source.info, id: sessionID, projectID: deps.Instance.project.id, - slug: data.info.slug, + slug: source.info.slug, directory: deps.Instance.directory, - version: data.info.version, + version: source.info.version, time, } delete info.workspaceID @@ -250,40 +273,37 @@ export function prepareSessionImport(data: Export, deps: PrepareDeps) { session_id: string data: Record }> = [] - for (const item of items) { - const id = ids.get(item.id) - if (!id) throw new SessionImportValidationError("Missing message ID") - const parentID = item.parent === undefined ? undefined : ids.get(item.parent) - if (item.parent !== undefined && !parentID) throw new SessionImportValidationError("Dangling message parent") + for (const msg of source.messages) { + const id = ids.get(msg.info.id)! + const parentID = msg.info.parentID === undefined ? undefined : ids.get(msg.info.parentID)! const next = { - ...item.info, + ...msg.info, id, sessionID, ...(parentID ? { parentID } : {}), } - messages.push({ id, session_id: sessionID, time_created: item.created, data: next }) + messages.push({ id, session_id: sessionID, time_created: msg.info.time.created, data: next }) - for (const part of item.parts) { - const partID = pids.get(part.id) - if (!partID) throw new SessionImportValidationError("Missing part ID") - const tail = part.tail === undefined ? undefined : ids.get(part.tail) - if (part.tail !== undefined && !tail) throw new SessionImportValidationError("Dangling compaction tail") + for (const part of msg.parts) { + const partID = pids.get(part.id)! + const tail = + part.type === "compaction" && typeof part.tail_start_id === "string" + ? ids.get(part.tail_start_id)! + : undefined const data: Record = { - ...part.data, + ...part, id: partID, messageID: id, sessionID, ...(tail ? { tail_start_id: tail } : {}), } - if (part.attachments) { - const state = part.data.state - if (!record(state)) throw new SessionImportValidationError("Invalid tool state") + const state = completed(part) + if (state?.attachments) { data.state = { ...state, - attachments: part.attachments.map((attachment) => { - const attachmentID = pids.get(attachment.id) - if (!attachmentID) throw new SessionImportValidationError("Missing attachment ID") - return { ...attachment.data, id: attachmentID, messageID: id, sessionID } + attachments: state.attachments.map((file) => { + const fileID = pids.get(file.id)! + return { ...file, id: fileID, messageID: id, sessionID } }), } } @@ -299,7 +319,7 @@ export function prepareSessionImport(data: Export, deps: PrepareDeps) { return { info, messages, parts } } -export function importSessionToDb(data: Export, deps: ImportDeps) { +export function importSessionToDb(data: unknown, deps: ImportDeps) { const prepared = prepareSessionImport(data, deps) deps.Database.transaction((db) => {