From a005a97fef26869afeb7aab989ae3f47d93b7cff Mon Sep 17 00:00:00 2001 From: Marc Terns Date: Sun, 11 Jan 2026 08:21:14 -0600 Subject: [PATCH] feat: add configurable DM history limit --- src/agents/pi-embedded-runner.test.ts | 91 +++++++++++++++++++++++++++ src/agents/pi-embedded-runner.ts | 48 ++++++++++++-- src/config/types.ts | 1 + src/config/zod-schema.ts | 1 + 4 files changed, 137 insertions(+), 4 deletions(-) diff --git a/src/agents/pi-embedded-runner.test.ts b/src/agents/pi-embedded-runner.test.ts index 6c85bff4b..c558b80e1 100644 --- a/src/agents/pi-embedded-runner.test.ts +++ b/src/agents/pi-embedded-runner.test.ts @@ -11,6 +11,7 @@ import { applyGoogleTurnOrderingFix, buildEmbeddedSandboxInfo, createSystemPromptOverride, + limitHistoryTurns, runEmbeddedPiAgent, splitSdkTools, } from "./pi-embedded-runner.js"; @@ -279,6 +280,96 @@ describe("applyGoogleTurnOrderingFix", () => { }); }); +describe("limitHistoryTurns", () => { + const makeMessages = (roles: ("user" | "assistant")[]): AgentMessage[] => + roles.map((role, i) => ({ + role, + content: [{ type: "text", text: `message ${i}` }], + })); + + it("returns all messages when limit is undefined", () => { + const messages = makeMessages(["user", "assistant", "user", "assistant"]); + expect(limitHistoryTurns(messages, undefined)).toBe(messages); + }); + + it("returns all messages when limit is 0", () => { + const messages = makeMessages(["user", "assistant", "user", "assistant"]); + expect(limitHistoryTurns(messages, 0)).toBe(messages); + }); + + it("returns all messages when limit is negative", () => { + const messages = makeMessages(["user", "assistant", "user", "assistant"]); + expect(limitHistoryTurns(messages, -1)).toBe(messages); + }); + + it("returns empty array when messages is empty", () => { + expect(limitHistoryTurns([], 5)).toEqual([]); + }); + + it("keeps all messages when fewer user turns than limit", () => { + const messages = makeMessages(["user", "assistant", "user", "assistant"]); + expect(limitHistoryTurns(messages, 10)).toBe(messages); + }); + + it("limits to last N user turns", () => { + const messages = makeMessages([ + "user", + "assistant", + "user", + "assistant", + "user", + "assistant", + ]); + const limited = limitHistoryTurns(messages, 2); + expect(limited.length).toBe(4); + expect(limited[0].content).toEqual([{ type: "text", text: "message 2" }]); + }); + + it("handles single user turn limit", () => { + const messages = makeMessages([ + "user", + "assistant", + "user", + "assistant", + "user", + "assistant", + ]); + const limited = limitHistoryTurns(messages, 1); + expect(limited.length).toBe(2); + expect(limited[0].content).toEqual([{ type: "text", text: "message 4" }]); + expect(limited[1].content).toEqual([{ type: "text", text: "message 5" }]); + }); + + it("handles messages with multiple assistant responses per user turn", () => { + const messages = makeMessages([ + "user", + "assistant", + "assistant", + "user", + "assistant", + ]); + const limited = limitHistoryTurns(messages, 1); + expect(limited.length).toBe(2); + expect(limited[0].role).toBe("user"); + expect(limited[1].role).toBe("assistant"); + }); + + it("preserves message content integrity", () => { + const messages: AgentMessage[] = [ + { role: "user", content: [{ type: "text", text: "first" }] }, + { + role: "assistant", + content: [{ type: "toolCall", id: "1", name: "bash", arguments: {} }], + }, + { role: "user", content: [{ type: "text", text: "second" }] }, + { role: "assistant", content: [{ type: "text", text: "response" }] }, + ]; + const limited = limitHistoryTurns(messages, 1); + expect(limited[0].content).toEqual([{ type: "text", text: "second" }]); + expect(limited[1].content).toEqual([{ type: "text", text: "response" }]); + }); +}); + describe("runEmbeddedPiAgent", () => { it("writes models.json into the provided agentDir", async () => { const agentDir = await fs.mkdtemp( diff --git a/src/agents/pi-embedded-runner.ts b/src/agents/pi-embedded-runner.ts index 48516bf54..98fb598a4 100644 --- a/src/agents/pi-embedded-runner.ts +++ b/src/agents/pi-embedded-runner.ts @@ -413,6 +413,38 @@ async function sanitizeSessionHistory(params: { }).messages; } +/** + * Limits conversation history to the last N user turns (and their associated + * assistant responses). This reduces token usage for long-running DM sessions. + * + * @param messages - The full message history + * @param limit - Max number of user turns to keep (undefined = no limit) + * @returns Messages trimmed to the last `limit` user turns + */ +export function limitHistoryTurns( + messages: AgentMessage[], + limit: number | undefined, +): AgentMessage[] { + if (!limit || limit <= 0 || messages.length === 0) return messages; + + // Count user messages from the end, find cutoff point + let userCount = 0; + let lastUserIndex = messages.length; + + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].role === "user") { + userCount++; + if (userCount > limit) { + // We exceeded the limit; keep from the last valid user turn onwards + return messages.slice(lastUserIndex); + } + lastUserIndex = i; + } + } + // Fewer than limit user turns, keep all + return messages; +} + const ACTIVE_EMBEDDED_RUNS = new Map(); type EmbeddedRunWaiter = { resolve: (ended: boolean) => void; @@ -1026,8 +1058,12 @@ export async function compactEmbeddedPiSession(params: { sessionId: params.sessionId, }); const validated = validateGeminiTurns(prior); - if (validated.length > 0) { - session.agent.replaceMessages(validated); + const limited = limitHistoryTurns( + validated, + params.config?.session?.dmHistoryLimit, + ); + if (limited.length > 0) { + session.agent.replaceMessages(limited); } const result = await session.compact(params.customInstructions); return { @@ -1417,8 +1453,12 @@ export async function runEmbeddedPiAgent(params: { sessionId: params.sessionId, }); const validated = validateGeminiTurns(prior); - if (validated.length > 0) { - session.agent.replaceMessages(validated); + const limited = limitHistoryTurns( + validated, + params.config?.session?.dmHistoryLimit, + ); + if (limited.length > 0) { + session.agent.replaceMessages(limited); } } catch (err) { session.dispose(); diff --git a/src/config/types.ts b/src/config/types.ts index 8a827a0ee..2419b9499 100644 --- a/src/config/types.ts +++ b/src/config/types.ts @@ -57,6 +57,7 @@ export type SessionConfig = { resetTriggers?: string[]; idleMinutes?: number; heartbeatIdleMinutes?: number; + dmHistoryLimit?: number; store?: string; typingIntervalSeconds?: number; typingMode?: TypingMode; diff --git a/src/config/zod-schema.ts b/src/config/zod-schema.ts index f7fa0c7c9..5e081336d 100644 --- a/src/config/zod-schema.ts +++ b/src/config/zod-schema.ts @@ -630,6 +630,7 @@ const SessionSchema = z resetTriggers: z.array(z.string()).optional(), idleMinutes: z.number().int().positive().optional(), heartbeatIdleMinutes: z.number().int().positive().optional(), + dmHistoryLimit: z.number().int().positive().optional(), store: z.string().optional(), typingIntervalSeconds: z.number().int().positive().optional(), typingMode: z