From e72238a6655bb495e24c588fa047b5b162da8f1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Wed, 22 Jul 2026 17:14:18 +0200 Subject: [PATCH] feat(cli): accept mobile file attachments in remote sessions (#12394) * feat(cli): accept mobile file attachments in remote sessions The mobile client uploads each attachment to R2 and sends a first-class FilePartInput with a server-issued . basename. The CLI fetches the file over HTTPS, re-emits it as a data: URL for text / image / PDF, or writes it to a per-session scratch directory for generic binaries so the agent's tools can read it. - Fetches are HTTPS-only, reject redirects, never forward credentials, are bounded to 5 MB + 1 byte (partial deleted on overflow), and time out. - Any per-attachment failure becomes an explanatory text part so the rest of the prompt still runs; the send_message ACK is unaffected because materialization happens inside the long-running dispatch before prompt(). - Scratch directory (0700 / files 0600) lives under Global.Path.tmp and is removed on session deletion and sender dispose. Basenames derive from the attachment id + validated extension, never the client-supplied filename. - The relay heartbeat now advertises capabilities.attachments so the mobile app only enables attachments for CLIs that support them. * fix(cli): secure remote attachment materialization * chore(cli): remove redundant change markers * fix(cli): coordinate remote attachment lifetime * fix(cli): fail closed during attachment cleanup * fix(cli): track idle attachment cleanup --- .changeset/remote-session-file-attachments.md | 5 + .../src/kilo-sessions/remote-protocol.ts | 9 + .../src/kilo-sessions/remote-sender.ts | 152 ++++- .../opencode/src/kilo-sessions/remote-ws.ts | 14 +- .../src/kilocode/remote-attachments.ts | 392 ++++++++++++ .../test/kilocode/remote-attachments.test.ts | 581 ++++++++++++++++++ .../kilocode/sessions/remote-protocol.test.ts | 56 ++ .../kilocode/sessions/remote-sender.test.ts | 377 ++++++++++++ .../test/kilocode/sessions/remote-ws.test.ts | 20 + 9 files changed, 1599 insertions(+), 7 deletions(-) create mode 100644 .changeset/remote-session-file-attachments.md create mode 100644 packages/opencode/src/kilocode/remote-attachments.ts create mode 100644 packages/opencode/test/kilocode/remote-attachments.test.ts diff --git a/.changeset/remote-session-file-attachments.md b/.changeset/remote-session-file-attachments.md new file mode 100644 index 0000000000..2b37b608ac --- /dev/null +++ b/.changeset/remote-session-file-attachments.md @@ -0,0 +1,5 @@ +--- +"@kilocode/cli": minor +--- + +Support file attachments in remote CLI sessions. diff --git a/packages/opencode/src/kilo-sessions/remote-protocol.ts b/packages/opencode/src/kilo-sessions/remote-protocol.ts index d7280b9882..926a025a9d 100644 --- a/packages/opencode/src/kilo-sessions/remote-protocol.ts +++ b/packages/opencode/src/kilo-sessions/remote-protocol.ts @@ -15,10 +15,19 @@ export namespace RemoteProtocol { // --- CLI → DO (Outbound) --- + // Capability flags advertised in the heartbeat so the relay can stop + // probing commands to discover what the CLI supports. Field name and + // nesting are an exact contract with the mobile ingest service. + export const Capabilities = z + .object({ + attachments: z.boolean().optional(), + }) + .optional() export const Heartbeat = z.object({ type: z.literal("heartbeat"), sessions: z.array(SessionInfo), protocolVersion: z.string().optional(), // lets relay detect CLI capabilities without probing commands + capabilities: Capabilities, }) export type Heartbeat = z.infer diff --git a/packages/opencode/src/kilo-sessions/remote-sender.ts b/packages/opencode/src/kilo-sessions/remote-sender.ts index 267677b996..165c221d57 100644 --- a/packages/opencode/src/kilo-sessions/remote-sender.ts +++ b/packages/opencode/src/kilo-sessions/remote-sender.ts @@ -4,6 +4,7 @@ import { RemoteModelCatalog } from "@/kilo-sessions/remote-model-catalog" import { RemoteProtocol } from "@/kilo-sessions/remote-protocol" import type { RemoteWS } from "@/kilo-sessions/remote-ws" import { GlobalBus } from "@/bus/global" +import { RemoteAttachments } from "@/kilocode/remote-attachments" import { Session } from "@/session/session" import type { MessageV2 } from "@/session/message-v2" import { SessionPrompt } from "@/session/prompt" @@ -149,6 +150,12 @@ export namespace RemoteSender { remoteExit?: { get: () => RemoteExit.Callback | undefined } + // Production wires this to RemoteAttachments.create so the scratch dir + // and Session.Event.Deleted cleanup are scoped to the session whose + // parts we are about to materialize. Tests pass a stub that simply + // returns the input so the existing remote-sender suite continues to + // exercise schema/ordering paths without touching the network. + attachments?: (sessionID: SessionID) => RemoteAttachments.Result | undefined } export type Sender = { @@ -273,6 +280,87 @@ export namespace RemoteSender { } }) + // The factory is resolved lazily so tests that never send http(s) file + // parts never trigger scratch-dir setup. The cache and its bus + // listener are owned by this RemoteSender instance and released by + // dispose(), so they survive relay subscribe/unsubscribe churn + // independent of the relay's view of the session set. The bus + // listener is also installed lazily on first use to keep the global + // bus listener count from inflating for senders that never handle + // attachments (the count would otherwise show up in unrelated tests + // that assert it stays at 0). + const attachments = + options.attachments ?? + ((sessionID: SessionID) => RemoteAttachments.create({ sessionID })) + const attachmentCache = new Map() + const pending = new Map() + const retired = new Map() + const cleaning = new Map>() + const deleted = new Set() + let closed = false + let attachmentBusUnsub: (() => void) | undefined + const ensureAttachmentListener = () => { + if (closed || attachmentBusUnsub) return + attachmentBusUnsub = sub((event: any) => { + if (event?.type !== Session.Event.Deleted.type) return + const sid = event?.properties?.sessionID + if (typeof sid !== "string") return + const id = SessionID.make(sid) + const result = attachmentCache.get(id) + if (!result) { + if (pending.has(id)) deleted.add(id) + return + } + deleted.add(id) + attachmentCache.delete(id) + if (pending.has(id)) { + retired.set(id, result) + return + } + void clean(id, result) + }) + } + function begin(id: SessionID) { + pending.set(id, (pending.get(id) ?? 0) + 1) + } + function clean(id: SessionID, result: RemoteAttachments.Result) { + const existing = cleaning.get(id) + if (existing) return existing + const cleanup = result + .dispose() + .catch((error) => options.log.warn("attachment cleanup failed", { error: String(error) })) + cleaning.set(id, cleanup) + void cleanup.finally(() => { + if (cleaning.get(id) === cleanup) cleaning.delete(id) + if (!pending.has(id)) deleted.delete(id) + }) + return cleanup + } + async function finish(id: SessionID) { + const count = pending.get(id) + if (!count) return + if (count > 1) { + pending.set(id, count - 1) + return + } + pending.delete(id) + const result = retired.get(id) + const cleanup = cleaning.get(id) ?? (result ? clean(id, result) : undefined) + if (result) retired.delete(id) + if (cleanup) await cleanup + if (!pending.has(id)) deleted.delete(id) + } + function attachmentFor(sessionID: SessionID): RemoteAttachments.Result | undefined { + if (closed || deleted.has(sessionID)) return undefined + const existing = attachmentCache.get(sessionID) + if (existing) return existing + const next = attachments(sessionID) + if (!next) return undefined + ensureAttachmentListener() + attachmentCache.set(sessionID, next) + return next + } + async function directoryFor(sid: string): Promise { const info = await session.get(SessionID.make(sid)).catch(() => undefined) return info?.directory ?? options.directory @@ -425,18 +513,40 @@ export namespace RemoteSender { }) } - function dispatchLongRunning(msg: RemoteProtocol.Command, dir: Promise, work: () => Promise) { + function dispatchLongRunning( + msg: RemoteProtocol.Command, + dir: Promise, + work: () => Promise, + settle?: () => void | Promise, + ) { const run = options.provide ?? provide + let settled = false + const complete = () => { + if (settled) return + settled = true + void settle?.() + } options.conn.send({ type: "response", id: msg.id, result: {} }) void (async () => { try { - await run({ directory: await dir, fn: work }) + await run({ + directory: await dir, + fn: async () => { + try { + await work() + } finally { + complete() + } + }, + }) } catch (e) { options.log.error("long-running command failed after ACK", { id: msg.id, command: msg.command, error: String(e), }) + } finally { + complete() } })() } @@ -691,9 +801,28 @@ export namespace RemoteSender { return } const promptInput = { ...input.data, ephemeralTools: normalized.ephemeralTools } as SessionPrompt.PromptInput - dispatchLongRunning(msg, directoryFor(promptInput.sessionID), async () => { - await prompt(promptInput) - }) + const remote = promptInput.parts.some((part) => part.type === "file" && RemoteAttachments.isFetchable(part.url)) + if (remote) { + begin(promptInput.sessionID) + ensureAttachmentListener() + } + dispatchLongRunning( + msg, + directoryFor(promptInput.sessionID), + async () => { + // Runs strictly after the synchronous ACK above and strictly before the + // existing prompt() call so the resolvePart boundary sees data: URLs + // and a scratch path instead of an http(s) URL it cannot fetch. + const materializer = remote ? attachmentFor(promptInput.sessionID) : undefined + if (materializer) { + promptInput.parts = await materializer.materialize(promptInput.parts) + } else if (remote) { + promptInput.parts = RemoteAttachments.failClosed(promptInput.parts) + } + await prompt(promptInput) + }, + remote ? () => finish(promptInput.sessionID) : undefined, + ) return } if (msg.command === "interrupt") { @@ -835,10 +964,23 @@ export namespace RemoteSender { } function dispose() { + closed = true if (unsub) { unsub() unsub = undefined } + // per-session materializers. Fire-and-forget the async dispose because + // RemoteAttachments.dispose() is best-effort scratch cleanup. + attachmentBusUnsub?.() + attachmentBusUnsub = undefined + for (const [id, result] of attachmentCache) { + if (pending.has(id)) { + retired.set(id, result) + continue + } + void result.dispose() + } + attachmentCache.clear() sessions.clear() children.clear() } diff --git a/packages/opencode/src/kilo-sessions/remote-ws.ts b/packages/opencode/src/kilo-sessions/remote-ws.ts index 0e8c6187c6..de631d23e1 100644 --- a/packages/opencode/src/kilo-sessions/remote-ws.ts +++ b/packages/opencode/src/kilo-sessions/remote-ws.ts @@ -233,7 +233,12 @@ export namespace RemoteWS { if (fresh !== undefined) { lastGood = fresh const sentLive = ws?.readyState === WebSocket.OPEN - send({ type: "heartbeat", protocolVersion: InstallationVersion, sessions: fresh }) + send({ + type: "heartbeat", + protocolVersion: InstallationVersion, + capabilities: { attachments: true }, + sessions: fresh, + }) if (sentLive) { // A waiter requiring a specific id is satisfied only when // the sent payload contains that id. Unsatisfied waiters @@ -264,7 +269,12 @@ export namespace RemoteWS { } else { // Degraded: preserve liveness with the last known-good list (empty // on cold start) and keep waiters pending for a future fresh send. - send({ type: "heartbeat", protocolVersion: InstallationVersion, sessions: lastGood ?? [] }) + send({ + type: "heartbeat", + protocolVersion: InstallationVersion, + capabilities: { attachments: true }, + sessions: lastGood ?? [], + }) waiters = cycleWaiters.concat(waiters) } } diff --git a/packages/opencode/src/kilocode/remote-attachments.ts b/packages/opencode/src/kilocode/remote-attachments.ts new file mode 100644 index 0000000000..9008065cc9 --- /dev/null +++ b/packages/opencode/src/kilocode/remote-attachments.ts @@ -0,0 +1,392 @@ +// Helper for materializing remote-session file attachments on the CLI side. +// +// Mobile uploads bytes to R2 and sends the CLI a first-class +// `FilePartInput = { id?, type: 'file', mime, filename?, url }` whose `url` +// is an HTTPS R2 presigned GET and whose `filename` is the server-issued +// `.` basename. The CLI must fetch those bytes and turn the part +// into a shape the existing `resolvePart` boundary can consume: +// +// - text/*, image/*, application/pdf → emit a `data:` URL file part +// (text canonicalizes to `text/plain`; PDF and image pass through to the +// existing PDF model modality / image normalization paths in +// `provider/transform.ts` and `resolvePart`) +// - any other extension → write to a per-session scratch directory and +// emit a synthetic text part describing the absolute path, filename, +// MIME, and size +// +// The helper is invoked only from `remote-sender.ts`; local prompts and +// other ingress paths never touch this code (decision 6). +import path from "node:path" +import fs from "node:fs/promises" +import { Global } from "@opencode-ai/core/global" +import * as Log from "@opencode-ai/core/util/log" +import { PartID, type SessionID } from "@/session/schema" +import type { SessionPrompt } from "@/session/prompt" + +const log = Log.create({ service: "remote-attachments" }) + +export namespace RemoteAttachments { + // Decision 4: canonical extension → MIME table. Text entries all + // canonicalize to `text/plain` at re-entry (per text caveat in the + // design). The binary fallback is `application/octet-stream` and is + // applied to any extension not present here AND to extensionless inputs. + export const EXTENSION_MIME: Record = { + png: "image/png", + jpg: "image/jpeg", + jpeg: "image/jpeg", + webp: "image/webp", + gif: "image/gif", + pdf: "application/pdf", + txt: "text/plain", + md: "text/plain", + csv: "text/plain", + log: "text/plain", + json: "text/plain", + xml: "text/plain", + yaml: "text/plain", + yml: "text/plain", + toml: "text/plain", + ini: "text/plain", + html: "text/plain", + css: "text/plain", + js: "text/plain", + jsx: "text/plain", + ts: "text/plain", + tsx: "text/plain", + py: "text/plain", + rb: "text/plain", + go: "text/plain", + rs: "text/plain", + java: "text/plain", + c: "text/plain", + h: "text/plain", + cpp: "text/plain", + hpp: "text/plain", + sh: "text/plain", + sql: "text/plain", + } + export const BINARY_MIME = "application/octet-stream" + export const TEXT_PLAIN = "text/plain" + // Hard cap on attachment bytes (5 MB + 1 byte so the helper aborts + // strictly when the body exceeds the agreed ceiling). + export const MAX_BYTES = 5 * 1024 * 1024 + 1 + // Per-attachment fetch budget. R2 presigned GETs in the same region + // complete in tens of ms; 15s is generous but bounded so a stalled + // connection can never hold the prompt open indefinitely. + export const FETCH_TIMEOUT_MS = 15_000 + export const SCRATCH_DIRNAME = "remote-attachments" + + export type Fetcher = (input: string, init?: RequestInit) => Promise + + export type Deps = { + sessionID: SessionID + /** Override the scratch root. Defaults to `Global.Path.tmp`. */ + tmpRoot?: string + /** Override `fetch` (used to inject mock responses in tests). */ + fetch?: Fetcher + /** Per-call logger. Defaults to the module logger. */ + log?: { + warn: (msg: string, meta?: unknown) => void + error: (msg: string, meta?: unknown) => void + } + } + + export type Result = { + materialize: (parts: SessionPrompt.PromptInput["parts"]) => Promise + dispose: () => Promise + } + + /** Extract the lowercased extension from a server-issued basename. */ + export function extensionOf(filename: string | undefined): string { + if (!filename) return "bin" + const i = filename.lastIndexOf(".") + if (i < 0 || i === filename.length - 1) return "bin" + return filename.slice(i + 1).toLowerCase() + } + + /** Return the canonical MIME for the given extension. */ + export function mimeFor(extension: string): string { + return EXTENSION_MIME[extension] ?? BINARY_MIME + } + + /** + * Validate that an extension is a safe single token usable as a file + * suffix. Anything that would expand to a path segment beyond a single + * component is rejected and falls back to "bin" at the call site. + */ + export function safeExtension(extension: string): string { + return /^[A-Za-z0-9]{1,16}$/.test(extension) ? extension : "bin" + } + + /** Classify by extension. Equivalent to `mimeFor(extensionOf(filename))`. */ + export function classify(filename: string | undefined): { mime: string; extension: string } { + const ext = safeExtension(extensionOf(filename)) + return { mime: mimeFor(ext), extension: ext } + } + + /** True when the URL points at an HTTP resource that needs validation. */ + export function isFetchable(url: string): boolean { + return /^https?:\/\//i.test(url) + } + + /** Reason a fetch failed. */ + export type FetchError = { + kind: "https" | "host" | "redirect" | "non-2xx" | "overflow" | "timeout" | "network" + message: string + status?: number + } + + function makeError(kind: FetchError["kind"], message: string, status?: number): FetchError & Error { + const e = new Error(message) as Error & FetchError + e.kind = kind + if (status !== undefined) e.status = status + return e + } + + async function readBounded(response: Response, signal: AbortSignal): Promise { + const body = response.body + if (!body) return new Uint8Array() + const chunks: Uint8Array[] = [] + let total = 0 + const reader = body.getReader() + const aborted = new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => reject(makeError("timeout", "attachment fetch timed out")), { once: true }) + }) + try { + while (true) { + const { done, value } = await Promise.race([reader.read(), aborted]) + if (done) break + if (!value) continue + total += value.byteLength + if (total >= MAX_BYTES) { + try { + await reader.cancel() + } catch { + // best-effort: the read is already failing the bound + } + throw makeError("overflow", `attachment exceeds ${MAX_BYTES - 1} bytes`) + } + chunks.push(value) + } + } catch (err) { + if (err && typeof err === "object" && "kind" in err) throw err + if (signal.aborted) { + throw makeError("timeout", "attachment fetch timed out") + } + throw makeError("network", "attachment fetch failed") + } finally { + try { + reader.releaseLock() + } catch { + // reader already detached; nothing to do + } + } + const out = new Uint8Array(total) + let offset = 0 + for (const chunk of chunks) { + out.set(chunk, offset) + offset += chunk.byteLength + } + return out + } + + /** + * Fetch a single R2 presigned URL with the full decision-6 safety net: + * - HTTPS only + * - redirects rejected + * - no credentials forwarded + * - body bounded to 5 MB + 1 byte + * - bounded timeout + * - non-2xx rejected + */ + export async function fetchOne(url: string, deps?: { fetch?: Fetcher; timeoutMs?: number }): Promise { + const parsed = (() => { + try { + return new URL(url) + } catch { + throw makeError("https", "attachment url is not valid") + } + })() + if (parsed.protocol !== "https:") { + throw makeError("https", `attachment url must use https (got ${parsed.protocol.replace(":", "")})`) + } + if (parsed.username || parsed.password) throw makeError("https", "attachment url must not include credentials") + if (!parsed.hostname.endsWith(".r2.cloudflarestorage.com")) { + throw makeError("host", "attachment url must use a Cloudflare R2 host") + } + const f = deps?.fetch ?? (globalThis.fetch as Fetcher | undefined) + if (!f) throw makeError("network", "no fetch implementation available") + const controller = new AbortController() + const timeoutMs = deps?.timeoutMs ?? FETCH_TIMEOUT_MS + const timer = setTimeout(() => controller.abort(new DOMException("timeout", "AbortError")), timeoutMs) + try { + const response = await f(url, { + method: "GET", + redirect: "error", + credentials: "omit", + signal: controller.signal, + }) + if (response.status < 200 || response.status >= 300) { + throw makeError("non-2xx", `attachment fetch returned ${response.status}`, response.status) + } + return await readBounded(response, controller.signal) + } catch (err) { + if (err && typeof err === "object" && "kind" in err) throw err + if (controller.signal.aborted) { + throw makeError("timeout", `attachment fetch timed out after ${timeoutMs}ms`) + } + const msg = err instanceof Error ? err.message : String(err) + if (/redirect/i.test(msg)) { + throw makeError("redirect", "attachment url redirected") + } + throw makeError("network", "attachment fetch failed") + } finally { + clearTimeout(timer) + } + } + + /** Build the explanatory text part that replaces a failed attachment. */ + export function failureText(filename: string | undefined, reason: string): { type: "text"; text: string } { + const name = filename ?? "attachment" + return { + type: "text", + text: `attachment ${name} could not be retrieved: ${reason}`, + } + } + + export function failClosed(parts: SessionPrompt.PromptInput["parts"]): SessionPrompt.PromptInput["parts"] { + return parts.map((part) => + part.type === "file" && isFetchable(part.url) ? failureText(part.filename, "attachment session is closed") : part, + ) + } + + /** + * Materialize a list of parts. Non-file parts are passed through + * unchanged. File parts whose URL is http(s) are fetched and replaced + * with a data: URL file part (text/image/pdf) or a scratch-file text + * part (binary). All other URL schemes (data:, file:, …) are passed + * through unchanged so the existing `resolvePart` boundary can decide. + * On any per-part failure the original part is replaced with an + * explanatory text part and the rest of the prompt still proceeds. + */ + export function create(deps: Deps): Result { + const sessionID = deps.sessionID + const root = deps.tmpRoot ?? Global.Path.tmp + const scratchDir = path.join(root, SCRATCH_DIRNAME, Buffer.from(sessionID).toString("base64url")) + const writer = deps.log ?? { + warn: (msg: string, meta?: unknown) => log.warn(msg, meta as never), + error: (msg: string, meta?: unknown) => log.error(msg, meta as never), + } + const f = deps.fetch ?? (globalThis.fetch as Fetcher | undefined) + let closed = false + let disposal: Promise | undefined + const active = new Set>() + const cleanup = async () => { + try { + await fs.rm(scratchDir, { recursive: true, force: true }) + } catch (err) { + writer.warn("scratch dir cleanup failed", { sessionID, error: String(err) }) + } + } + const run = async (parts: SessionPrompt.PromptInput["parts"]): Promise => { + const out: SessionPrompt.PromptInput["parts"] = [] + for (const part of parts) { + if (!part || typeof part !== "object" || part.type !== "file") { + out.push(part) + continue + } + const url = part.url + if (!isFetchable(url)) { + out.push(part) + continue + } + const filename = part.filename + const { extension } = classify(filename) + const id = part.id ?? PartID.make(`prt_${crypto.randomUUID()}`) + const basename = `${crypto.randomUUID()}.${extension}` + const target = path.join(scratchDir, basename) + try { + const bytes = await fetchOne(url, { fetch: f }) + if (extension === "pdf") { + // PDF falls through to the existing PDF modality + // (`provider/transform.ts`); the helper hands it back as an + // application/pdf file part with a data: URL. + out.push({ + id, + type: "file" as const, + mime: "application/pdf", + filename, + url: dataUrl("application/pdf", bytes), + }) + continue + } + if (mimeFor(extension) === BINARY_MIME) { + // Binary fallback — persist to scratch and surface a text part. + try { + await fs.mkdir(scratchDir, { recursive: true, mode: 0o700 }) + await fs.writeFile(target, bytes, { mode: 0o600 }) + } catch (err) { + await fs + .rm(target, { force: true }) + .catch((cleanupError) => + writer.warn("partial scratch file cleanup failed", { sessionID, error: String(cleanupError) }), + ) + writer.error("scratch write failed", { sessionID, error: String(err) }) + out.push(failureText(filename, `local write failed: ${err instanceof Error ? err.message : String(err)}`)) + continue + } + out.push({ + type: "text" as const, + text: + `attachment saved to ${target} (filename: ${filename ?? basename}, mime: ${BINARY_MIME}, size: ${bytes.byteLength} bytes). ` + + `Use the read tool on that path to inspect it.`, + }) + continue + } + // text or image — re-enter as a data: URL. + const mime = mimeFor(extension) + out.push({ + id, + type: "file" as const, + mime, + filename, + url: dataUrl(mime, bytes), + }) + } catch (err) { + const reason = + err && typeof err === "object" && "message" in err ? String((err as Error).message) : String(err) + writer.warn("attachment fetch failed", { sessionID, filename, error: reason }) + out.push(failureText(filename, reason)) + } + } + return out + } + + const materialize = (parts: SessionPrompt.PromptInput["parts"]): Promise => { + if (closed) { + return Promise.resolve(failClosed(parts)) + } + const job = run(parts) + active.add(job) + void job.then( + () => active.delete(job), + () => active.delete(job), + ) + return job + } + + const dispose = async () => { + if (disposal) return disposal + closed = true + disposal = Promise.allSettled([...active]).then(cleanup) + return disposal + } + + return { materialize, dispose } + } + + /** Build a `data:` URL from a buffer. */ + export function dataUrl(mime: string, bytes: Uint8Array): string { + return `data:${mime};base64,${Buffer.from(bytes).toString("base64")}` + } +} diff --git a/packages/opencode/test/kilocode/remote-attachments.test.ts b/packages/opencode/test/kilocode/remote-attachments.test.ts new file mode 100644 index 0000000000..4c1289f691 --- /dev/null +++ b/packages/opencode/test/kilocode/remote-attachments.test.ts @@ -0,0 +1,581 @@ +import { describe, expect, test } from "bun:test" +import fs from "node:fs/promises" +import os from "node:os" +import path from "node:path" +import { RemoteAttachments } from "../../src/kilocode/remote-attachments" +import { PartID, SessionID } from "../../src/session/schema" + +async function tmpRoot() { + return fs.mkdtemp(path.join(os.tmpdir(), "remote-attachments-test-")) +} + +function scratch(root: string, sessionID: string) { + return path.join(root, RemoteAttachments.SCRATCH_DIRNAME, Buffer.from(sessionID).toString("base64url")) +} + +const nolog = { + warn: () => {}, + error: () => {}, +} + +function okResponse(body: Uint8Array | string, init: ResponseInit = {}): Response { + const bytes = typeof body === "string" ? new TextEncoder().encode(body) : body + return new Response(bytes as BodyInit, { status: 200, ...init }) +} + +function jsonResponse(status: number) { + return new Response(JSON.stringify({ ok: false }), { status }) +} + +describe("RemoteAttachments.classify / EXTENSION_MIME", () => { + test("every table entry classifies to its declared MIME", () => { + for (const [ext, mime] of Object.entries(RemoteAttachments.EXTENSION_MIME)) { + const r = RemoteAttachments.classify(`file.${ext}`) + expect(r.mime).toBe(mime) + expect(r.extension).toBe(ext) + } + }) + + test("fallback returns application/octet-stream for unknown extensions", () => { + const r = RemoteAttachments.classify("file.xyz") + expect(r.mime).toBe("application/octet-stream") + expect(r.extension).toBe("xyz") + }) + + test("extensionless file falls back to bin → application/octet-stream", () => { + const r = RemoteAttachments.classify("README") + expect(r.extension).toBe("bin") + expect(r.mime).toBe("application/octet-stream") + }) + + test("undefined filename falls back to bin → application/octet-stream", () => { + const r = RemoteAttachments.classify(undefined) + expect(r.extension).toBe("bin") + expect(r.mime).toBe("application/octet-stream") + }) + + test("dot-suffix file (no extension chars) falls back to bin", () => { + const r = RemoteAttachments.classify("file.") + expect(r.extension).toBe("bin") + expect(r.mime).toBe("application/octet-stream") + }) + + test("case-insensitive extension lookup", () => { + const r = RemoteAttachments.classify("PHOTO.PNG") + expect(r.extension).toBe("png") + expect(r.mime).toBe("image/png") + }) + + test("cross-surface: a valid extension not in the canonical table is the binary fallback", () => { + // `.weirdo` is a real, safe single-token extension but not in EXTENSION_MIME, + // so it must classify as the binary fallback (no MIME-based extension inference). + const r = RemoteAttachments.classify("attachment.weirdo") + expect(r.extension).toBe("weirdo") + expect(r.mime).toBe("application/octet-stream") + expect(RemoteAttachments.mimeFor("weirdo")).toBe("application/octet-stream") + }) + + test("safeExtension rejects path-traversal and oversized tokens", () => { + expect(RemoteAttachments.safeExtension("../../etc/passwd")).toBe("bin") + expect(RemoteAttachments.safeExtension("a/b")).toBe("bin") + expect(RemoteAttachments.safeExtension("a b")).toBe("bin") + expect(RemoteAttachments.safeExtension("a".repeat(17))).toBe("bin") + expect(RemoteAttachments.safeExtension("png")).toBe("png") + }) +}) + +describe("RemoteAttachments.isFetchable", () => { + test("matches http and https", () => { + expect(RemoteAttachments.isFetchable("https://acct.r2.cloudflarestorage.com/abc")).toBe(true) + expect(RemoteAttachments.isFetchable("http://acct.r2.cloudflarestorage.com/abc")).toBe(true) + expect(RemoteAttachments.isFetchable("HTTPS://acct.r2.cloudflarestorage.com/abc")).toBe(true) + }) + + test("rejects non-http schemes", () => { + expect(RemoteAttachments.isFetchable("data:text/plain,hi")).toBe(false) + expect(RemoteAttachments.isFetchable("file:///etc/passwd")).toBe(false) + expect(RemoteAttachments.isFetchable("ftp://acct.r2.cloudflarestorage.com/abc")).toBe(false) + }) +}) + +describe("RemoteAttachments.failureText", () => { + test("uses filename when provided", () => { + const t = RemoteAttachments.failureText("report.csv", "boom") + expect(t).toEqual({ type: "text", text: "attachment report.csv could not be retrieved: boom" }) + }) + + test("uses generic label when filename is missing", () => { + const t = RemoteAttachments.failureText(undefined, "boom") + expect(t.text).toContain("attachment attachment could not be retrieved") + }) +}) + +describe("RemoteAttachments.dataUrl", () => { + test("encodes bytes as base64 with the right mime prefix", () => { + const bytes = new TextEncoder().encode("hello world") + const url = RemoteAttachments.dataUrl("text/plain", bytes) + expect(url.startsWith("data:text/plain;base64,")).toBe(true) + const base64 = url.slice("data:text/plain;base64,".length) + expect(Buffer.from(base64, "base64").toString("utf8")).toBe("hello world") + }) +}) + +describe("RemoteAttachments.fetchOne safety", () => { + test("rejects http:// (not https)", async () => { + let f: any + try { + await expect( + RemoteAttachments.fetchOne("http://acct.r2.cloudflarestorage.com/abc", { + fetch: (f = async () => okResponse("x")), + }), + ).rejects.toMatchObject({ kind: "https" }) + } finally { + void f + } + }) + + test("rejects malformed URL", async () => { + const secret = "secret-token" + const err = await RemoteAttachments.fetchOne(`not a url?token=${secret}`, { + fetch: async () => okResponse("x"), + }).then( + () => undefined, + (error) => error as Error, + ) + expect(err).toMatchObject({ kind: "https" }) + expect(err?.message).not.toContain(secret) + }) + + test("accepts only Cloudflare R2 hosts", async () => { + const f = async () => okResponse("x") + await expect( + RemoteAttachments.fetchOne("https://acct.r2.cloudflarestorage.com/file", { fetch: f }), + ).resolves.toBeInstanceOf(Uint8Array) + await expect( + RemoteAttachments.fetchOne("https://bucket.acct.r2.cloudflarestorage.com/file", { fetch: f }), + ).resolves.toBeInstanceOf(Uint8Array) + for (const url of [ + "https://r2.cloudflarestorage.com/file", + "https://r2.cloudflarestorage.com.evil.test/file", + "https://internal.example/file", + "https://user:pass@acct.r2.cloudflarestorage.com/file", + ]) { + await expect(RemoteAttachments.fetchOne(url, { fetch: f })).rejects.toBeDefined() + } + }) + + test("rejects redirects by passing redirect: error", async () => { + let captured: RequestInit | undefined + const f = async (_url: string, init?: RequestInit) => { + captured = init + // Simulate fetch rejecting when redirect: error + a redirect response + const err = new TypeError("Failed to fetch: redirect not allowed") + throw err + } + await expect( + RemoteAttachments.fetchOne("https://acct.r2.cloudflarestorage.com/abc", { fetch: f }), + ).rejects.toMatchObject({ + kind: "redirect", + }) + expect(captured?.redirect).toBe("error") + }) + + test("rejects non-2xx responses", async () => { + const f = async () => jsonResponse(404) + await expect( + RemoteAttachments.fetchOne("https://acct.r2.cloudflarestorage.com/missing", { fetch: f }), + ).rejects.toMatchObject({ + kind: "non-2xx", + status: 404, + }) + }) + + test("aborts + rejects when body exceeds MAX_BYTES", async () => { + const oversize = new Uint8Array(RemoteAttachments.MAX_BYTES) + // Build a stream that yields the entire oversize buffer in one chunk so the + // bounded reader trips the overflow guard. + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(oversize) + controller.close() + }, + }) + const f = async () => new Response(stream, { status: 200 }) + await expect( + RemoteAttachments.fetchOne("https://acct.r2.cloudflarestorage.com/big", { fetch: f }), + ).rejects.toMatchObject({ + kind: "overflow", + }) + }) + + test("aborts + rejects on timeout", async () => { + // A never-resolving fetch should be aborted by the timeout. + const f = (_url: string, init?: RequestInit) => + new Promise((_resolve, reject) => { + init?.signal?.addEventListener("abort", () => { + reject(new DOMException("aborted", "AbortError")) + }) + }) + await expect( + RemoteAttachments.fetchOne("https://acct.r2.cloudflarestorage.com/slow", { fetch: f, timeoutMs: 25 }), + ).rejects.toMatchObject({ kind: "timeout" }) + }) + + test("keeps the timeout active while reading the body", async () => { + const stream = new ReadableStream({ pull: () => new Promise(() => {}) }) + const f = async () => new Response(stream, { status: 200 }) + await expect( + RemoteAttachments.fetchOne("https://acct.r2.cloudflarestorage.com/stalled", { fetch: f, timeoutMs: 25 }), + ).rejects.toMatchObject({ kind: "timeout" }) + }) + + test("forwards credentials: omit (no cookies on the wire)", async () => { + let captured: RequestInit | undefined + const f = async (_url: string, init?: RequestInit) => { + captured = init + return okResponse("hi") + } + await RemoteAttachments.fetchOne("https://acct.r2.cloudflarestorage.com/abc", { fetch: f }) + expect(captured?.credentials).toBe("omit") + }) +}) + +describe("RemoteAttachments.create().materialize", () => { + test("returns the input list when it has no file parts", async () => { + const root = await tmpRoot() + try { + const r = RemoteAttachments.create({ sessionID: SessionID.make("ses_a"), tmpRoot: root, log: nolog }) + const out = await r.materialize([{ type: "text", text: "hi" }]) + expect(out).toEqual([{ type: "text", text: "hi" }]) + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("passes through non-fetchable URLs unchanged (data:, file:)", async () => { + const root = await tmpRoot() + try { + const r = RemoteAttachments.create({ sessionID: SessionID.make("ses_a"), tmpRoot: root, log: nolog }) + const part = { type: "file" as const, mime: "image/png", filename: "a.png", url: "data:image/png;base64,AAAA" } + const out = await r.materialize([part]) + expect(out).toEqual([part]) + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("fetches a CSV and canonicalizes it to a text/plain data: URL", async () => { + const root = await tmpRoot() + try { + const f = async (url: string) => okResponse("a,b\n1,2\n") + const r = RemoteAttachments.create({ + sessionID: SessionID.make("ses_csv"), + tmpRoot: root, + fetch: f, + log: nolog, + }) + const out = await r.materialize([ + { + id: PartID.make("prt_csv"), + type: "file", + mime: "text/csv", + filename: "report.csv", + url: "https://acct.r2.cloudflarestorage.com/report.csv", + }, + ]) + expect(out).toHaveLength(1) + const file = out[0] as any + expect(file.type).toBe("file") + expect(file.id).toBe("prt_csv") + expect(file.mime).toBe("text/plain") + expect(file.filename).toBe("report.csv") + expect(file.url).toStartWith("data:text/plain;base64,") + expect(file).not.toHaveProperty("source") + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("fetches a PNG and emits an image/png data: URL file part", async () => { + const root = await tmpRoot() + try { + const pngBytes = new Uint8Array([0x89, 0x50, 0x4e, 0x47]) + const f = async () => okResponse(pngBytes) + const r = RemoteAttachments.create({ + sessionID: SessionID.make("ses_png"), + tmpRoot: root, + fetch: f, + log: nolog, + }) + const out = await r.materialize([ + { + id: PartID.make("prt_png"), + type: "file", + mime: "image/png", + filename: "photo.png", + url: "https://acct.r2.cloudflarestorage.com/photo.png", + }, + ]) + expect(out).toHaveLength(1) + const file = out[0] as any + expect(file.type).toBe("file") + expect(file.mime).toBe("image/png") + expect(file.url).toStartWith("data:image/png;base64,") + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("fetches a PDF and emits an application/pdf data: URL file part (NOT generic binary)", async () => { + const root = await tmpRoot() + try { + const pdfBytes = new Uint8Array([0x25, 0x50, 0x44, 0x46]) // %PDF + const f = async () => okResponse(pdfBytes) + const r = RemoteAttachments.create({ + sessionID: SessionID.make("ses_pdf"), + tmpRoot: root, + fetch: f, + log: nolog, + }) + const out = await r.materialize([ + { + id: PartID.make("prt_pdf"), + type: "file", + mime: "application/pdf", + filename: "doc.pdf", + url: "https://acct.r2.cloudflarestorage.com/doc.pdf", + }, + ]) + expect(out).toHaveLength(1) + const file = out[0] as any + expect(file.type).toBe("file") + expect(file.mime).toBe("application/pdf") + expect(file.url).toStartWith("data:application/pdf;base64,") + // No text part was emitted — the PDF falls through to the existing resolvePart path. + const dir = scratch(root, "ses_pdf") + const entries = await fs.readdir(dir).catch(() => [] as string[]) + expect(entries).toEqual([]) + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("writes a generic binary attachment to the scratch dir and emits a text part with absolute path", async () => { + const root = await tmpRoot() + try { + const bin = new Uint8Array([0, 1, 2, 3, 4, 5, 6, 7]) + const f = async () => okResponse(bin) + const r = RemoteAttachments.create({ + sessionID: SessionID.make("ses_bin"), + tmpRoot: root, + fetch: f, + log: nolog, + }) + const out = await r.materialize([ + { + id: PartID.make("prt_bin"), + type: "file", + mime: "application/octet-stream", + filename: "blob.bin", + url: "https://acct.r2.cloudflarestorage.com/blob.bin", + }, + ]) + expect(out).toHaveLength(1) + const text = out[0] as any + expect(text.type).toBe("text") + const dir = scratch(root, "ses_bin") + expect(text.text).toContain(dir) + expect(text.text).toContain("filename: blob.bin") + expect(text.text).toContain("mime: application/octet-stream") + expect(text.text).toContain(`size: ${bin.byteLength} bytes`) + + const entries = await fs.readdir(dir) + expect(entries).toHaveLength(1) + expect(entries[0]).toMatch(/^[0-9a-f-]{36}\.bin$/) + const written = await fs.readFile(path.join(dir, entries[0]!)) + expect(Array.from(written)).toEqual(Array.from(bin)) + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("confines schema-valid malicious attachment and session ids", async () => { + const root = await tmpRoot() + try { + const f = async () => okResponse(new Uint8Array([0])) + const sessionID = SessionID.make("ses_../../escaped") + const r = RemoteAttachments.create({ + sessionID, + tmpRoot: root, + fetch: f, + log: nolog, + }) + // Both branded IDs satisfy their schemas because only the prefix is checked. + const out = await r.materialize([ + { + id: PartID.make("prt_../../escaped"), + type: "file", + mime: "application/octet-stream", + filename: "../../../etc/passwd", + url: "https://acct.r2.cloudflarestorage.com/x.bin", + }, + ]) + const text = out[0] as any + const base = path.join(root, RemoteAttachments.SCRATCH_DIRNAME) + const dir = scratch(root, sessionID) + const entries = await fs.readdir(dir) + expect(path.relative(base, dir)).not.toStartWith("..") + expect(entries).toHaveLength(1) + expect(entries[0]).toMatch(/^[0-9a-f-]{36}\.bin$/) + expect(text.text).toContain(path.join(dir, entries[0]!)) + expect(text.text).not.toContain("prt_../../escaped") + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("generates a uuid basename when attachmentId is missing", async () => { + const root = await tmpRoot() + try { + const f = async () => okResponse(new Uint8Array([0])) + const r = RemoteAttachments.create({ + sessionID: SessionID.make("ses_noid"), + tmpRoot: root, + fetch: f, + log: nolog, + }) + const out = await r.materialize([ + { + type: "file", + mime: "application/octet-stream", + filename: "blob.bin", + url: "https://acct.r2.cloudflarestorage.com/blob.bin", + }, + ]) + const text = out[0] as any + const dir = scratch(root, "ses_noid") + const entries = await fs.readdir(dir) + expect(entries).toHaveLength(1) + const name = entries[0]! + expect(name.endsWith(".bin")).toBe(true) + expect(name).not.toContain("blob") + expect(text.text).toContain(name) + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("replaces a failed attachment with an explanatory text part and keeps the rest of the prompt", async () => { + const root = await tmpRoot() + try { + const f = async (url: string) => { + if (url.endsWith("good.png")) return okResponse(new Uint8Array([1, 2, 3])) + return jsonResponse(500) + } + const r = RemoteAttachments.create({ + sessionID: SessionID.make("ses_mix"), + tmpRoot: root, + fetch: f, + log: nolog, + }) + const out = await r.materialize([ + { type: "text", text: "see attached" }, + { + type: "file", + mime: "image/png", + filename: "good.png", + url: "https://acct.r2.cloudflarestorage.com/good.png", + }, + { type: "file", mime: "image/png", filename: "bad.png", url: "https://acct.r2.cloudflarestorage.com/bad.png" }, + ]) + expect(out).toHaveLength(3) + expect((out[0] as any).type).toBe("text") + expect((out[1] as any).type).toBe("file") + expect((out[2] as any).type).toBe("text") + expect((out[2] as any).text).toContain("attachment bad.png could not be retrieved") + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("does not expose attachment URLs in errors or logs", async () => { + const root = await tmpRoot() + try { + const secret = "secret-token" + const url = `https://acct.r2.cloudflarestorage.com/file?token=${secret}` + const logs: unknown[] = [] + const r = RemoteAttachments.create({ + sessionID: SessionID.make("ses_secret"), + tmpRoot: root, + fetch: async () => { + throw new Error(url) + }, + log: { + warn: (_msg, meta) => logs.push(meta), + error: (_msg, meta) => logs.push(meta), + }, + }) + const out = await r.materialize([{ type: "file", mime: "image/png", filename: "safe.png", url }]) + expect(JSON.stringify({ out, logs })).not.toContain(secret) + expect(JSON.stringify({ out, logs })).not.toContain(url) + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("materialize after dispose fails fetchable parts closed without exposing the URL", async () => { + const root = await tmpRoot() + try { + const r = RemoteAttachments.create({ sessionID: SessionID.make("ses_d"), tmpRoot: root, log: nolog }) + await r.dispose() + const out = await r.materialize([ + { type: "file", mime: "image/png", filename: "x.png", url: "https://acct.r2.cloudflarestorage.com/x.png" }, + ]) + expect(out).toEqual([ + { type: "text", text: "attachment x.png could not be retrieved: attachment session is closed" }, + ]) + expect(JSON.stringify(out)).not.toContain("cloudflarestorage.com") + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) + + test("dispose waits for an already-started materialize before removing scratch", async () => { + const root = await tmpRoot() + try { + let ready!: () => void + const started = new Promise((resolve) => { + ready = resolve + }) + let release!: (response: Response) => void + const response = new Promise((resolve) => { + release = resolve + }) + const r = RemoteAttachments.create({ + sessionID: SessionID.make("ses_concurrent"), + tmpRoot: root, + fetch: () => { + ready() + return response + }, + log: nolog, + }) + const materialize = r.materialize([ + { + type: "file", + mime: "application/octet-stream", + filename: "blob.bin", + url: "https://acct.r2.cloudflarestorage.com/blob.bin", + }, + ]) + await started + const dispose = r.dispose() + release(okResponse(new Uint8Array([1]))) + await materialize + await dispose + await expect(fs.stat(scratch(root, "ses_concurrent"))).rejects.toMatchObject({ code: "ENOENT" }) + } finally { + await fs.rm(root, { recursive: true, force: true }) + } + }) +}) diff --git a/packages/opencode/test/kilocode/sessions/remote-protocol.test.ts b/packages/opencode/test/kilocode/sessions/remote-protocol.test.ts index f2c8adfa5c..6de374028d 100644 --- a/packages/opencode/test/kilocode/sessions/remote-protocol.test.ts +++ b/packages/opencode/test/kilocode/sessions/remote-protocol.test.ts @@ -265,4 +265,60 @@ describe("RemoteProtocol", () => { expect(result.data.type).toBe("heartbeat_ack") } }) + + test("heartbeat without capabilities parses", () => { + const result = RemoteProtocol.Heartbeat.safeParse({ + type: "heartbeat", + sessions: [], + }) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.capabilities).toBeUndefined() + } + }) + + test("heartbeat with capabilities.attachments parses", () => { + const result = RemoteProtocol.Heartbeat.safeParse({ + type: "heartbeat", + sessions: [], + capabilities: { attachments: true }, + }) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.capabilities?.attachments).toBe(true) + } + }) + + test("heartbeat with capabilities and no attachments key parses", () => { + const result = RemoteProtocol.Heartbeat.safeParse({ + type: "heartbeat", + sessions: [], + capabilities: {}, + }) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.capabilities?.attachments).toBeUndefined() + } + }) + + test("heartbeat rejects non-boolean capabilities.attachments", () => { + const result = RemoteProtocol.Heartbeat.safeParse({ + type: "heartbeat", + sessions: [], + capabilities: { attachments: "yes" }, + }) + expect(result.success).toBe(false) + }) + + test("outbound union accepts heartbeat with capabilities", () => { + const result = RemoteProtocol.Outbound.safeParse({ + type: "heartbeat", + sessions: [], + capabilities: { attachments: true }, + }) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.type).toBe("heartbeat") + } + }) }) diff --git a/packages/opencode/test/kilocode/sessions/remote-sender.test.ts b/packages/opencode/test/kilocode/sessions/remote-sender.test.ts index 214793c00a..de47de4f06 100644 --- a/packages/opencode/test/kilocode/sessions/remote-sender.test.ts +++ b/packages/opencode/test/kilocode/sessions/remote-sender.test.ts @@ -294,6 +294,383 @@ describe("RemoteSender", () => { await provideStarted }) + test("send_message ACKs before the attachment materializer resolves", async () => { + const { conn, sent } = fakeConn() + let resolveMaterialize!: (parts: any[]) => void + const materializeStarted = new Promise((r) => { + // signal when the materializer has been invoked + r() + }) + const materializeInvoked = Promise.withResolvers() + const sender = RemoteSender.create({ + conn, + directory: "/tmp/test", + log: nolog, + subscribe: fakeBus().subscribe, + provide: async (input: { directory: string; fn: () => R }) => input.fn(), + attachments: () => ({ + materialize: (parts: readonly any[]) => + new Promise((resolve) => { + resolveMaterialize = resolve + materializeInvoked.resolve() + // never resolves on its own — proves the ACK is sent first + }), + dispose: async () => {}, + }), + }) + + sender.handle({ + type: "command", + id: "req_attach_ack", + command: "send_message", + data: { + sessionID: "ses_attach", + parts: [{ type: "file", mime: "image/png", filename: "a.png", url: "https://r2.example/a.png" }], + }, + }) + + // ACK is sent synchronously, BEFORE the materializer (or provide) completes + expect(sent).toHaveLength(1) + expect(sent[0]).toEqual({ type: "response", id: "req_attach_ack", result: {} }) + + // The materializer IS invoked after the ACK, confirming the work is queued + // but does not block the synchronous response. + await materializeInvoked.promise + // Resolve to let the trailing microtask settle. + resolveMaterialize([]) + await Promise.resolve() + await materializeStarted + }) + + test("does not create attachments when delayed send resumes after dispose", async () => { + const { conn } = fakeConn() + const bus = fakeBus() + const entered = Promise.withResolvers() + const release = Promise.withResolvers() + const finished = Promise.withResolvers() + let factories = 0 + let materialized = 0 + let subscriptions = 0 + const sender = RemoteSender.create({ + conn, + directory: "/tmp/test", + log: nolog, + subscribe: (callback) => { + subscriptions++ + return bus.subscribe(callback) + }, + session: { + get: async (id) => info(id), + children: async () => [], + }, + provide: async (input: { directory: string; fn: () => R }) => { + entered.resolve() + await release.promise + try { + return await input.fn() + } finally { + finished.resolve() + } + }, + prompt: async () => {}, + attachments: () => { + factories++ + return { + materialize: async (parts) => { + materialized++ + return parts + }, + dispose: async () => {}, + } + }, + }) + + sender.handle({ + type: "command", + id: "req_disposed_attachment", + command: "send_message", + data: { + sessionID: "ses_disposed_attachment", + parts: [{ type: "file", mime: "image/png", filename: "a.png", url: "https://example.com/a.png" }], + }, + }) + await entered.promise + sender.dispose() + release.resolve() + await finished.promise + + expect(factories).toBe(0) + expect(materialized).toBe(0) + expect(subscriptions).toBe(1) + expect(bus.count()).toBe(0) + }) + + test("does not create first attachments when delayed send resumes after session deletion", async () => { + const { conn } = fakeConn() + const bus = fakeBus() + const entered = Promise.withResolvers() + const release = Promise.withResolvers() + const finished = Promise.withResolvers() + const again = Promise.withResolvers() + const cleaned = Promise.withResolvers() + let factories = 0 + let materialized = 0 + let disposed = 0 + let subscriptions = 0 + const prompts: SessionPrompt.PromptInput["parts"][] = [] + const sender = RemoteSender.create({ + conn, + directory: "/tmp/test", + log: nolog, + subscribe: (callback) => { + subscriptions++ + return bus.subscribe(callback) + }, + session: { + get: async (id) => info(id), + children: async () => [], + }, + provide: async (input: { directory: string; fn: () => R }) => { + entered.resolve() + await release.promise + try { + return await input.fn() + } finally { + finished.resolve() + } + }, + prompt: async (input) => { + prompts.push(input.parts) + }, + attachments: () => { + factories++ + return { + materialize: async (parts) => { + materialized++ + again.resolve() + return parts + }, + dispose: async () => { + disposed++ + cleaned.resolve() + }, + } + }, + }) + sender.handle({ + type: "command", + id: "req_deleted_attachment", + command: "send_message", + data: { + sessionID: "ses_deleted_attachment", + parts: [{ type: "file", mime: "image/png", filename: "a.png", url: "https://example.com/a.png" }], + }, + }) + await entered.promise + bus.fire({ type: Session.Event.Deleted.type, properties: { sessionID: "ses_deleted_attachment" } }) + release.resolve() + await finished.promise + + expect(factories).toBe(0) + expect(materialized).toBe(0) + expect(disposed).toBe(0) + expect(subscriptions).toBe(1) + expect(bus.count()).toBe(1) + expect(prompts[0]).toEqual([ + { type: "text", text: "attachment a.png could not be retrieved: attachment session is closed" }, + ]) + expect(JSON.stringify(prompts[0])).not.toContain("example.com") + + sender.handle({ + type: "command", + id: "req_reused_attachment", + command: "send_message", + data: { + sessionID: "ses_deleted_attachment", + parts: [{ type: "file", mime: "image/png", filename: "a.png", url: "https://example.com/a.png" }], + }, + }) + await again.promise + expect(factories).toBe(1) + expect(materialized).toBe(1) + sender.dispose() + await cleaned.promise + expect(disposed).toBe(1) + }) + + test("keeps materialized scratch owned until an in-flight prompt settles after deletion", async () => { + const { conn } = fakeConn() + const bus = fakeBus() + const started = Promise.withResolvers() + const release = Promise.withResolvers() + const cleanupStarted = Promise.withResolvers() + const cleanupRelease = Promise.withResolvers() + const cleaned = Promise.withResolvers() + const repeated = Promise.withResolvers() + let factories = 0 + let materialized = 0 + let disposed = 0 + const seen: SessionPrompt.PromptInput["parts"][] = [] + const sender = RemoteSender.create({ + conn, + directory: "/tmp/test", + log: nolog, + subscribe: bus.subscribe, + session: { + get: async (id) => info(id), + children: async () => [], + }, + provide: async (input: { directory: string; fn: () => R }) => input.fn(), + prompt: async (input) => { + seen.push(input.parts) + if (seen.length > 1) { + repeated.resolve() + return + } + started.resolve() + await release.promise + }, + attachments: () => { + factories++ + return { + materialize: async () => { + materialized++ + return [{ type: "text", text: "attachment saved to /tmp/scratch/file.bin" }] + }, + dispose: async () => { + disposed++ + cleanupStarted.resolve() + await cleanupRelease.promise + cleaned.resolve() + }, + } + }, + }) + + sender.handle({ + type: "command", + id: "req_prompt_attachment", + command: "send_message", + data: { + sessionID: "ses_prompt_attachment", + parts: [ + { type: "file", mime: "application/octet-stream", filename: "a.bin", url: "https://example.com/a.bin" }, + ], + }, + }) + await started.promise + bus.fire({ type: Session.Event.Deleted.type, properties: { sessionID: "ses_prompt_attachment" } }) + + expect(seen[0]).toEqual([{ type: "text", text: "attachment saved to /tmp/scratch/file.bin" }]) + expect(disposed).toBe(0) + release.resolve() + await cleanupStarted.promise + + sender.handle({ + type: "command", + id: "req_prompt_attachment_reused", + command: "send_message", + data: { + sessionID: "ses_prompt_attachment", + parts: [ + { type: "file", mime: "application/octet-stream", filename: "a.bin", url: "https://example.com/a.bin" }, + ], + }, + }) + await repeated.promise + expect(factories).toBe(1) + expect(materialized).toBe(1) + expect(disposed).toBe(1) + expect(seen[1]).toEqual([ + { type: "text", text: "attachment a.bin could not be retrieved: attachment session is closed" }, + ]) + expect(JSON.stringify(seen[1])).not.toContain("example.com") + + cleanupRelease.resolve() + await cleaned.promise + expect(disposed).toBe(1) + sender.dispose() + }) + + test("blocks a new attachment generation while idle-cache deletion cleanup runs", async () => { + const { conn } = fakeConn() + const bus = fakeBus() + const first = Promise.withResolvers() + const cleanupStarted = Promise.withResolvers() + const cleanupRelease = Promise.withResolvers() + const cleaned = Promise.withResolvers() + const repeated = Promise.withResolvers() + let runs = 0 + let factories = 0 + let materialized = 0 + const seen: SessionPrompt.PromptInput["parts"][] = [] + const sender = RemoteSender.create({ + conn, + directory: "/tmp/test", + log: nolog, + subscribe: bus.subscribe, + session: { + get: async (id) => info(id), + children: async () => [], + }, + provide: async (input: { directory: string; fn: () => R }) => { + try { + return await input.fn() + } finally { + runs++ + if (runs === 1) first.resolve() + } + }, + prompt: async (input) => { + seen.push(input.parts) + if (seen.length === 2) repeated.resolve() + }, + attachments: () => { + factories++ + return { + materialize: async () => { + materialized++ + return [{ type: "text", text: "attachment saved to /tmp/scratch/file.bin" }] + }, + dispose: async () => { + cleanupStarted.resolve() + await cleanupRelease.promise + cleaned.resolve() + }, + } + }, + }) + const send = (id: string) => + sender.handle({ + type: "command", + id, + command: "send_message", + data: { + sessionID: "ses_idle_cleanup", + parts: [ + { type: "file", mime: "application/octet-stream", filename: "a.bin", url: "https://example.com/a.bin" }, + ], + }, + }) + + send("req_idle_first") + await first.promise + bus.fire({ type: Session.Event.Deleted.type, properties: { sessionID: "ses_idle_cleanup" } }) + await cleanupStarted.promise + send("req_idle_repeated") + await repeated.promise + + expect(factories).toBe(1) + expect(materialized).toBe(1) + expect(seen[1]).toEqual([ + { type: "text", text: "attachment a.bin could not be retrieved: attachment session is closed" }, + ]) + expect(JSON.stringify(seen[1])).not.toContain("example.com") + cleanupRelease.resolve() + await cleaned.promise + sender.dispose() + }) + test("send_message keeps client toggles persistent and terminal restriction ephemeral", async () => { const { conn } = fakeConn() const calls: SessionPrompt.PromptInput[] = [] diff --git a/packages/opencode/test/kilocode/sessions/remote-ws.test.ts b/packages/opencode/test/kilocode/sessions/remote-ws.test.ts index f114abace6..c8d2951b54 100644 --- a/packages/opencode/test/kilocode/sessions/remote-ws.test.ts +++ b/packages/opencode/test/kilocode/sessions/remote-ws.test.ts @@ -235,6 +235,26 @@ describe("RemoteWS", () => { expect(parsed.sessions).toEqual([{ id: "s1", status: "active", title: "Test" }]) }) + test("heartbeat advertises capabilities.attachments = true", async () => { + server = createServer() + const connecting = server.waitForConnect() + const msg = server.waitForMessage() + + conn = RemoteWS.connect({ + url: server.url, + getToken: async () => "tok", + getSessions: async () => ({ sessions: [] }), + log: nolog(), + heartbeat: 100, + }) + + await connecting + await settled() + const raw = await msg + const parsed = JSON.parse(raw) + expect(parsed.capabilities).toEqual({ attachments: true }) + }) + test("serializes concurrent heartbeat snapshots", async () => { server = createServer() const connecting = server.waitForConnect()