diff --git a/packages/core/schema.json b/packages/core/schema.json index d0eeeebd5c41..3e5f5c9431b3 100644 --- a/packages/core/schema.json +++ b/packages/core/schema.json @@ -1,8 +1,8 @@ { "version": "7", "dialect": "sqlite", - "id": "f14a9b18-8207-487e-a3d3-227e629ba9ad", - "prevIds": ["169a0f0f-d58f-479f-b024-fa1c7b9a09db"], + "id": "19be7275-7d21-4de7-a12b-2addde7991da", + "prevIds": ["f14a9b18-8207-487e-a3d3-227e629ba9ad"], "ddl": [ { "name": "workspace", @@ -1220,6 +1220,16 @@ "entityType": "columns", "table": "session" }, + { + "type": "real", + "notNull": false, + "autoincrement": false, + "default": null, + "generated": null, + "name": "budget", + "entityType": "columns", + "table": "session" + }, { "type": "integer", "notNull": true, diff --git a/packages/core/script/migration.ts b/packages/core/script/migration.ts index 4f383f5f8f68..4e05ce7a95d6 100644 --- a/packages/core/script/migration.ts +++ b/packages/core/script/migration.ts @@ -107,7 +107,7 @@ export default { ...config, out: ${JSON.stringify(output)} } async function generatedMigrations(directory: string) { return (await Array.fromAsync(new Bun.Glob("*/migration.sql").scan({ cwd: directory }))) - .map((file) => file.split("/")[0]) + .map((file) => file.split(/[\\/]/)[0]) .filter((name): name is string => name !== undefined) .sort() } diff --git a/packages/core/src/database/migration.gen.ts b/packages/core/src/database/migration.gen.ts index e6ea4eaa1477..a2c2b318a353 100644 --- a/packages/core/src/database/migration.gen.ts +++ b/packages/core/src/database/migration.gen.ts @@ -40,5 +40,6 @@ export const migrations = ( import("./migration/20260622142730_simplify_session_context_epoch"), import("./migration/20260622170816_reset_v2_session_state"), import("./migration/20260622202450_simplify_session_input"), + import("./migration/20260812223059_session_budget"), ]) ).map((module) => module.default) satisfies DatabaseMigration.Migration[] diff --git a/packages/core/src/database/migration/20260812223059_session_budget.ts b/packages/core/src/database/migration/20260812223059_session_budget.ts new file mode 100644 index 000000000000..b94c27a0919c --- /dev/null +++ b/packages/core/src/database/migration/20260812223059_session_budget.ts @@ -0,0 +1,11 @@ +import { Effect } from "effect" +import type { DatabaseMigration } from "../migration" + +export default { + id: "20260812223059_session_budget", + up(tx) { + return Effect.gen(function* () { + yield* tx.run(`ALTER TABLE \`session\` ADD \`budget\` real;`) + }) + }, +} satisfies DatabaseMigration.Migration diff --git a/packages/core/src/database/schema.gen.ts b/packages/core/src/database/schema.gen.ts index ed60fde6c55f..3f3c564edcbd 100644 --- a/packages/core/src/database/schema.gen.ts +++ b/packages/core/src/database/schema.gen.ts @@ -196,6 +196,7 @@ export default { \`summary_diffs\` text, \`metadata\` text, \`cost\` real DEFAULT 0 NOT NULL, + \`budget\` real, \`tokens_input\` integer DEFAULT 0 NOT NULL, \`tokens_output\` integer DEFAULT 0 NOT NULL, \`tokens_reasoning\` integer DEFAULT 0 NOT NULL, diff --git a/packages/core/src/session/info.ts b/packages/core/src/session/info.ts index 66832750fd65..b6cb3b7f12c4 100644 --- a/packages/core/src/session/info.ts +++ b/packages/core/src/session/info.ts @@ -26,6 +26,7 @@ export function fromRow(row: typeof SessionTable.$inferSelect): SessionSchema.In } : undefined, cost: row.cost, + budget: row.budget ?? undefined, tokens: { input: row.tokens_input, output: row.tokens_output, diff --git a/packages/core/src/session/projector.ts b/packages/core/src/session/projector.ts index 792067017d14..31b0f9130988 100644 --- a/packages/core/src/session/projector.ts +++ b/packages/core/src/session/projector.ts @@ -60,6 +60,7 @@ function sessionRow(info: SessionV1.SessionInfo): typeof SessionTable.$inferInse summary_diffs: info.summary?.diffs ? [...info.summary.diffs] : undefined, metadata: info.metadata, cost: info.cost ?? 0, + budget: info.budget ?? null, tokens_input: (info.tokens ?? { input: 0 }).input, tokens_output: (info.tokens ?? { output: 0 }).output, tokens_reasoning: (info.tokens ?? { reasoning: 0 }).reasoning, diff --git a/packages/core/src/session/sql.ts b/packages/core/src/session/sql.ts index 264a1d2cca0a..1dfe6d6ff23e 100644 --- a/packages/core/src/session/sql.ts +++ b/packages/core/src/session/sql.ts @@ -41,6 +41,7 @@ export const SessionTable = sqliteTable( summary_diffs: text({ mode: "json" }).$type(), metadata: text({ mode: "json" }).$type>(), cost: real().notNull().default(0), + budget: real(), tokens_input: integer().notNull().default(0), tokens_output: integer().notNull().default(0), tokens_reasoning: integer().notNull().default(0), diff --git a/packages/opencode/src/server/routes/instance/httpapi/groups/session.ts b/packages/opencode/src/server/routes/instance/httpapi/groups/session.ts index 959a303dc964..bd7703d0b719 100644 --- a/packages/opencode/src/server/routes/instance/httpapi/groups/session.ts +++ b/packages/opencode/src/server/routes/instance/httpapi/groups/session.ts @@ -50,6 +50,7 @@ export const UpdatePayload = Schema.Struct({ title: Schema.optional(Schema.String), metadata: Schema.optional(Session.Metadata), permission: Schema.optional(PermissionV1.Ruleset), + budget: Schema.optional(Schema.NullOr(Schema.Finite)), time: Schema.optional( Schema.Struct({ archived: Schema.optional(Session.ArchivedTimestamp), diff --git a/packages/opencode/src/server/routes/instance/httpapi/handlers/session.ts b/packages/opencode/src/server/routes/instance/httpapi/handlers/session.ts index 662585020a64..8bd9b81c53a8 100644 --- a/packages/opencode/src/server/routes/instance/httpapi/handlers/session.ts +++ b/packages/opencode/src/server/routes/instance/httpapi/handlers/session.ts @@ -197,6 +197,9 @@ export const sessionHandlers = HttpApiBuilder.group(InstanceHttpApi, "session", permission: Permission.merge(current.permission ?? [], ctx.payload.permission), }) } + if (ctx.payload.budget !== undefined) { + yield* session.setBudget({ sessionID: ctx.params.sessionID, budget: ctx.payload.budget ?? undefined }) + } if (ctx.payload.time?.archived !== undefined) { yield* session.setArchived({ sessionID: ctx.params.sessionID, time: ctx.payload.time.archived }) } diff --git a/packages/opencode/src/session/processor.ts b/packages/opencode/src/session/processor.ts index 20aa8a8404d8..0092c8a5a01d 100644 --- a/packages/opencode/src/session/processor.ts +++ b/packages/opencode/src/session/processor.ts @@ -72,6 +72,10 @@ interface ProcessorContext extends Input { needsCompaction: boolean currentText: SessionV1.TextPart | undefined reasoningMap: Record + baseCost: number + budget: number | undefined + turnCost: number + budgetExceeded: boolean } type StreamEvent = LLMEvent @@ -100,6 +104,7 @@ const layer = Layer.effect( // may execute tools internally before emitting start-step events, // so capturing inside the event handler can be too late. const initialSnapshot = yield* snapshot.track() + const current = yield* session.get(input.sessionID).pipe(Effect.orDie) const ctx: ProcessorContext = { assistantMessage: input.assistantMessage, sessionID: input.sessionID, @@ -111,6 +116,10 @@ const layer = Layer.effect( needsCompaction: false, currentText: undefined, reasoningMap: {}, + baseCost: current.cost ?? 0, + budget: current.budget, + turnCost: 0, + budgetExceeded: false, } let aborted = false @@ -454,6 +463,18 @@ const layer = Layer.effect( cost: usage.cost, }) yield* session.updateMessage(ctx.assistantMessage) + ctx.turnCost += usage.cost + if (ctx.budget !== undefined && ctx.baseCost + ctx.turnCost >= ctx.budget && !ctx.budgetExceeded) { + ctx.budgetExceeded = true + yield* session.updatePart({ + id: PartID.ascending(), + messageID: ctx.assistantMessage.id, + sessionID: ctx.sessionID, + type: "text", + text: "Session budget reached. Increase the budget to continue.", + time: { start: Date.now(), end: Date.now() }, + }) + } if (ctx.snapshot) { const patch = yield* snapshot.patch(ctx.snapshot) if (patch.files.length) { @@ -677,7 +698,7 @@ const layer = Layer.effect( ) if (ctx.needsCompaction) return "compact" - if (ctx.blocked || ctx.assistantMessage.error) return "stop" + if (ctx.budgetExceeded || ctx.blocked || ctx.assistantMessage.error) return "stop" return "continue" }) }) diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 22b1d7d99a2a..00b3745bf57a 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -1129,6 +1129,12 @@ const layer = Layer.effect( break } + const sessionState = yield* sessions.get(sessionID).pipe(Effect.orDie) + if (sessionState.budget !== undefined && (sessionState.cost ?? 0) >= sessionState.budget) { + yield* status.set(sessionID, { type: "idle" }) + break + } + step++ if (step === 1) yield* title({ diff --git a/packages/opencode/src/session/session.ts b/packages/opencode/src/session/session.ts index cfe034146c1b..af049f2efe3d 100644 --- a/packages/opencode/src/session/session.ts +++ b/packages/opencode/src/session/session.ts @@ -95,6 +95,7 @@ export function fromRow(row: SessionRow): Info { version: row.version, summary, cost: row.cost, + budget: row.budget ?? undefined, tokens: { input: row.tokens_input, output: row.tokens_output, @@ -137,6 +138,7 @@ export function toRow(info: Info) { summary_diffs: info.summary?.diffs, metadata: info.metadata, cost: info.cost ?? 0, + budget: info.budget ?? null, tokens_input: (info.tokens ?? EmptyTokens).input, tokens_output: (info.tokens ?? EmptyTokens).output, tokens_reasoning: (info.tokens ?? EmptyTokens).reasoning, @@ -231,6 +233,7 @@ export const Info = Schema.Struct({ parentID: optional(SessionID), summary: optional(Summary), cost: optional(Schema.Finite), + budget: optional(Schema.Finite), tokens: optional(Tokens), share: optional(Share), title: Schema.String, @@ -265,6 +268,7 @@ export const CreateInput = Schema.optional( model: Schema.optional(Model), metadata: Schema.optional(Metadata), permission: Schema.optional(PermissionV1.Ruleset), + budget: Schema.optional(Schema.Finite), workspaceID: Schema.optional(WorkspaceV2.ID), }), ) @@ -422,6 +426,7 @@ export interface Interface { model?: Schema.Schema.Type metadata?: typeof Metadata.Type permission?: PermissionV1.Ruleset + budget?: number workspaceID?: WorkspaceV2.ID }) => Effect.Effect readonly fork: (input: { sessionID: SessionID; messageID?: MessageID }) => Effect.Effect @@ -430,6 +435,7 @@ export interface Interface { readonly setTitle: (input: { sessionID: SessionID; title: string }) => Effect.Effect readonly setArchived: (input: { sessionID: SessionID; time?: number }) => Effect.Effect readonly setMetadata: (input: typeof SetMetadataInput.Type) => Effect.Effect + readonly setBudget: (input: { sessionID: SessionID; budget?: number }) => Effect.Effect readonly setAgentModel: (input: { sessionID: SessionID agent: string @@ -509,6 +515,7 @@ const layer: Layer.Layer< path?: string metadata?: typeof Metadata.Type permission?: PermissionV1.Ruleset + budget?: number }) { const ctx = yield* InstanceState.context const result: Info = { @@ -526,6 +533,7 @@ const layer: Layer.Layer< metadata: input.metadata, permission: input.permission ? [...input.permission] : undefined, cost: 0, + budget: input.budget, tokens: EmptyTokens, time: { created: Date.now(), @@ -673,6 +681,7 @@ const layer: Layer.Layer< model?: Schema.Schema.Type metadata?: typeof Metadata.Type permission?: PermissionV1.Ruleset + budget?: number workspaceID?: WorkspaceV2.ID }) { const ctx = yield* InstanceState.context @@ -686,6 +695,7 @@ const layer: Layer.Layer< model: input?.model, metadata: input?.metadata, permission: input?.permission, + budget: input?.budget, workspaceID: input?.workspaceID ?? workspace, }) }) @@ -764,6 +774,10 @@ const layer: Layer.Layer< yield* patch(input.sessionID, { metadata: input.metadata, time: { updated: Date.now() } }).pipe(Effect.orDie) }) + const setBudget = Effect.fn("Session.setBudget")(function* (input: { sessionID: SessionID; budget?: number }) { + yield* patch(input.sessionID, { budget: input.budget, time: { updated: Date.now() } }).pipe(Effect.orDie) + }) + const setAgentModel = Effect.fn("Session.setAgentModel")(function* (input: { sessionID: SessionID agent: string @@ -915,6 +929,7 @@ const layer: Layer.Layer< setTitle, setArchived, setMetadata, + setBudget, setAgentModel, setPermission, setRevert, diff --git a/packages/opencode/test/session/processor-effect.test.ts b/packages/opencode/test/session/processor-effect.test.ts index 052477d0a2e7..4e5bf9c3d183 100644 --- a/packages/opencode/test/session/processor-effect.test.ts +++ b/packages/opencode/test/session/processor-effect.test.ts @@ -285,8 +285,53 @@ it.live("session.processor effect tests capture llm input cleanly", () => ), ) -it.live("session.processor effect tests preserve text start time", () => +it.live("session.processor effect tests stop when the session budget is exceeded", () => provideTmpdirServer( + ({ dir, llm }) => + Effect.gen(function* () { + const { processors, session, provider } = yield* boot() + + yield* llm.text("hello") + + const chat = yield* session.create({ budget: 0 }) + const parent = yield* user(chat.id, "hi") + const msg = yield* assistant(chat.id, parent.id, path.resolve(dir)) + const mdl = yield* provider.getModel(ref.providerID, ref.modelID) + const handle = yield* processors.create({ + assistantMessage: msg, + sessionID: chat.id, + model: mdl, + }) + + const value = yield* handle.process({ + user: { + id: parent.id, + sessionID: chat.id, + role: "user", + time: parent.time, + agent: parent.agent, + model: { providerID: ref.providerID, modelID: ref.modelID }, + } satisfies SessionV1.User, + sessionID: chat.id, + model: mdl, + agent: agent(), + system: [], + messages: [{ role: "user", content: "hi" }], + tools: {}, + } satisfies LLM.StreamInput) + const parts = yield* MessageV2.parts(msg.id) + + expect(value).toBe("stop") + expect( + parts.some((part) => part.type === "text" && part.text === "Session budget reached. Increase the budget to continue."), + ).toBe(true) + expect(msg.finish).toBe("stop") + }), + { config: (url) => providerCfg(url) }, + ), +) + +it.live("session.processor effect tests preserve text start time", () => provideTmpdirServer( ({ dir, llm }) => Effect.gen(function* () { const database = yield* Database.Service diff --git a/packages/opencode/test/session/prompt.test.ts b/packages/opencode/test/session/prompt.test.ts index 5a0176abc9b0..c41075f95c17 100644 --- a/packages/opencode/test/session/prompt.test.ts +++ b/packages/opencode/test/session/prompt.test.ts @@ -460,6 +460,27 @@ noLLMServer.instance( { config: cfg }, ) +it.instance("loop stops without an LLM request when the session budget is exhausted", () => + Effect.gen(function* () { + const { llm } = yield* useServerConfig(providerCfg) + const prompt = yield* SessionPrompt.Service + const sessions = yield* Session.Service + const chat = yield* sessions.create({ title: "Pinned", budget: 0 }) + const seeded = yield* seed(chat.id, { finish: "stop" }) + yield* prompt.prompt({ + sessionID: chat.id, + agent: "build", + noReply: true, + parts: [{ type: "text", text: "continue please" }], + }) + + const result = yield* prompt.loop({ sessionID: chat.id }) + + expect(result.info.id).toBe(seeded.assistant.id) + expect(yield* llm.hits).toHaveLength(0) + }), +) + noLLMServer.instance( "loop exits for a completed parent turn with nonmonotonic message IDs", () => diff --git a/packages/opencode/test/session/schema-decoding.test.ts b/packages/opencode/test/session/schema-decoding.test.ts index 3bcbd55d4b6a..420eeebe6ac1 100644 --- a/packages/opencode/test/session/schema-decoding.test.ts +++ b/packages/opencode/test/session/schema-decoding.test.ts @@ -65,6 +65,7 @@ describe("Session.Info", () => { title: "Full session", version: "1.0.0", metadata: { source: "test" }, + budget: 2.5, time: { created: 100, updated: 200, compacting: 150, archived: 300 }, permission: [{ action: "allow" as const, pattern: "*", permission: "read" }], revert: { diff --git a/packages/schema/src/session.ts b/packages/schema/src/session.ts index 937705eeb2f4..b8850ee394b1 100644 --- a/packages/schema/src/session.ts +++ b/packages/schema/src/session.ts @@ -23,6 +23,7 @@ export const Info = Schema.Struct({ agent: Agent.ID.pipe(optional), model: Model.Ref.pipe(optional), cost: Schema.Finite, + budget: Schema.Finite.pipe(optional), tokens: Schema.Struct({ input: Schema.Finite, output: Schema.Finite, diff --git a/packages/schema/src/v1/session.ts b/packages/schema/src/v1/session.ts index 75e9282f117c..3fe1b983e05d 100644 --- a/packages/schema/src/v1/session.ts +++ b/packages/schema/src/v1/session.ts @@ -550,6 +550,7 @@ export const SessionInfo = Schema.Struct({ parentID: optional(SessionID), summary: optional(SessionSummary), cost: optional(Schema.Finite), + budget: optional(Schema.Finite), tokens: optional(SessionTokens), share: optional(SessionShare), title: Schema.String, diff --git a/packages/sdk/js/script/build.ts b/packages/sdk/js/script/build.ts index 79e0879c9e14..b58962f78f40 100755 --- a/packages/sdk/js/script/build.ts +++ b/packages/sdk/js/script/build.ts @@ -82,7 +82,18 @@ const historyTypesPatched = generatedTypes.replace( if (historyTypesPatched === generatedTypes) { throw new Error("Session history numeric query patch did not apply") } -await Bun.write("./src/v2/gen/types.gen.ts", historyTypesPatched) +// The OpenAPI encoder drops the null arm of `Schema.NullOr` when it is wrapped +// in `Schema.optional`, so the generated update data types lose `null`. The +// runtime decoder still accepts null (it clears the session budget), so widen +// the update-side budget types to match the wire contract. +const budgetTypesPatched = historyTypesPatched.replace( + /(export type SessionUpdateData = \{[\s\S]*?budget\?: )number/, + "$1number | null", +) +if (budgetTypesPatched === historyTypesPatched) { + throw new Error("Session update budget nullable patch did not apply") +} +await Bun.write("./src/v2/gen/types.gen.ts", budgetTypesPatched) const generatedSdk = await Bun.file("./src/v2/gen/sdk.gen.ts").text() const historySdkPatched = generatedSdk.replace( @@ -92,7 +103,14 @@ const historySdkPatched = generatedSdk.replace( if (historySdkPatched === generatedSdk) { throw new Error("Session history numeric SDK patch did not apply") } -await Bun.write("./src/v2/gen/sdk.gen.ts", historySdkPatched) +const budgetSdkPatched = historySdkPatched.replace( + /(budget\?: )number([;,]\s*time\?: \{\s*archived\?: number)/, + "$1number | null$2", +) +if (budgetSdkPatched === historySdkPatched) { + throw new Error("Session update budget nullable SDK patch did not apply") +} +await Bun.write("./src/v2/gen/sdk.gen.ts", budgetSdkPatched) // Patch a @hey-api/openapi-ts codegen bug: SseFn incorrectly passes the // endpoint's TError into the second generic of ServerSentEventsResult, which diff --git a/packages/sdk/js/src/v2/gen/sdk.gen.ts b/packages/sdk/js/src/v2/gen/sdk.gen.ts index 9ed0084aac84..108cbc1b5435 100644 --- a/packages/sdk/js/src/v2/gen/sdk.gen.ts +++ b/packages/sdk/js/src/v2/gen/sdk.gen.ts @@ -3423,6 +3423,7 @@ export class Session2 extends HeyApiClient { [key: string]: unknown } permission?: PermissionRuleset + budget?: number workspaceID?: string }, options?: Options, @@ -3440,6 +3441,7 @@ export class Session2 extends HeyApiClient { { in: "body", key: "model" }, { in: "body", key: "metadata" }, { in: "body", key: "permission" }, + { in: "body", key: "budget" }, { in: "body", key: "workspaceID" }, ], }, @@ -3566,6 +3568,7 @@ export class Session2 extends HeyApiClient { [key: string]: unknown } permission?: PermissionRuleset + budget?: number | null time?: { archived?: number } @@ -3583,6 +3586,7 @@ export class Session2 extends HeyApiClient { { in: "body", key: "title" }, { in: "body", key: "metadata" }, { in: "body", key: "permission" }, + { in: "body", key: "budget" }, { in: "body", key: "time" }, ], }, diff --git a/packages/sdk/js/src/v2/gen/types.gen.ts b/packages/sdk/js/src/v2/gen/types.gen.ts index 90c91e9158cc..87a27612c5de 100644 --- a/packages/sdk/js/src/v2/gen/types.gen.ts +++ b/packages/sdk/js/src/v2/gen/types.gen.ts @@ -182,6 +182,7 @@ export type Session = { diffs?: Array } cost?: number + budget?: number tokens?: { input: number output: number @@ -2207,6 +2208,7 @@ export type GlobalSession = { diffs?: Array } cost?: number + budget?: number tokens?: { input: number output: number @@ -3909,6 +3911,7 @@ export type SessionV2Info = { agent?: string model?: ModelRef cost: number + budget?: number tokens: { input: number output: number @@ -9484,6 +9487,7 @@ export type SessionCreateData = { [key: string]: unknown } permission?: PermissionRuleset + budget?: number workspaceID?: string } path?: never @@ -9617,6 +9621,7 @@ export type SessionUpdateData = { [key: string]: unknown } permission?: PermissionRuleset + budget?: number | null time?: { archived?: number } diff --git a/packages/tui/src/feature-plugins/sidebar/context.tsx b/packages/tui/src/feature-plugins/sidebar/context.tsx index f1c99d9679ce..9778970758a2 100644 --- a/packages/tui/src/feature-plugins/sidebar/context.tsx +++ b/packages/tui/src/feature-plugins/sidebar/context.tsx @@ -1,7 +1,8 @@ import type { AssistantMessage } from "@opencode-ai/sdk/v2" import type { TuiPlugin, TuiPluginApi } from "@opencode-ai/plugin/tui" import type { BuiltinTuiPlugin } from "../builtins" -import { createMemo } from "solid-js" +import type { InputRenderable } from "@opentui/core" +import { createMemo, createSignal, Show } from "solid-js" const id = "internal:sidebar-context" @@ -10,11 +11,17 @@ const money = new Intl.NumberFormat("en-US", { currency: "USD", }) +const BUDGET_STEP = 0.25 + function View(props: { api: TuiPluginApi; session_id: string }) { const theme = () => props.api.theme.current const msg = createMemo(() => props.api.state.session.messages(props.session_id)) const session = createMemo(() => props.api.state.session.get(props.session_id)) const cost = createMemo(() => session()?.cost ?? 0) + const budget = createMemo(() => session()?.budget) + + const [editing, setEditing] = createSignal(false) + let input: InputRenderable | undefined const state = createMemo(() => { const last = msg().findLast((item): item is AssistantMessage => item.role === "assistant" && item.tokens.output > 0) @@ -34,6 +41,42 @@ function View(props: { api: TuiPluginApi; session_id: string }) { } }) + const saveBudget = (value: number | undefined) => { + void props.api.client.session.update({ + sessionID: props.session_id, + budget: value ?? null, + }) + } + + const startEditing = () => { + setEditing(true) + setTimeout(() => { + if (!input || input.isDestroyed) return + input.focus() + }, 1) + } + + const adjust = (direction: "up" | "down") => { + const current = budget() + const next = + (current ?? Math.max(cost() + BUDGET_STEP, BUDGET_STEP)) + (direction === "up" ? BUDGET_STEP : -BUDGET_STEP) + if (next < BUDGET_STEP) { + if (current !== undefined) saveBudget(undefined) + return + } + saveBudget(next) + } + + const budgetLabel = createMemo(() => { + const value = budget() + return value !== undefined ? `${money.format(value)} budget` : "unlimited budget" + }) + const exceeded = createMemo(() => { + const value = budget() + return value !== undefined && cost() >= value + }) + const budgetColor = createMemo(() => (exceeded() ? theme().error : theme().textMuted)) + return ( @@ -42,6 +85,40 @@ function View(props: { api: TuiPluginApi; session_id: string }) { {state().tokens.toLocaleString()} tokens {state().percent ?? 0}% used {money.format(cost())} spent + { + const parsed = Number.parseFloat(input?.value ?? "") + saveBudget(Number.isFinite(parsed) && parsed > 0 ? parsed : undefined) + setEditing(false) + }} + onKeyDown={(event) => { + if (event.name === "escape") setEditing(false) + }} + ref={(ref) => (input = ref)} + /> + } + > + { + if (event.scroll?.direction === "up") adjust("up") + else if (event.scroll?.direction === "down") adjust("down") + }} + onMouseUp={() => startEditing()} + > + + {money.format(cost())} / {budgetLabel()} + + + ) } diff --git a/packages/tui/test/feature-plugins/sidebar-context.test.tsx b/packages/tui/test/feature-plugins/sidebar-context.test.tsx new file mode 100644 index 000000000000..a1762e7beeb1 --- /dev/null +++ b/packages/tui/test/feature-plugins/sidebar-context.test.tsx @@ -0,0 +1,135 @@ +/** @jsxImportSource @opentui/solid */ +import { expect, test } from "bun:test" +import { testRender, type JSX } from "@opentui/solid" +import type { TuiPluginApi, TuiPluginMeta } from "@opencode-ai/plugin/tui" +import { createSignal } from "solid-js" + +const theme = { + text: "#ffffff", + textMuted: "#888888", + error: "#ff5555", + primary: "#007acc", + backgroundPanel: "#1e1e1e", +} + +function createHarness(budget: number | null) { + const updates: Array<{ sessionID: string; budget: number | null }> = [] + const session = () => ({ + id: "dummy", + slug: "budget-test", + title: "Budget session", + projectID: "project", + directory: "/tmp/opencode", + version: "0.0.0-test", + cost: 0.5, + ...(budget === null ? {} : { budget }), + tokens: { input: 100, output: 100, reasoning: 0, cache: { read: 0, write: 0 } }, + time: { created: 0, updated: 0 }, + }) + const [current, setCurrent] = createSignal(session()) + const registered: { + slots: Record unknown> + }[] = [] + const api = { + theme: { current: theme }, + state: { + session: { + messages: () => [], + get: () => current(), + }, + provider: [], + }, + client: { + session: { + update: async (params: { sessionID: string; budget?: number | null }) => { + updates.push({ sessionID: params.sessionID, budget: params.budget ?? null }) + budget = params.budget ?? null + setCurrent(session()) + return { data: session() } + }, + }, + }, + slots: { + register: (plugin: (typeof registered)[number]) => { + registered.push(plugin) + return "test-plugin" + }, + }, + } as unknown as TuiPluginApi + return { api, registered, updates } +} + +const waitForFrame = async (setup: Awaited>, text: string) => { + let frame = "" + for (let index = 0; index < 60; index++) { + await setup.renderOnce() + frame = setup.captureCharFrame() + if (frame.includes(text)) return frame + await new Promise((resolve) => setTimeout(resolve, 25)) + } + throw new Error(`frame never included "${text}"\n${frame}`) +} + +const findBudgetRow = (frame: string) => { + const rows = frame.split("\n") + const rowIndex = rows.findIndex((row) => row.includes("budget")) + const col = rows[rowIndex].indexOf("$0.50") + return { rowIndex, col } +} + +test("sidebar context widget renders, scrolls, and clears the session budget", async () => { + const { api, registered, updates } = createHarness(2) + const plugin = (await import("../../src/feature-plugins/sidebar/context")).default + await plugin.tui(api, undefined, {} as TuiPluginMeta) + const slot = registered[0].slots.sidebar_content + + const setup = await testRender(() => slot({}, { session_id: "dummy" }) as JSX.Element, { width: 60, height: 10 }) + + try { + let frame = await waitForFrame(setup, "$2.00 budget") + expect(frame).toContain("$0.50 / $2.00 budget") + const { rowIndex, col } = findBudgetRow(frame) + expect(rowIndex).toBeGreaterThan(-1) + expect(col).toBeGreaterThan(-1) + + await setup.mockMouse.scroll(col + 2, rowIndex, "up") + frame = await waitForFrame(setup, "$2.25 budget") + expect(updates).toContainEqual({ sessionID: "dummy", budget: 2.25 }) + + for (let index = 0; index < 12; index++) { + await setup.mockMouse.scroll(col + 2, rowIndex, "down") + await new Promise((resolve) => setTimeout(resolve, 10)) + } + frame = await waitForFrame(setup, "unlimited budget") + expect(updates).toContainEqual({ sessionID: "dummy", budget: null }) + } finally { + if (!setup.renderer.isDestroyed) setup.renderer.destroy() + } +}) + +test("sidebar context widget edits the budget by typing", async () => { + const { api, registered, updates } = createHarness(2) + const plugin = (await import("../../src/feature-plugins/sidebar/context")).default + await plugin.tui(api, undefined, {} as TuiPluginMeta) + const slot = registered[0].slots.sidebar_content + + const setup = await testRender(() => slot({}, { session_id: "dummy" }) as JSX.Element, { width: 60, height: 10 }) + + try { + const frame = await waitForFrame(setup, "$2.00 budget") + const { rowIndex, col } = findBudgetRow(frame) + + await setup.mockMouse.click(col + 2, rowIndex) + await setup.mockInput.typeText(".5") + setup.mockInput.pressEnter() + await waitForFrame(setup, "$2.50 budget") + expect(updates).toContainEqual({ sessionID: "dummy", budget: 2.5 }) + + await setup.mockMouse.click(col + 2, rowIndex) + setup.mockInput.pressEscape() + await waitForFrame(setup, "$2.50 budget") + expect(updates).toHaveLength(1) + } finally { + if (!setup.renderer.isDestroyed) setup.renderer.destroy() + } +})