diff --git a/packages/opencode/src/kilocode/plugins/model-usage.ts b/packages/opencode/src/kilocode/plugins/model-usage.ts index 99f58031f0..5ce052d3ee 100644 --- a/packages/opencode/src/kilocode/plugins/model-usage.ts +++ b/packages/opencode/src/kilocode/plugins/model-usage.ts @@ -1,4 +1,5 @@ -import type { KilocodeSessionModelUsageResponse } from "@kilocode/sdk/v2" +import type { KilocodeSessionModelUsageResponse, Session } from "@kilocode/sdk/v2" +import { KiloRoutedModel } from "@/kilocode/session/routed-model" export type SessionModelUsage = KilocodeSessionModelUsageResponse export type UsageResult = { sessionID: string; data?: SessionModelUsage } @@ -12,6 +13,28 @@ export function failed(result: UsageResult | undefined, sessionID: string) { return result?.sessionID === sessionID && !result.data } +export function member(input: { + root: string + sessionID: string + get: (sessionID: string) => Session | undefined + info?: Session +}) { + const seen = new Set() + const visit = (sessionID: string, info?: Session): boolean => { + if (sessionID === input.root) return true + if (seen.has(sessionID)) return false + seen.add(sessionID) + const session = info ?? input.get(sessionID) + if (!session?.parentID) return false + return visit(session.parentID) + } + return visit(input.sessionID, input.info) +} + +export function label(model: SessionModelUsage["models"][number]) { + return KiloRoutedModel.displayName(KiloRoutedModel.display(model.modelID)) +} + const count = new Intl.NumberFormat("en-US") const currency = new Intl.NumberFormat("en-US", { style: "currency", diff --git a/packages/opencode/src/kilocode/plugins/sidebar-usage.tsx b/packages/opencode/src/kilocode/plugins/sidebar-usage.tsx index 8606fe8370..741368da2f 100644 --- a/packages/opencode/src/kilocode/plugins/sidebar-usage.tsx +++ b/packages/opencode/src/kilocode/plugins/sidebar-usage.tsx @@ -7,18 +7,14 @@ import { formatCost, formatCount, formatRate, + label, + member, select, - type SessionModelUsage, type UsageResult, } from "@/kilocode/plugins/model-usage" const id = "internal:kilo-sidebar-usage" -function identity(model: SessionModelUsage["models"][number]) { - if (model.providerID === "kilo" && model.modelID.includes("/")) return model.modelID - return `${model.providerID}/${model.modelID}` -} - function View(props: { api: TuiPluginApi; session_id: string }) { const theme = () => props.api.theme.current const local = useLocal() @@ -47,14 +43,24 @@ function View(props: { api: TuiPluginApi; session_id: string }) { onMount(() => { const refresh = () => void refetch() + const related = (sessionID: string, info?: ReturnType) => + member({ root: props.session_id, sessionID, info, get: props.api.state.session.get }) const offs = [ props.api.event.on("message.part.updated", (event) => { - if (event.properties.part.type === "step-finish") refresh() + if (event.properties.part.type === "step-finish" && related(event.properties.sessionID)) refresh() + }), + props.api.event.on("message.part.removed", (event) => { + if (related(event.properties.sessionID)) refresh() + }), + props.api.event.on("message.removed", (event) => { + if (related(event.properties.sessionID)) refresh() + }), + props.api.event.on("session.created", (event) => { + if (related(event.properties.sessionID, event.properties.info)) refresh() + }), + props.api.event.on("session.deleted", (event) => { + if (related(event.properties.sessionID, event.properties.info)) refresh() }), - props.api.event.on("message.part.removed", refresh), - props.api.event.on("message.removed", refresh), - props.api.event.on("session.created", refresh), - props.api.event.on("session.deleted", refresh), props.api.event.on("server.connected", refresh), ] onCleanup(() => { @@ -108,7 +114,7 @@ function View(props: { api: TuiPluginApi; session_id: string }) { {(model) => ( - {identity(model)} + {label(model)} Steps {formatCount(model.steps)} | Cost {formatCost(model.cost)} diff --git a/packages/opencode/src/kilocode/session/model-usage.ts b/packages/opencode/src/kilocode/session/model-usage.ts index 05ab22bdf1..47467be1d7 100644 --- a/packages/opencode/src/kilocode/session/model-usage.ts +++ b/packages/opencode/src/kilocode/session/model-usage.ts @@ -1,7 +1,7 @@ import { NonNegativeInt } from "@opencode-ai/core/schema" import { Effect, Schema } from "effect" -import { InstanceState } from "@/effect/instance-state" import { ModelID, ProviderID } from "@/provider/schema" +import { ProjectID } from "@/project/schema" import { SessionID } from "@/session/schema" import { Database } from "@/storage/db" @@ -37,6 +37,10 @@ export namespace ModelUsage { type Info = typeof Info.Type + type Anchor = { + projectID: ProjectID + } + type Ancestor = { id: SessionID parentID: SessionID | null @@ -54,6 +58,8 @@ export namespace ModelUsage { write: number } + const ANCHOR_SQL = "SELECT project_id AS projectID FROM session WHERE id = ?" + const ANCESTORS_SQL = ` WITH RECURSIVE ancestor(id, parent_id) AS ( SELECT id, parent_id @@ -125,16 +131,16 @@ export namespace ModelUsage { }) export const get = Effect.fn("ModelUsage.get")(function* (sessionID: SessionID) { - const ctx = yield* InstanceState.context return yield* Effect.sync(() => { const db = Database.Client().$client - const args = [sessionID, ctx.project.id, ctx.project.id] as const - const ancestors = db.prepare(ANCESTORS_SQL).all(...args) - if (ancestors.length === 0) return undefined + const anchor = db.prepare(ANCHOR_SQL).get(sessionID) + if (!anchor) return undefined + const args = [sessionID, anchor.projectID, anchor.projectID] as const + const ancestors = db.prepare(ANCESTORS_SQL).all(...args) const ids = new Set(ancestors.map((item) => item.id)) const rootID = ancestors.find((item) => !item.parentID || !ids.has(item.parentID))?.id ?? sessionID - const familyArgs = [rootID, ctx.project.id, ctx.project.id] as const + const familyArgs = [rootID, anchor.projectID, anchor.projectID] as const const rows = db.prepare(USAGE_SQL).all(...familyArgs) const totals = empty() const models = rows.map((row): Model => { diff --git a/packages/opencode/test/kilocode/cli/cmd/tui/model-usage.test.ts b/packages/opencode/test/kilocode/cli/cmd/tui/model-usage.test.ts index babaf4d40a..8bc593d9b6 100644 --- a/packages/opencode/test/kilocode/cli/cmd/tui/model-usage.test.ts +++ b/packages/opencode/test/kilocode/cli/cmd/tui/model-usage.test.ts @@ -1,5 +1,18 @@ import { describe, expect, test } from "bun:test" -import { failed, formatRate, select, type SessionModelUsage } from "@/kilocode/plugins/model-usage" +import type { Session } from "@kilocode/sdk/v2" +import { failed, formatRate, label, member, select, type SessionModelUsage } from "@/kilocode/plugins/model-usage" + +const session = (id: string, parentID?: string) => + ({ + id, + parentID, + slug: id, + projectID: "project", + directory: "/project", + title: id, + version: "1", + time: { created: 0, updated: 0 }, + }) satisfies Session const data = { totals: { @@ -11,11 +24,34 @@ const data = { } satisfies SessionModelUsage describe("TUI model usage", () => { - test("rejects stale session results and computes the cache rate", () => { + test("filters session results and formats usage labels", () => { + const root = session("ses_root") + const child = session("ses_child", root.id) + const sessions = new Map([root, child].map((item) => [item.id, item])) + expect(select({ sessionID: "ses_old", data }, "ses_current")).toBeUndefined() expect(failed({ sessionID: "ses_old" }, "ses_current")).toBeFalse() expect(select({ sessionID: "ses_current", data }, "ses_current")).toBe(data) expect(failed({ sessionID: "ses_current" }, "ses_current")).toBeTrue() + expect(member({ root: root.id, sessionID: child.id, get: (id) => sessions.get(id) })).toBeTrue() + expect( + member({ + root: root.id, + sessionID: "ses_new", + info: session("ses_new", child.id), + get: (id) => sessions.get(id), + }), + ).toBeTrue() + expect(member({ root: root.id, sessionID: "ses_other", get: () => undefined })).toBeFalse() + expect( + label({ + providerID: "kilo", + modelID: "openai/gpt-5.5-20260423", + steps: 1, + cost: 0, + tokens: data.totals.tokens, + }), + ).toBe("gpt-5.5") expect(formatRate({ input: 100, output: 0, reasoning: 0, cache: { read: 300, write: 100 } })).toBe("60.0%") }) }) diff --git a/packages/opencode/test/kilocode/session-model-usage.test.ts b/packages/opencode/test/kilocode/session-model-usage.test.ts index 7a3b6a855f..e47f9d9ac5 100644 --- a/packages/opencode/test/kilocode/session-model-usage.test.ts +++ b/packages/opencode/test/kilocode/session-model-usage.test.ts @@ -1,10 +1,15 @@ import { describe, expect } from "bun:test" import { Effect } from "effect" import { ModelUsage } from "@/kilocode/session/model-usage" +import { ProjectTable } from "@/project/project.sql" +import { ProjectID } from "@/project/schema" import { MessageV2 } from "@/session/message-v2" import { Session } from "@/session/session" +import { SessionTable } from "@/session/session.sql" import { MessageID, PartID, SessionID } from "@/session/schema" import { ModelID, ProviderID } from "@/provider/schema" +import { Database, eq } from "@/storage/db" +import { TestInstance } from "../fixture/fixture" import { testEffect } from "../lib/effect" const it = testEffect(Session.defaultLayer) @@ -64,6 +69,7 @@ describe("session model usage", () => { it.instance("aggregates direct step usage by model across the top-level session tree", () => Effect.gen(function* () { const sessions = yield* Session.Service + const test = yield* TestInstance const root = yield* sessions.create({ title: "root" }) const child = yield* sessions.create({ title: "child", parentID: root.id }) const sibling = yield* sessions.create({ title: "sibling", parentID: root.id }) @@ -105,6 +111,23 @@ describe("session model usage", () => { tokens: { input: 9_000, output: 9_000, reasoning: 9_000, cache: { read: 9_000, write: 9_000 } }, }) + const project = ProjectID.make("legacy-project") + Database.use((db) => { + db.insert(ProjectTable) + .values({ + id: project, + worktree: test.directory, + vcs: "git", + time_created: Date.now(), + time_updated: Date.now(), + sandboxes: [], + }) + .run() + for (const session of [root, child, sibling]) { + db.update(SessionTable).set({ project_id: project }).where(eq(SessionTable.id, session.id)).run() + } + }) + expect(yield* ModelUsage.get(child.id)).toEqual({ totals: { steps: 3,