diff --git a/.changeset/read-notebook-cells.md b/.changeset/read-notebook-cells.md new file mode 100644 index 00000000000..00df018eb2a --- /dev/null +++ b/.changeset/read-notebook-cells.md @@ -0,0 +1,5 @@ +--- +"@kilocode/cli": patch +--- + +Read Jupyter notebooks as ordered markdown and code cell content instead of raw notebook payloads. diff --git a/packages/opencode/src/kilocode/tool/notebook.ts b/packages/opencode/src/kilocode/tool/notebook.ts new file mode 100644 index 00000000000..51d1006f7a2 --- /dev/null +++ b/packages/opencode/src/kilocode/tool/notebook.ts @@ -0,0 +1,46 @@ +import * as path from "path" +import { Readable } from "stream" +import * as Encoding from "../encoding" + +type ObjectValue = Record + +const object = (value: unknown): value is ObjectValue => typeof value === "object" && value !== null && !Array.isArray(value) + +const parse = (text: string): unknown => { + try { + return JSON.parse(text) + } catch { + return undefined + } +} + +const source = (value: unknown): string | undefined => { + if (typeof value === "string") return value + if (!Array.isArray(value) || !value.every((line) => typeof line === "string")) return undefined + return value.join("") +} + +const render = (kind: "markdown" | "code", text: string) => { + const body = text.endsWith("\n") ? text : `${text}\n` + return `<${kind}_cell>\n${body}` +} + +export async function open(filepath: string): Promise { + if (path.extname(filepath).toLowerCase() !== ".ipynb") return undefined + + const raw = (await Encoding.read(filepath)).text + const data = parse(raw) + if (!object(data) || !Array.isArray(data.cells)) return Readable.from([raw]) + + const cells: string[] = [] + for (const cell of data.cells) { + if (!object(cell)) continue + if (cell.cell_type !== "markdown" && cell.cell_type !== "code") continue + + const text = source(cell.source) + if (text === undefined) continue + cells.push(render(cell.cell_type, text)) + } + + return Readable.from([cells.length ? cells.join("\n\n") : "(Notebook contains no markdown or code cell content.)"]) +} diff --git a/packages/opencode/src/tool/read.ts b/packages/opencode/src/tool/read.ts index 551e3744553..88c9f216e63 100644 --- a/packages/opencode/src/tool/read.ts +++ b/packages/opencode/src/tool/read.ts @@ -15,6 +15,7 @@ import { isPdfAttachment, sniffAttachmentMime } from "@/util/media" // kilocode_change start import * as Encoding from "../kilocode/encoding" import * as TextStream from "../kilocode/text-stream" +import * as Notebook from "../kilocode/tool/notebook" // kilocode_change end const DEFAULT_READ_LIMIT = 2000 @@ -358,6 +359,8 @@ export const ReadTool = Tool.define( // routed through TextStream.withFallback so non-UTF-8 files are decoded via // iconv. The body otherwise matches upstream. export async function lines(filepath: string, opts: { limit: number; offset: number }) { + const extracted = await Notebook.open(filepath) // kilocode_change - extract readable notebook cells before paging + if (extracted) return readLines(extracted, opts) // kilocode_change return TextStream.withFallback(filepath, (stream) => readLines(stream, opts)) } diff --git a/packages/opencode/test/kilocode/read-notebook.test.ts b/packages/opencode/test/kilocode/read-notebook.test.ts new file mode 100644 index 00000000000..e4a6fa2da1e --- /dev/null +++ b/packages/opencode/test/kilocode/read-notebook.test.ts @@ -0,0 +1,196 @@ +import { describe, expect } from "bun:test" +import { Effect, Layer } from "effect" +import path from "path" +import { Agent } from "../../src/agent/agent" +import * as CrossSpawnSpawner from "@opencode-ai/core/cross-spawn-spawner" +import { AppFileSystem } from "@opencode-ai/core/filesystem" +import { LSP } from "../../src/lsp/lsp" +import { Instruction } from "../../src/session/instruction" +import { MessageID, SessionID } from "../../src/session/schema" +import { ReadTool } from "../../src/tool/read" +import * as Tool from "../../src/tool/tool" +import { Truncate } from "../../src/tool/truncate" +import { provideInstance, tmpdirScoped } from "../fixture/fixture" +import { testEffect } from "../lib/effect" + +const ctx = { + sessionID: SessionID.make("ses_test-notebook"), + messageID: MessageID.make(""), + callID: "", + agent: "code", + abort: AbortSignal.any([]), + messages: [], + metadata: () => Effect.void, + ask: () => Effect.void, +} + +const it = testEffect( + Layer.mergeAll( + Agent.defaultLayer, + AppFileSystem.defaultLayer, + CrossSpawnSpawner.defaultLayer, + Instruction.defaultLayer, + LSP.defaultLayer, + Truncate.defaultLayer, + ), +) + +const run = Effect.fn("NotebookReadTest.run")(function* (dir: string, args: Tool.InferParameters) { + return yield* provideInstance(dir)( + Effect.gen(function* () { + const info = yield* ReadTool + const tool = yield* Tool.init(info) + return yield* tool.execute(args, ctx) + }), + ) +}) + +const put = Effect.fn("NotebookReadTest.put")(function* (filepath: string, content: string | Uint8Array) { + const fs = yield* AppFileSystem.Service + yield* fs.writeWithDirs(filepath, content) +}) + +const notebook = JSON.stringify({ + metadata: { secret: "ignore-notebook-metadata" }, + cells: [ + { + cell_type: "markdown", + metadata: { private: "ignore-cell-metadata" }, + source: ["# Analysis\n", "Useful introduction"], + }, + { + cell_type: "raw", + source: ["ignore raw cell"], + }, + { + cell_type: "code", + execution_count: 7, + metadata: {}, + source: ["value = 42\n", "print(value)"], + outputs: [{ output_type: "stream", text: ["ignore-output-payload"] }], + }, + ], +}) + +describe("kilocode notebook reads", () => { + it.live("extracts markdown and code cells without notebook payloads", () => + Effect.gen(function* () { + const dir = yield* tmpdirScoped() + const filepath = path.join(dir, "analysis.ipynb") + yield* put(filepath, notebook) + + const result = yield* run(dir, { filePath: filepath }) + + expect(result.output).toContain("") + expect(result.output).toContain("# Analysis") + expect(result.output).toContain("") + expect(result.output).toContain("value = 42") + expect(result.output.indexOf("# Analysis")).toBeLessThan(result.output.indexOf("value = 42")) + expect(result.output).not.toContain("ignore-output-payload") + expect(result.output).not.toContain("ignore-notebook-metadata") + expect(result.output).not.toContain("ignore-cell-metadata") + expect(result.output).not.toContain("ignore raw cell") + }), + ) + + it.live("skips invalid cells without exposing raw notebook payloads", () => + Effect.gen(function* () { + const dir = yield* tmpdirScoped() + const filepath = path.join(dir, "partial.ipynb") + const content = JSON.stringify({ + cells: [ + null, + { cell_type: "code", source: null, outputs: ["INVALID_OUTPUT_SHOULD_NOT_APPEAR"] }, + { cell_type: "markdown", source: ["Readable cell"], metadata: { marker: "CELL_METADATA_SHOULD_NOT_APPEAR" } }, + ], + }) + yield* put(filepath, content) + + const result = yield* run(dir, { filePath: filepath }) + + expect(result.output).toContain("Readable cell") + expect(result.output).not.toContain("INVALID_OUTPUT_SHOULD_NOT_APPEAR") + expect(result.output).not.toContain("CELL_METADATA_SHOULD_NOT_APPEAR") + }), + ) + + it.live("reports valid notebooks with no readable cells without exposing payloads", () => + Effect.gen(function* () { + const dir = yield* tmpdirScoped() + const filepath = path.join(dir, "empty-content.ipynb") + const content = JSON.stringify({ + metadata: { marker: "NOTEBOOK_METADATA_SHOULD_NOT_APPEAR" }, + cells: [ + null, + { cell_type: "raw", source: ["RAW_CONTENT_SHOULD_NOT_APPEAR"] }, + { cell_type: "code", source: null, outputs: ["INVALID_OUTPUT_SHOULD_NOT_APPEAR"] }, + ], + }) + yield* put(filepath, content) + + const result = yield* run(dir, { filePath: filepath }) + + expect(result.output).toContain("Notebook contains no markdown or code cell content") + expect(result.output).not.toContain("NOTEBOOK_METADATA_SHOULD_NOT_APPEAR") + expect(result.output).not.toContain("RAW_CONTENT_SHOULD_NOT_APPEAR") + expect(result.output).not.toContain("INVALID_OUTPUT_SHOULD_NOT_APPEAR") + }), + ) + + it.live("applies read pagination to extracted cell text", () => + Effect.gen(function* () { + const dir = yield* tmpdirScoped() + const filepath = path.join(dir, "paged.ipynb") + yield* put(filepath, notebook) + + const result = yield* run(dir, { filePath: filepath, offset: 2, limit: 2 }) + + expect(result.output).toContain("2: # Analysis") + expect(result.output).toContain("3: Useful introduction") + expect(result.output).not.toContain("value = 42") + expect(result.metadata.preview).toBe("# Analysis\nUseful introduction") + expect(result.metadata.truncated).toBe(true) + }), + ) + + it.live("falls back to raw text for malformed notebooks", () => + Effect.gen(function* () { + const dir = yield* tmpdirScoped() + const filepath = path.join(dir, "broken.ipynb") + const content = '{"cells":[{"cell_type":"markdown","source":["unfinished"]}' + yield* put(filepath, content) + + const result = yield* run(dir, { filePath: filepath }) + + expect(result.output).toContain(content) + expect(result.output).not.toContain("") + }), + ) + + it.live("keeps ordinary text reads unchanged", () => + Effect.gen(function* () { + const dir = yield* tmpdirScoped() + const filepath = path.join(dir, "notes.txt") + yield* put(filepath, "plain text") + + const result = yield* run(dir, { filePath: filepath }) + + expect(result.output).toContain("1: plain text") + expect(result.output).not.toContain("") + }), + ) + + it.live("keeps PDF files as native attachments", () => + Effect.gen(function* () { + const dir = yield* tmpdirScoped() + const filepath = path.join(dir, "document.pdf") + yield* put(filepath, "%PDF-1.4\nminimal content") + + const result = yield* run(dir, { filePath: filepath }) + + expect(result.output).toBe("PDF read successfully") + expect(result.attachments?.[0].mime).toBe("application/pdf") + expect(result.metadata.truncated).toBe(false) + }), + ) +})