opencode/packages/core/script/benchmark-snapshot-session.ts

/**
 * Drives the real SessionStep.attempt through a scripted agent trace against a real
 * snapshot store, then repeats it with snapshots disabled. The difference is the
 * wall-clock cost snapshots add to the session.
 *
 *   bun run script/benchmark-snapshot-session.ts [--fixture opencode] [--ttft 400] [--rounds 2]
 *
 * Uses the fixtures generated by benchmark-snapshot.ts.
 */
import { $ } from "bun"
import fs from "fs/promises"
import os from "os"
import path from "path"
import { LanguageModel, LLM, LLMEvent } from "@opencode/ai"
import { OpenAIChat } from "@opencode/ai/protocols/openai-chat"
import { TestLLM } from "@opencode/ai/testing"
import { Global } from "@opencode/util/global"
import { AppProcess } from "@opencode/util/process"
import { LayerNode } from "@opencode/util/effect/layer-node"
import { Money } from "@opencode/schema/money"
import { Effect, Layer, Logger, Stream } from "effect"
import { Agent } from "../src/agent"
import { Bus } from "../src/bus"
import { Database } from "../src/database/database"
import { AppNodeBuilder } from "../src/effect/app-node-builder"
import { Location } from "../src/location"
import { Project } from "../src/project"
import { ProjectTable } from "../src/project/sql"
import { AbsolutePath } from "../src/schema"
import { Session } from "../src/session"
import { SessionMessage } from "../src/session/message"
import { SessionProjector } from "../src/session/projector"
import { SessionRunnerModel } from "../src/session/runner/model"
import { SessionStep } from "../src/session/runner/step"
import { SessionTable } from "../src/session/sql"
import { Snapshot } from "../src/snapshot"
import { ToolOutput } from "../src/tool-output"

const args = process.argv.slice(2)
const flag = (name: string) => {
  const index = args.indexOf(`--${name}`)
  return index === -1 ? undefined : args[index + 1]
}
const fixture = flag("fixture") ?? "opencode"
const ttft = Number(flag("ttft") ?? 400)
const rounds = Number(flag("rounds") ?? 2)
const root = path.join(process.env.SNAPSHOT_BENCH_ROOT ?? os.tmpdir(), "opencode-snapshot-bench")
const directory = path.join(root, "fixtures", fixture)

// A read-heavy coding loop: explore, edit, run checks, answer.
const trace = [
  "read",
  "grep",
  "read",
  "glob",
  "read",
  "edit",
  "shell",
  "read",
  "edit",
  "shell",
  "grep",
  "read",
  "edit",
  "shell",
  "text",
] as const

let spawns = 0
const countingProcess = Layer.effect(
  AppProcess.Service,
  Effect.gen(function* () {
    const real = yield* AppProcess.Service
    return AppProcess.Service.of({
      ...real,
      run: (command, options) => {
        spawns++
        return real.run(command, options)
      },
    })
  }),
).pipe(Layer.provide(AppNodeBuilder.build(AppProcess.node)))

const snapshotLayer = (data: string) =>
  AppNodeBuilder.build(Snapshot.node, [
    Location.node.replace(Location.boundNode(Location.Ref.make({ directory: AbsolutePath.make(directory) }))),
    Global.node.replace(Global.layerWith({ data, config: path.join(data, "config") })),
    AppProcess.node.replace(countingProcess),
  ])

const files = (await $`git ls-files -z -- '*.ts'`.cwd(directory).text()).split("\0").filter(Boolean).slice(0, 50)

const model = SessionRunnerModel.resolved(
  LanguageModel.make({ id: "bench-model", provider: "test", route: OpenAIChat.route }),
  {
    capabilities: { tools: true, input: ["text"], output: ["text"] },
    limit: { context: 100_000, output: 1_000 },
    cost: [
      {
        input: Money.USDPerMillionTokens.make(1),
        output: Money.USDPerMillionTokens.make(2),
        cache: { read: Money.USDPerMillionTokens.make(0.1), write: Money.USDPerMillionTokens.make(0.5) },
      },
    ],
  },
)

const delayed = (response: ReturnType<typeof TestLLM.stop>) =>
  Stream.unwrap(Effect.sleep(ttft).pipe(Effect.as(Stream.fromIterable(response as Iterable<LLMEvent>))))

const run = (snapshots: Layer.Layer<Snapshot.Service>) =>
  Effect.gen(function* () {
    const db = (yield* Database.Service).db
    const llm = yield* TestLLM.Test
    const sessionID = Session.ID.create()
    yield* db
      .insert(ProjectTable)
      .values({ id: Project.ID.global, worktree: AbsolutePath.make(directory), sandboxes: [] })
      .onConflictDoNothing()
      .run()
    yield* db
      .insert(SessionTable)
      .values({ id: sessionID, project_id: Project.ID.global, slug: "bench", directory, version: "bench" })
      .run()
    const steps = yield* SessionStep.make
    const snapshot = yield* Snapshot.Service
    // Warm the snapshot store so the first step does not pay repository creation.
    yield* snapshot.capture()
    const before = spawns
    const start = performance.now()
    let edits = 0
    for (let round = 0; round < rounds; round++)
      for (const kind of trace) {
        const call = `call-${round}-${edits}-${kind}-${Math.random().toString(36).slice(2)}`
        yield* llm.push(
          delayed(
            kind === "text"
              ? TestLLM.text("Done.", `text-${call}`)
              : TestLLM.tool(call, kind, { path: files[edits % files.length] }),
          ),
        )
        yield* steps.attempt({
          isLocationClosed: () => false,
          sessionID,
          assistantMessageID: SessionMessage.ID.create(),
          agent: Agent.defaultID,
          model,
          prepared: {
            retry: () => Effect.void,
            request: LLM.request({ model: model.model, prompt: "bench" }),
            options: {},
            executeTool: (input) =>
              Effect.gen(function* () {
                if (input.call.name === "edit")
                  yield* Effect.promise(() =>
                    Bun.write(path.join(directory, files[edits++ % files.length]!), `// bench ${edits}\n`),
                  )
                return { content: [{ type: "text" as const, text: "ok" }] }
              }),
          },
          retry: () => Effect.succeed({ retry: false as const }),
          recoverContinuation: false,
          recoverOverflow: Effect.succeed(false),
        })
      }
    return { ms: performance.now() - start, spawns: spawns - before }
  }).pipe(
    Effect.provide(snapshots),
    Effect.provide(
      Layer.merge(
        AppNodeBuilder.build(LayerNode.group([Database.node, Bus.node, SessionProjector.node, ToolOutput.node]), [
          Bus.node.replace(Bus.configured({ persist: true })),
        ]),
        TestLLM.testLayer(),
      ),
    ),
    Effect.provide(Logger.layer([])),
  )

await $`git -c core.fsmonitor=false reset -q --hard`.cwd(directory).quiet()
const data = await fs.mkdtemp(path.join(root, "session-data-"))
const noop = await Effect.runPromise(run(Snapshot.noopLayer))
await $`git -c core.fsmonitor=false reset -q --hard`.cwd(directory).quiet()
const real = await Effect.runPromise(run(snapshotLayer(data)))
await $`git -c core.fsmonitor=false reset -q --hard`.cwd(directory).quiet()
await fs.rm(data, { recursive: true, force: true })

const count = trace.length * rounds
console.log(`fixture ${fixture}, ${count} steps, simulated time to first event ${ttft} ms`)
console.log(`no snapshots   ${noop.ms.toFixed(0).padStart(7)} ms`)
console.log(`snapshots      ${real.ms.toFixed(0).padStart(7)} ms   git ${real.spawns}`)
console.log(
  `overhead       ${(real.ms - noop.ms).toFixed(0).padStart(7)} ms   ${((real.ms - noop.ms) / count).toFixed(1)} ms/step   ${(real.spawns / count).toFixed(1)} git/step`,
)