feat(task): use small model for subagents

This commit is contained in:
Shoubhit Dash 2026-04-29 17:10:19 +05:30
commit 6ba2bbc91c
2 changed files with 61 additions and 6 deletions

View file

@ -6,6 +6,7 @@ import { MessageV2 } from "../session/message-v2"
import { Agent } from "../agent/agent" import { Agent } from "../agent/agent"
import type { SessionPrompt } from "../session/prompt" import type { SessionPrompt } from "../session/prompt"
import { Config } from "@/config/config" import { Config } from "@/config/config"
import { Provider } from "@/provider/provider"
import { Effect, Schema } from "effect" import { Effect, Schema } from "effect"
export interface TaskPromptOps { export interface TaskPromptOps {
@ -32,6 +33,7 @@ export const TaskTool = Tool.define(
Effect.gen(function* () { Effect.gen(function* () {
const agent = yield* Agent.Service const agent = yield* Agent.Service
const config = yield* Config.Service const config = yield* Config.Service
const provider = yield* Provider.Service
const sessions = yield* Session.Service const sessions = yield* Session.Service
const run = Effect.fn("TaskTool.execute")(function* ( const run = Effect.fn("TaskTool.execute")(function* (
@ -99,9 +101,10 @@ export const TaskTool = Tool.define(
const msg = yield* Effect.sync(() => MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID })) const msg = yield* Effect.sync(() => MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID }))
if (msg.info.role !== "assistant") return yield* Effect.fail(new Error("Not an assistant message")) if (msg.info.role !== "assistant") return yield* Effect.fail(new Error("Not an assistant message"))
const smallModel = next.model ? undefined : yield* provider.getSmallModel(msg.info.providerID)
const model = next.model ?? { const model = next.model ?? {
modelID: msg.info.modelID, modelID: smallModel?.id ?? msg.info.modelID,
providerID: msg.info.providerID, providerID: smallModel?.providerID ?? msg.info.providerID,
} }
yield* ctx.metadata({ yield* ctx.metadata({

View file

@ -2,6 +2,7 @@ import { afterEach, describe, expect } from "bun:test"
import { Effect, Layer } from "effect" import { Effect, Layer } from "effect"
import { Agent } from "../../src/agent/agent" import { Agent } from "../../src/agent/agent"
import { Config } from "@/config/config" import { Config } from "@/config/config"
import { Provider } from "@/provider/provider"
import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner" import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
import { Instance } from "../../src/project/instance" import { Instance } from "../../src/project/instance"
import { Session } from "@/session/session" import { Session } from "@/session/session"
@ -28,6 +29,7 @@ const it = testEffect(
Layer.mergeAll( Layer.mergeAll(
Agent.defaultLayer, Agent.defaultLayer,
Config.defaultLayer, Config.defaultLayer,
Provider.defaultLayer,
CrossSpawnSpawner.defaultLayer, CrossSpawnSpawner.defaultLayer,
Session.defaultLayer, Session.defaultLayer,
Truncate.defaultLayer, Truncate.defaultLayer,
@ -35,7 +37,7 @@ const it = testEffect(
), ),
) )
const seed = Effect.fn("TaskToolTest.seed")(function* (title = "Pinned") { const seed = Effect.fn("TaskToolTest.seed")(function* (title = "Pinned", model = ref) {
const session = yield* Session.Service const session = yield* Session.Service
const chat = yield* session.create({ title }) const chat = yield* session.create({ title })
const user = yield* session.updateMessage({ const user = yield* session.updateMessage({
@ -43,7 +45,7 @@ const seed = Effect.fn("TaskToolTest.seed")(function* (title = "Pinned") {
role: "user", role: "user",
sessionID: chat.id, sessionID: chat.id,
agent: "build", agent: "build",
model: ref, model,
time: { created: Date.now() }, time: { created: Date.now() },
}) })
const assistant: MessageV2.Assistant = { const assistant: MessageV2.Assistant = {
@ -56,8 +58,8 @@ const seed = Effect.fn("TaskToolTest.seed")(function* (title = "Pinned") {
cost: 0, cost: 0,
path: { cwd: "/tmp", root: "/tmp" }, path: { cwd: "/tmp", root: "/tmp" },
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
modelID: ref.modelID, modelID: model.modelID,
providerID: ref.providerID, providerID: model.providerID,
time: { created: Date.now() }, time: { created: Date.now() },
} }
yield* session.updateMessage(assistant) yield* session.updateMessage(assistant)
@ -274,6 +276,56 @@ describe("tool.task", () => {
), ),
) )
it.live("execute defaults subagents to the provider small model", () =>
provideTmpdirInstance(
() =>
Effect.gen(function* () {
const current = {
providerID: ProviderID.openai,
modelID: ModelID.make("gpt-5"),
}
const { chat, assistant } = yield* seed("Pinned", current)
const tool = yield* TaskTool
const def = yield* tool.init()
let seen: SessionPrompt.PromptInput | undefined
const promptOps = stubOps({ onPrompt: (input) => (seen = input) })
const result = yield* def.execute(
{
description: "inspect bug",
prompt: "look into the cache key path",
subagent_type: "general",
},
{
sessionID: chat.id,
messageID: assistant.id,
agent: "build",
abort: new AbortController().signal,
extra: { promptOps },
messages: [],
metadata: () => Effect.void,
ask: () => Effect.void,
},
)
expect(result.metadata.model).toEqual({
providerID: ProviderID.openai,
modelID: ModelID.make("gpt-5.4-mini"),
})
expect(seen?.model).toEqual(result.metadata.model)
}),
{
config: {
provider: {
openai: {
options: { apiKey: "test" },
},
},
},
},
),
)
it.live("execute creates a child when task_id does not exist", () => it.live("execute creates a child when task_id does not exist", () =>
provideTmpdirInstance(() => provideTmpdirInstance(() =>
Effect.gen(function* () { Effect.gen(function* () {