feat: add configurable DM history limit
This commit is contained in:
@@ -11,6 +11,7 @@ import {
|
|||||||
applyGoogleTurnOrderingFix,
|
applyGoogleTurnOrderingFix,
|
||||||
buildEmbeddedSandboxInfo,
|
buildEmbeddedSandboxInfo,
|
||||||
createSystemPromptOverride,
|
createSystemPromptOverride,
|
||||||
|
limitHistoryTurns,
|
||||||
runEmbeddedPiAgent,
|
runEmbeddedPiAgent,
|
||||||
splitSdkTools,
|
splitSdkTools,
|
||||||
} from "./pi-embedded-runner.js";
|
} 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", () => {
|
describe("runEmbeddedPiAgent", () => {
|
||||||
it("writes models.json into the provided agentDir", async () => {
|
it("writes models.json into the provided agentDir", async () => {
|
||||||
const agentDir = await fs.mkdtemp(
|
const agentDir = await fs.mkdtemp(
|
||||||
|
|||||||
@@ -413,6 +413,38 @@ async function sanitizeSessionHistory(params: {
|
|||||||
}).messages;
|
}).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<string, EmbeddedPiQueueHandle>();
|
const ACTIVE_EMBEDDED_RUNS = new Map<string, EmbeddedPiQueueHandle>();
|
||||||
type EmbeddedRunWaiter = {
|
type EmbeddedRunWaiter = {
|
||||||
resolve: (ended: boolean) => void;
|
resolve: (ended: boolean) => void;
|
||||||
@@ -1026,8 +1058,12 @@ export async function compactEmbeddedPiSession(params: {
|
|||||||
sessionId: params.sessionId,
|
sessionId: params.sessionId,
|
||||||
});
|
});
|
||||||
const validated = validateGeminiTurns(prior);
|
const validated = validateGeminiTurns(prior);
|
||||||
if (validated.length > 0) {
|
const limited = limitHistoryTurns(
|
||||||
session.agent.replaceMessages(validated);
|
validated,
|
||||||
|
params.config?.session?.dmHistoryLimit,
|
||||||
|
);
|
||||||
|
if (limited.length > 0) {
|
||||||
|
session.agent.replaceMessages(limited);
|
||||||
}
|
}
|
||||||
const result = await session.compact(params.customInstructions);
|
const result = await session.compact(params.customInstructions);
|
||||||
return {
|
return {
|
||||||
@@ -1417,8 +1453,12 @@ export async function runEmbeddedPiAgent(params: {
|
|||||||
sessionId: params.sessionId,
|
sessionId: params.sessionId,
|
||||||
});
|
});
|
||||||
const validated = validateGeminiTurns(prior);
|
const validated = validateGeminiTurns(prior);
|
||||||
if (validated.length > 0) {
|
const limited = limitHistoryTurns(
|
||||||
session.agent.replaceMessages(validated);
|
validated,
|
||||||
|
params.config?.session?.dmHistoryLimit,
|
||||||
|
);
|
||||||
|
if (limited.length > 0) {
|
||||||
|
session.agent.replaceMessages(limited);
|
||||||
}
|
}
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
session.dispose();
|
session.dispose();
|
||||||
|
|||||||
@@ -57,6 +57,7 @@ export type SessionConfig = {
|
|||||||
resetTriggers?: string[];
|
resetTriggers?: string[];
|
||||||
idleMinutes?: number;
|
idleMinutes?: number;
|
||||||
heartbeatIdleMinutes?: number;
|
heartbeatIdleMinutes?: number;
|
||||||
|
dmHistoryLimit?: number;
|
||||||
store?: string;
|
store?: string;
|
||||||
typingIntervalSeconds?: number;
|
typingIntervalSeconds?: number;
|
||||||
typingMode?: TypingMode;
|
typingMode?: TypingMode;
|
||||||
|
|||||||
@@ -630,6 +630,7 @@ const SessionSchema = z
|
|||||||
resetTriggers: z.array(z.string()).optional(),
|
resetTriggers: z.array(z.string()).optional(),
|
||||||
idleMinutes: z.number().int().positive().optional(),
|
idleMinutes: z.number().int().positive().optional(),
|
||||||
heartbeatIdleMinutes: z.number().int().positive().optional(),
|
heartbeatIdleMinutes: z.number().int().positive().optional(),
|
||||||
|
dmHistoryLimit: z.number().int().positive().optional(),
|
||||||
store: z.string().optional(),
|
store: z.string().optional(),
|
||||||
typingIntervalSeconds: z.number().int().positive().optional(),
|
typingIntervalSeconds: z.number().int().positive().optional(),
|
||||||
typingMode: z
|
typingMode: z
|
||||||
|
|||||||
Reference in New Issue
Block a user