Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/read-notebook-cells.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@kilocode/cli": patch
---

Read Jupyter notebooks as ordered markdown and code cell content instead of raw notebook payloads.
46 changes: 46 additions & 0 deletions packages/opencode/src/kilocode/tool/notebook.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
import * as path from "path"
import { Readable } from "stream"
import * as Encoding from "../encoding"

type ObjectValue = Record<string, unknown>

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}</${kind}_cell>`
}

export async function open(filepath: string): Promise<Readable | undefined> {
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.)"])
}
3 changes: 3 additions & 0 deletions packages/opencode/src/tool/read.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Comment thread
marius-kilocode marked this conversation as resolved.
if (extracted) return readLines(extracted, opts) // kilocode_change
return TextStream.withFallback(filepath, (stream) => readLines(stream, opts))
}

Expand Down
196 changes: 196 additions & 0 deletions packages/opencode/test/kilocode/read-notebook.test.ts
Original file line number Diff line number Diff line change
@@ -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<typeof ReadTool>) {
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("<markdown_cell>")
expect(result.output).toContain("# Analysis")
expect(result.output).toContain("<code_cell>")
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("<markdown_cell>")
}),
)

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("<markdown_cell>")
}),
)

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)
}),
)
})
Loading