fix: propagate variant field in all promptAsync continuation paths (#3081)

All 5 continuation paths now send variant as top-level body field:
- boulder-continuation-injector.ts
- ralph-loop/continuation-prompt-injector.ts
- todo-continuation-enforcer/continuation-injection.ts
- unstable-agent-babysitter-hook.ts
- session-recovery/resume.ts

Plus type/helper updates in atlas, todo-continuation-enforcer,
unstable-agent-babysitter, and session-recovery.

TDD: 18 regression tests added, all pass. tsc clean.
This commit is contained in:
YeonGyu-Kim
2026-04-07 15:10:38 +09:00
parent 3e8fd5ff18
commit aa528e42c0
15 changed files with 310 additions and 23 deletions
@@ -121,4 +121,64 @@ describe("injectBoulderContinuation", () => {
expect(result).toBe("skipped_agent_unavailable")
expect(promptAsyncMock).not.toHaveBeenCalled()
})
test("#given recent prompt context includes variant #when injecting boulder continuation #then promptAsync receives variant as a top-level field", async () => {
// given
registerAgentName("atlas")
const capturedRequests: Array<{
body?: {
model?: { providerID: string; modelID: string }
variant?: string
}
}> = []
const promptAsyncMock = mock(async (request: unknown) => {
capturedRequests.push(request as typeof capturedRequests[number])
return undefined
})
const recentModel = {
providerID: "anthropic",
modelID: "claude-sonnet-4-20250514",
variant: "max",
}
const messagesMock = mock(async () => ({
data: [{
id: "msg_1",
info: {
agent: "atlas",
model: recentModel,
time: { created: Date.now() },
},
}],
}))
const ctx = {
directory: "/tmp",
client: {
session: {
messages: messagesMock,
promptAsync: promptAsyncMock,
},
},
} as unknown as PluginInput
// when
const result = await injectBoulderContinuation({
ctx,
sessionID: "ses_test_variant",
planName: "test-plan",
remaining: 1,
total: 2,
agent: "atlas",
sessionState: { promptFailureCount: 0 },
})
// then
expect(result).toBe("injected")
expect(capturedRequests).toHaveLength(1)
expect(capturedRequests[0]?.body?.model).toEqual({
providerID: "anthropic",
modelID: "claude-sonnet-4-20250514",
})
expect(capturedRequests[0]?.body?.variant).toBe("max")
})
})
@@ -71,12 +71,18 @@ export async function injectBoulderContinuation(input: {
const promptContext = await resolveRecentPromptContextForSession(ctx, sessionID)
const inheritedTools = resolveInheritedPromptTools(sessionID, promptContext.tools)
await ctx.client.session.promptAsync({
path: { id: sessionID },
body: {
agent: continuationAgent,
...(promptContext.model !== undefined ? { model: promptContext.model } : {}),
...(inheritedTools ? { tools: inheritedTools } : {}),
const launchModel = promptContext.model
? { providerID: promptContext.model.providerID, modelID: promptContext.model.modelID }
: undefined
const launchVariant = promptContext.model?.variant
await ctx.client.session.promptAsync({
path: { id: sessionID },
body: {
agent: continuationAgent,
...(launchModel ? { model: launchModel } : {}),
...(launchVariant ? { variant: launchVariant } : {}),
...(inheritedTools ? { tools: inheritedTools } : {}),
parts: [createInternalAgentTextPart(prompt)],
},
query: { directory: ctx.directory },
+16 -2
View File
@@ -40,7 +40,14 @@ export async function resolveRecentPromptContextForSession(
const model = info?.model
const tools = normalizePromptTools(info?.tools)
if (model?.providerID && model?.modelID) {
return { model: { providerID: model.providerID, modelID: model.modelID }, tools }
return {
model: {
providerID: model.providerID,
modelID: model.modelID,
...(model.variant ? { variant: model.variant } : {}),
},
tools,
}
}
if (info?.providerID && info?.modelID) {
@@ -63,7 +70,14 @@ export async function resolveRecentPromptContextForSession(
if (!model?.providerID || !model?.modelID) {
return { tools }
}
return { model: { providerID: model.providerID, modelID: model.modelID }, tools }
return {
model: {
providerID: model.providerID,
modelID: model.modelID,
...(model.variant ? { variant: model.variant } : {}),
},
tools,
}
}
export async function resolveRecentModelForSession(
+1 -1
View File
@@ -2,7 +2,7 @@ import type { AgentOverrides } from "../../config"
import type { BackgroundManager } from "../../features/background-agent"
import type { TopLevelTaskRef } from "../../features/boulder-state"
export type ModelInfo = { providerID: string; modelID: string }
export type ModelInfo = { providerID: string; modelID: string; variant?: string }
export interface AtlasHookOptions {
directory: string
@@ -0,0 +1,52 @@
import { describe, expect, test } from "bun:test"
import { injectContinuationPrompt } from "./continuation-prompt-injector"
describe("ralph-loop continuation prompt injector", () => {
test("#given inherited message model includes variant #when injecting continuation prompt #then promptAsync receives variant as a top-level field", async () => {
// given
let promptBody:
| {
model?: { providerID: string; modelID: string }
variant?: string
}
| undefined
const model = {
providerID: "openai",
modelID: "gpt-5.3-codex",
variant: "max",
}
const ctx = {
client: {
session: {
messages: async () => ({
data: [{ info: { agent: "sisyphus", model } }],
}),
promptAsync: async (input: {
body: {
model?: { providerID: string; modelID: string }
variant?: string
}
}) => {
promptBody = input.body
return {}
},
},
},
}
// when
await injectContinuationPrompt(ctx as never, {
sessionID: "ses_ralph_variant",
prompt: "continue",
directory: "/tmp/test",
apiTimeoutMs: 50,
})
// then
expect(promptBody?.model).toEqual({
providerID: "openai",
modelID: "gpt-5.3-codex",
})
expect(promptBody?.variant).toBe("max")
})
})
@@ -11,7 +11,7 @@ import {
type MessageInfo = {
agent?: string
model?: { providerID: string; modelID: string }
model?: { providerID: string; modelID: string; variant?: string }
modelID?: string
providerID?: string
tools?: Record<string, boolean | "allow" | "deny" | "ask">
@@ -28,7 +28,7 @@ export async function injectContinuationPrompt(
},
): Promise<void> {
let agent: string | undefined
let model: { providerID: string; modelID: string } | undefined
let model: { providerID: string; modelID: string; variant?: string } | undefined
let tools: Record<string, boolean | "allow" | "deny" | "ask"> | undefined
const sourceSessionID = options.inheritFromSessionID ?? options.sessionID
@@ -62,6 +62,7 @@ export async function injectContinuationPrompt(
? {
providerID: currentMessage.model.providerID,
modelID: currentMessage.model.modelID,
...(currentMessage.model.variant ? { variant: currentMessage.model.variant } : {}),
}
: undefined
tools = currentMessage?.tools
@@ -69,11 +70,17 @@ export async function injectContinuationPrompt(
const inheritedTools = resolveInheritedPromptTools(sourceSessionID, tools)
const launchModel = model
? { providerID: model.providerID, modelID: model.modelID }
: undefined
const launchVariant = model?.variant
await ctx.client.session.promptAsync({
path: { id: options.sessionID },
body: {
...(agent !== undefined ? { agent } : {}),
...(model !== undefined ? { model } : {}),
...(launchModel ? { model: launchModel } : {}),
...(launchVariant ? { variant: launchVariant } : {}),
...(inheritedTools ? { tools: inheritedTools } : {}),
parts: [createInternalAgentTextPart(options.prompt)],
},
+30 -2
View File
@@ -22,9 +22,35 @@ describe("session-recovery resume", () => {
expect(config.tools).toEqual({ question: false, bash: true })
})
test("resumeSession sends inherited tools with continuation prompt", async () => {
test("#given the last user message includes model variant #when extracting resume config #then the variant is preserved", () => {
// given
const model = {
providerID: "openai",
modelID: "gpt-5.3-codex",
variant: "max",
}
const userMessage: MessageData = {
info: {
agent: "Hephaestus",
model,
},
}
// when
const config = extractResumeConfig(userMessage, "ses_resume_variant")
// then
expect(config.model).toEqual(model)
})
test("resumeSession sends inherited tools and variant with continuation prompt", async () => {
// given
let promptBody: Record<string, unknown> | undefined
const model = {
providerID: "openai",
modelID: "gpt-5.3-codex",
variant: "max",
}
const client = {
session: {
promptAsync: async (input: { body: Record<string, unknown> }) => {
@@ -38,12 +64,14 @@ describe("session-recovery resume", () => {
const ok = await resumeSession(client as never, {
sessionID: "ses_resume_prompt",
agent: "Hephaestus",
model: { providerID: "openai", modelID: "gpt-5.3-codex" },
model,
tools: { question: false, bash: true },
})
// then
expect(ok).toBe(true)
expect(promptBody?.model).toEqual({ providerID: "openai", modelID: "gpt-5.3-codex" })
expect(promptBody?.variant).toBe("max")
expect(promptBody?.tools).toEqual({ question: false, bash: true })
expect(Array.isArray(promptBody?.parts)).toBe(true)
const firstPart = (promptBody?.parts as Array<{ text?: string }>)?.[0]
+7 -1
View File
@@ -27,12 +27,18 @@ export function extractResumeConfig(userMessage: MessageData | undefined, sessio
export async function resumeSession(client: Client, config: ResumeConfig): Promise<boolean> {
try {
const inheritedTools = resolveInheritedPromptTools(config.sessionID, config.tools)
const launchModel = config.model
? { providerID: config.model.providerID, modelID: config.model.modelID }
: undefined
const launchVariant = config.model?.variant
await client.session.promptAsync({
path: { id: config.sessionID },
body: {
parts: [createInternalAgentTextPart(RECOVERY_RESUME_TEXT)],
agent: config.agent,
model: config.model,
...(launchModel ? { model: launchModel } : {}),
...(launchVariant ? { variant: launchVariant } : {}),
...(inheritedTools ? { tools: inheritedTools } : {}),
},
})
+2
View File
@@ -73,6 +73,7 @@ export interface MessageData {
model?: {
providerID: string
modelID: string
variant?: string
}
system?: string
tools?: Record<string, boolean>
@@ -94,6 +95,7 @@ export interface ResumeConfig {
model?: {
providerID: string
modelID: string
variant?: string
}
tools?: Record<string, boolean>
}
@@ -119,4 +119,57 @@ describe("injectContinuation", () => {
// then
expect(injected).toBe(false)
})
test("#given resolved model info includes variant #when reinjecting continuation #then promptAsync receives variant as a top-level field", async () => {
// given
let capturedBody:
| {
model?: { providerID: string; modelID: string }
variant?: string
}
| undefined
const ctx = {
directory: "/tmp/test",
client: {
session: {
todo: async () => ({ data: [{ id: "1", content: "todo", status: "pending", priority: "high" }] }),
promptAsync: async (input: {
body: {
model?: { providerID: string; modelID: string }
variant?: string
}
}) => {
capturedBody = input.body
return {}
},
},
},
}
const sessionStateStore = {
getExistingState: () => ({ inFlight: false, lastInjectedAt: 0, consecutiveFailures: 0 }),
}
const model = {
providerID: "openai",
modelID: "gpt-5.3-codex",
variant: "max",
}
// when
await injectContinuation({
ctx: ctx as never,
sessionID: "ses_continuation_variant",
resolvedInfo: {
agent: "Hephaestus",
model,
},
sessionStateStore: sessionStateStore as never,
})
// then
expect(capturedBody?.model).toEqual({
providerID: "openai",
modelID: "gpt-5.3-codex",
})
expect(capturedBody?.variant).toBe("max")
})
})
@@ -174,11 +174,17 @@ ${todoList}`
const inheritedTools = resolveInheritedPromptTools(sessionID, tools)
const launchModel = model
? { providerID: model.providerID, modelID: model.modelID }
: undefined
const launchVariant = model?.variant
await ctx.client.session.promptAsync({
path: { id: sessionID },
body: {
agent: promptAgent,
...(model !== undefined ? { model } : {}),
...(launchModel ? { model: launchModel } : {}),
...(launchVariant ? { variant: launchVariant } : {}),
...(inheritedTools ? { tools: inheritedTools } : {}),
parts: [createInternalAgentTextPart(prompt)],
},
@@ -45,7 +45,7 @@ export interface MessageInfo {
role?: string
error?: { name?: string; data?: unknown }
agent?: string
model?: { providerID: string; modelID: string }
model?: { providerID: string; modelID: string; variant?: string }
providerID?: string
modelID?: string
tools?: Record<string, ToolPermission>
@@ -57,7 +57,7 @@ export interface MessageWithInfo {
export interface ResolvedMessageInfo {
agent?: string
model?: { providerID: string; modelID: string }
model?: { providerID: string; modelID: string; variant?: string }
tools?: Record<string, ToolPermission>
}
@@ -214,4 +214,45 @@ describe("unstable-agent-babysitter hook", () => {
expect(promptCalls.length).toBe(1)
Date.now = originalNow
})
test("#given the main session model includes variant #when injecting a babysitter reminder #then promptAsync receives variant as a top-level field", async () => {
// given
setMainSession("main-1")
const promptCalls: Array<{ input: unknown }> = []
const mainModel = {
providerID: "openai",
modelID: "gpt-4",
variant: "max",
}
const ctx = createMockPluginInput({
messagesBySession: {
"main-1": [
{ info: { agent: "sisyphus", model: mainModel } },
],
"bg-1": [
{ info: { role: "assistant" }, parts: [{ type: "thinking", thinking: "deep thought" }] },
],
},
promptCalls,
})
const backgroundManager = createBackgroundManager([createTask()])
const hook = createUnstableAgentBabysitterHook(ctx, {
backgroundManager,
config: { timeout_ms: 120000 },
})
// when
await hook.event({ event: { type: "session.idle", properties: { sessionID: "main-1" } } })
// then
expect(promptCalls.length).toBe(1)
const payload = promptCalls[0].input as {
body?: {
model?: { providerID: string; modelID: string }
variant?: string
}
}
expect(payload.body?.model).toEqual({ providerID: "openai", modelID: "gpt-4" })
expect(payload.body?.variant).toBe("max")
})
})
@@ -5,7 +5,7 @@ export const THINKING_SUMMARY_MAX_CHARS = 500 as const
type MessageInfo = {
role?: string
agent?: string
model?: { providerID: string; modelID: string }
model?: { providerID: string; modelID: string; variant?: string }
providerID?: string
modelID?: string
tools?: Record<string, boolean | "allow" | "deny" | "ask">
@@ -33,7 +33,11 @@ export function getMessageInfo(value: unknown): MessageInfo | undefined {
? info.model
: undefined
const model = modelValue && typeof modelValue.providerID === "string" && typeof modelValue.modelID === "string"
? { providerID: modelValue.providerID, modelID: modelValue.modelID }
? {
providerID: modelValue.providerID,
modelID: modelValue.modelID,
...(typeof modelValue.variant === "string" ? { variant: modelValue.variant } : {}),
}
: undefined
return {
role: typeof info.role === "string" ? info.role : undefined,
@@ -30,6 +30,7 @@ type BabysitterContext = {
body: {
parts: Array<{ type: "text"; text: string }>
agent?: string
variant?: string
model?: { providerID: string; modelID: string }
tools?: Record<string, boolean>
}
@@ -40,6 +41,7 @@ type BabysitterContext = {
body: {
parts: Array<{ type: "text"; text: string }>
agent?: string
variant?: string
model?: { providerID: string; modelID: string }
tools?: Record<string, boolean>
}
@@ -58,9 +60,9 @@ type BabysitterOptions = {
async function resolveMainSessionTarget(
ctx: BabysitterContext,
sessionID: string
): Promise<{ agent?: string; model?: { providerID: string; modelID: string }; tools?: Record<string, boolean> }> {
): Promise<{ agent?: string; model?: { providerID: string; modelID: string; variant?: string }; tools?: Record<string, boolean> }> {
let agent = getSessionAgent(sessionID)
let model: { providerID: string; modelID: string } | undefined
let model: { providerID: string; modelID: string; variant?: string } | undefined
let tools: Record<string, boolean> | undefined
try {
@@ -206,11 +208,17 @@ export function createUnstableAgentBabysitterHook(ctx: BabysitterContext, option
const { agent, model, tools } = await resolveMainSessionTarget(ctx, mainSessionID)
try {
const launchModel = model
? { providerID: model.providerID, modelID: model.modelID }
: undefined
const launchVariant = model?.variant
await ctx.client.session.promptAsync({
path: { id: mainSessionID },
body: {
...(agent ? { agent } : {}),
...(model ? { model } : {}),
...(launchModel ? { model: launchModel } : {}),
...(launchVariant ? { variant: launchVariant } : {}),
...(tools ? { tools } : {}),
parts: [createInternalAgentTextPart(reminder)],
},