diff --git a/src/features/background-agent/compaction-aware-message-resolver.test.ts b/src/features/background-agent/compaction-aware-message-resolver.test.ts index 17cc07024..5b9bed5af 100644 --- a/src/features/background-agent/compaction-aware-message-resolver.test.ts +++ b/src/features/background-agent/compaction-aware-message-resolver.test.ts @@ -2,7 +2,15 @@ import { describe, test, expect, beforeEach, afterEach } from "bun:test" import { mkdtempSync, writeFileSync, rmSync } from "node:fs" import { join } from "node:path" import { tmpdir } from "node:os" -import { isCompactionAgent, findNearestMessageExcludingCompaction } from "./compaction-aware-message-resolver" +import { + isCompactionAgent, + findNearestMessageExcludingCompaction, + resolvePromptContextFromSessionMessages, +} from "./compaction-aware-message-resolver" +import { + clearCompactionAgentConfigCheckpoint, + setCompactionAgentConfigCheckpoint, +} from "../../shared/compaction-agent-config-checkpoint" describe("isCompactionAgent", () => { describe("#given agent name variations", () => { @@ -65,6 +73,7 @@ describe("findNearestMessageExcludingCompaction", () => { afterEach(() => { rmSync(tempDir, { force: true, recursive: true }) + clearCompactionAgentConfigCheckpoint("ses_checkpoint") }) describe("#given directory with messages", () => { @@ -186,5 +195,65 @@ describe("findNearestMessageExcludingCompaction", () => { expect(result).not.toBeNull() expect(result?.agent).toBe("newer") }) + + test("merges partial metadata from multiple recent messages", () => { + // given + writeFileSync( + join(tempDir, "003.json"), + JSON.stringify({ model: { providerID: "anthropic", modelID: "claude-opus-4-1" } }), + ) + writeFileSync(join(tempDir, "002.json"), JSON.stringify({ agent: "atlas" })) + writeFileSync(join(tempDir, "001.json"), JSON.stringify({ tools: { bash: true } })) + + // when + const result = findNearestMessageExcludingCompaction(tempDir) + + // then + expect(result).toEqual({ + agent: "atlas", + model: { providerID: "anthropic", modelID: "claude-opus-4-1" }, + tools: { bash: true }, + }) + }) + + test("fills missing metadata from compaction checkpoint", () => { + // given + setCompactionAgentConfigCheckpoint("ses_checkpoint", { + agent: "sisyphus", + model: { providerID: "openai", modelID: "gpt-5" }, + }) + writeFileSync(join(tempDir, "001.json"), JSON.stringify({ tools: { bash: true } })) + + // when + const result = findNearestMessageExcludingCompaction(tempDir, "ses_checkpoint") + + // then + expect(result).toEqual({ + agent: "sisyphus", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: true }, + }) + }) + }) +}) + +describe("resolvePromptContextFromSessionMessages", () => { + test("merges partial prompt context from recent SDK messages", () => { + // given + const messages = [ + { info: { agent: "atlas" } }, + { info: { model: { providerID: "anthropic", modelID: "claude-opus-4-1" } } }, + { info: { tools: { bash: true } } }, + ] + + // when + const result = resolvePromptContextFromSessionMessages(messages) + + // then + expect(result).toEqual({ + agent: "atlas", + model: { providerID: "anthropic", modelID: "claude-opus-4-1" }, + tools: { bash: true }, + }) }) }) diff --git a/src/features/background-agent/compaction-aware-message-resolver.ts b/src/features/background-agent/compaction-aware-message-resolver.ts index 1bf94bfdb..60b3949b3 100644 --- a/src/features/background-agent/compaction-aware-message-resolver.ts +++ b/src/features/background-agent/compaction-aware-message-resolver.ts @@ -1,6 +1,21 @@ import { readdirSync, readFileSync } from "node:fs" import { join } from "node:path" import type { StoredMessage } from "../hook-message-injector" +import { getCompactionAgentConfigCheckpoint } from "../../shared/compaction-agent-config-checkpoint" + +type SessionMessage = { + info?: { + agent?: string + model?: { + providerID?: string + modelID?: string + variant?: string + } + providerID?: string + modelID?: string + tools?: StoredMessage["tools"] + } +} export function isCompactionAgent(agent: string | undefined): boolean { return agent?.trim().toLowerCase() === "compaction" @@ -16,42 +31,121 @@ function hasFullAgentAndModel(message: StoredMessage): boolean { function hasPartialAgentOrModel(message: StoredMessage): boolean { const hasAgent = !!message.agent && !isCompactionAgent(message.agent) const hasModel = !!message.model?.providerID && !!message.model?.modelID - return hasAgent || hasModel + return hasAgent || hasModel || !!message.tools } -export function findNearestMessageExcludingCompaction(messageDir: string): StoredMessage | null { - try { - const files = readdirSync(messageDir) - .filter((name) => name.endsWith(".json")) - .sort() - .reverse() - - for (const file of files) { - try { - const content = readFileSync(join(messageDir, file), "utf-8") - const parsed = JSON.parse(content) as StoredMessage - if (hasFullAgentAndModel(parsed)) { - return parsed - } - } catch { - continue - } - } - - for (const file of files) { - try { - const content = readFileSync(join(messageDir, file), "utf-8") - const parsed = JSON.parse(content) as StoredMessage - if (hasPartialAgentOrModel(parsed)) { - return parsed - } - } catch { - continue - } - } - } catch { +function convertSessionMessageToStoredMessage(message: SessionMessage): StoredMessage | null { + const info = message.info + if (!info) { return null } - return null + const providerID = info.model?.providerID ?? info.providerID + const modelID = info.model?.modelID ?? info.modelID + + return { + ...(info.agent ? { agent: info.agent } : {}), + ...(providerID && modelID + ? { + model: { + providerID, + modelID, + ...(info.model?.variant ? { variant: info.model.variant } : {}), + }, + } + : {}), + ...(info.tools ? { tools: info.tools } : {}), + } +} + +function mergeStoredMessages( + messages: Array, + sessionID?: string, +): StoredMessage | null { + const merged: StoredMessage = {} + + for (const message of messages) { + if (!message || isCompactionAgent(message.agent)) { + continue + } + + if (!merged.agent && message.agent) { + merged.agent = message.agent + } + + if (!merged.model?.providerID && message.model?.providerID && message.model.modelID) { + merged.model = { + providerID: message.model.providerID, + modelID: message.model.modelID, + ...(message.model.variant ? { variant: message.model.variant } : {}), + } + } + + if (!merged.tools && message.tools) { + merged.tools = message.tools + } + + if (hasFullAgentAndModel(merged) && merged.tools) { + break + } + } + + const checkpoint = sessionID + ? getCompactionAgentConfigCheckpoint(sessionID) + : undefined + + if (!merged.agent && checkpoint?.agent) { + merged.agent = checkpoint.agent + } + + if (!merged.model && checkpoint?.model) { + merged.model = { + providerID: checkpoint.model.providerID, + modelID: checkpoint.model.modelID, + } + } + + if (!merged.tools && checkpoint?.tools) { + merged.tools = checkpoint.tools + } + + return hasPartialAgentOrModel(merged) ? merged : null +} + +export function resolvePromptContextFromSessionMessages( + messages: SessionMessage[], + sessionID?: string, +): StoredMessage | null { + const convertedMessages = messages + .map(convertSessionMessageToStoredMessage) + .reverse() + + return mergeStoredMessages(convertedMessages, sessionID) +} + +export function findNearestMessageExcludingCompaction( + messageDir: string, + sessionID?: string, +): StoredMessage | null { + try { + const files = readdirSync(messageDir) + .filter((name: string) => name.endsWith(".json")) + .sort() + .reverse() + + const messages: Array = [] + + for (const file of files) { + try { + const content = readFileSync(join(messageDir, file), "utf-8") + messages.push(JSON.parse(content) as StoredMessage) + } catch { + continue + } + } + + return mergeStoredMessages(messages, sessionID) + } catch { + return null + } }