diff --git a/src/auto-reply/reply/agent-runner-memory.test.ts b/src/auto-reply/reply/agent-runner-memory.test.ts index d8ca2221298..500d777cf6b 100644 --- a/src/auto-reply/reply/agent-runner-memory.test.ts +++ b/src/auto-reply/reply/agent-runner-memory.test.ts @@ -67,8 +67,15 @@ describe("runMemoryFlushIfNeeded", () => { }; if (typeof params.newSessionId === "string" && params.newSessionId) { nextEntry.sessionId = params.newSessionId; - const storePath = typeof params.storePath === "string" ? params.storePath : rootDir; - nextEntry.sessionFile = path.join(path.dirname(storePath), `${params.newSessionId}.jsonl`); + if (typeof params.newSessionFile === "string" && params.newSessionFile) { + nextEntry.sessionFile = params.newSessionFile; + } else { + const storePath = typeof params.storePath === "string" ? params.storePath : rootDir; + nextEntry.sessionFile = path.join( + path.dirname(storePath), + `${params.newSessionId}.jsonl`, + ); + } } params.sessionStore[sessionKey] = nextEntry; if (typeof params.storePath === "string") { @@ -287,6 +294,76 @@ describe("runMemoryFlushIfNeeded", () => { ); }); + it("updates the active preflight run after transcript rotation", async () => { + const sessionFile = path.join(rootDir, "session.jsonl"); + const successorFile = path.join(rootDir, "session-rotated.jsonl"); + await fs.writeFile( + sessionFile, + `${JSON.stringify({ message: { role: "user", content: "x".repeat(5_000) } })}\n`, + "utf8", + ); + registerMemoryFlushPlanResolver(() => ({ + softThresholdTokens: 1, + forceFlushTranscriptBytes: 1_000_000_000, + reserveTokensFloor: 0, + prompt: "Pre-compaction memory flush.\nNO_REPLY", + systemPrompt: "Write memory to memory/YYYY-MM-DD.md.", + relativePath: "memory/2023-11-14.md", + })); + compactEmbeddedPiSessionMock.mockResolvedValueOnce({ + ok: true, + compacted: true, + result: { + tokensAfter: 42, + sessionId: "session-rotated", + sessionFile: successorFile, + }, + }); + const sessionEntry: SessionEntry = { + sessionId: "session", + sessionFile, + updatedAt: Date.now(), + totalTokensFresh: false, + }; + const sessionStore = { "agent:main:main": sessionEntry }; + const followupRun = createTestFollowupRun({ + sessionId: "session", + sessionFile, + sessionKey: "agent:main:main", + }); + const updateSessionId = vi.fn(); + const replyOperation = { + abortSignal: new AbortController().signal, + setPhase: vi.fn(), + updateSessionId, + } as never; + + const entry = await runPreflightCompactionIfNeeded({ + cfg: { agents: { defaults: { compaction: { memoryFlush: {} } } } }, + followupRun, + defaultModel: "anthropic/claude-opus-4-6", + agentCfgContextTokens: 100, + sessionEntry, + sessionStore, + sessionKey: "agent:main:main", + storePath: path.join(rootDir, "sessions.json"), + isHeartbeat: false, + replyOperation, + }); + + expect(entry?.sessionId).toBe("session-rotated"); + expect(entry?.sessionFile).toBe(successorFile); + expect(followupRun.run.sessionId).toBe("session-rotated"); + expect(followupRun.run.sessionFile).toBe(successorFile); + expect(updateSessionId).toHaveBeenCalledWith("session-rotated"); + expect(refreshQueuedFollowupSessionMock).toHaveBeenCalledWith({ + key: "agent:main:main", + previousSessionId: "session", + nextSessionId: "session-rotated", + nextSessionFile: successorFile, + }); + }); + it("uses configured prompts and stored bootstrap warning signatures", async () => { const sessionEntry: SessionEntry = { sessionId: "session", diff --git a/src/auto-reply/reply/agent-runner-memory.ts b/src/auto-reply/reply/agent-runner-memory.ts index 366e6febff7..640dbe18daa 100644 --- a/src/auto-reply/reply/agent-runner-memory.ts +++ b/src/auto-reply/reply/agent-runner-memory.ts @@ -514,6 +514,23 @@ export async function runPreflightCompactionIfNeeded(params: { followupRun: params.followupRun, }); entry = params.sessionStore?.[params.sessionKey] ?? entry; + if (entry) { + const previousSessionId = params.followupRun.run.sessionId; + params.followupRun.run.sessionId = entry.sessionId; + params.replyOperation.updateSessionId(entry.sessionId); + if (entry.sessionFile) { + params.followupRun.run.sessionFile = entry.sessionFile; + } + const queueKey = params.followupRun.run.sessionKey ?? params.sessionKey; + if (queueKey) { + memoryDeps.refreshQueuedFollowupSession({ + key: queueKey, + previousSessionId, + nextSessionId: entry.sessionId, + nextSessionFile: entry.sessionFile, + }); + } + } return entry ?? params.sessionEntry; }