diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 874086a06bdb..bf2ee0154f22 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -201,9 +201,13 @@ const layer = Layer.effect( const context = entries.map((entry) => entry.message) const isLastStep = agent.info?.steps !== undefined && currentStep >= agent.info.steps const toolMaterialization = isLastStep ? undefined : yield* tools.materialize(agent.info?.permissions) + // Agent opted into a fast model for lightweight steps. Prefer the small model when one is + // available; fall back to the session model otherwise (never fail a turn for a missing small model). + const resolvedSmall = agent.info?.small ? yield* models.resolveSmall(session) : undefined + const modelForStep = resolvedSmall ?? model const promptCacheKey = /^ses_[0-9a-f]{64}$/.test(session.id) ? session.id.slice(4) : session.id const request = LLM.request({ - model, + model: modelForStep, http: { headers: { "x-session-affinity": session.id, diff --git a/packages/core/src/session/runner/model.ts b/packages/core/src/session/runner/model.ts index 74e78120c20e..185a7b46d01f 100644 --- a/packages/core/src/session/runner/model.ts +++ b/packages/core/src/session/runner/model.ts @@ -73,12 +73,14 @@ export type Error = export interface Interface { readonly resolve: (session: SessionSchema.Info) => Effect.Effect + readonly resolveSmall: (session: SessionSchema.Info) => Effect.Effect } export class Service extends Context.Service()("@opencode/v2/SessionRunnerModel") {} /** Test or embedding seam for supplying a model resolver directly. */ -export const layerWith = (resolve: Interface["resolve"]) => Layer.succeed(Service, Service.of({ resolve })) +export const layerWith = (resolve: Interface["resolve"], resolveSmall?: Interface["resolveSmall"]) => + Layer.succeed(Service, Service.of({ resolve, resolveSmall: resolveSmall ?? (() => Effect.succeed(undefined)) })) const apiKey = (model: ModelV2.Info, credential?: Credential.Value) => { if (credential?.type === "key") return Auth.value(credential.key) @@ -211,6 +213,31 @@ export const locationLayer = Layer.effect( connection ? yield* integrations.connection.resolve(connection) : undefined, ) }), + resolveSmall: Effect.fn("SessionRunnerModel.resolveSmall")(function* (session) { + // Resolve the small model for the session's active provider, used for lightweight + // steps (title/status/confirmations) when the agent opts in via `small: true`. + const defaultModel = session.model ? undefined : yield* catalog.model.default() + const selected = session.model + ? (yield* catalog.model.available()).find( + (model) => model.providerID === session.model?.providerID && model.id === session.model.id, + ) + : defaultModel && supported(defaultModel) + ? defaultModel + : (yield* catalog.model.available()).find(supported) + if (!selected) return undefined + const small = yield* catalog.model.small(selected.providerID) + if (!small || !supported(small)) return undefined + if (small.id === selected.id) return undefined + const provider = yield* catalog.provider.get(small.providerID) + const connection = yield* integrations.connection.active( + provider?.integrationID ?? Integration.ID.make(small.providerID), + ) + return yield* resolve( + session, + small, + connection ? yield* integrations.connection.resolve(connection) : undefined, + ) + }), }) }), ) diff --git a/packages/core/test/session-runner-model.test.ts b/packages/core/test/session-runner-model.test.ts index 49bbce95a381..47a647b99a06 100644 --- a/packages/core/test/session-runner-model.test.ts +++ b/packages/core/test/session-runner-model.test.ts @@ -1,7 +1,7 @@ import { describe, expect } from "bun:test" import { LLM } from "@opencode-ai/llm" import { LLMClient } from "@opencode-ai/llm/route" -import { DateTime, Effect } from "effect" +import { DateTime, Effect, Layer } from "effect" import { Headers } from "effect/unstable/http" import { Credential } from "@opencode-ai/core/credential" import { Integration } from "@opencode-ai/core/integration" @@ -11,7 +11,7 @@ import { ProjectV2 } from "@opencode-ai/core/project" import { SessionRunnerModel } from "@opencode-ai/core/session/runner/model" import { SessionV2 } from "@opencode-ai/core/session" import { AbsolutePath } from "@opencode-ai/core/schema" -import { it } from "./lib/effect" +import { it, testEffect } from "./lib/effect" type Api = | { @@ -345,3 +345,51 @@ describe("SessionRunnerModel", () => { }), ) }) + +describe("SessionRunnerModel.resolveSmall", () => { + const session = SessionV2.Info.make({ + id: SessionV2.ID.make("ses_small"), + projectID: ProjectV2.ID.global, + title: "test", + cost: 0, + tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + time: { created: DateTime.makeUnsafe(0), updated: DateTime.makeUnsafe(0) }, + location: { directory: AbsolutePath.make("/project") }, + }) + + const smallIt = testEffect( + Layer.succeed( + SessionRunnerModel.Service, + SessionRunnerModel.Service.of({ + resolve: () => Effect.succeed({} as never), + resolveSmall: () => Effect.succeed(undefined), + }), + ), + ) + + smallIt.effect("returns undefined when no small model is injected", () => + Effect.gen(function* () { + const svc = yield* SessionRunnerModel.Service + const resolved = yield* svc.resolveSmall(session) + expect(resolved).toBeUndefined() + }), + ) + + const withSmall = testEffect( + Layer.succeed( + SessionRunnerModel.Service, + SessionRunnerModel.Service.of({ + resolve: () => Effect.succeed({} as never), + resolveSmall: () => Effect.succeed({ id: "small-model" } as never), + }), + ), + ) + + withSmall.effect("returns the injected small model when available", () => + Effect.gen(function* () { + const svc = yield* SessionRunnerModel.Service + const resolved = yield* svc.resolveSmall(session) + expect(resolved).toMatchObject({ id: "small-model" }) + }), + ) +}) diff --git a/packages/schema/src/agent.ts b/packages/schema/src/agent.ts index adf7aec2823f..2df7da4f6b7a 100644 --- a/packages/schema/src/agent.ts +++ b/packages/schema/src/agent.ts @@ -20,6 +20,7 @@ export interface Info extends Schema.Schema.Type {} export const Info = Schema.Struct({ id: ID, model: Model.Ref.pipe(optional), + small: Schema.Boolean.pipe(optional), request: Provider.Request, system: Schema.String.pipe(optional), description: Schema.String.pipe(optional),