test: share embedded compaction hook helpers

This commit is contained in:
Peter Steinberger
2026-03-13 23:58:43 +00:00
parent 34a552383f
commit 91b9c47dad

View File

@@ -93,6 +93,8 @@ vi.mock("@mariozechner/pi-ai/oauth", () => ({
vi.mock("@mariozechner/pi-coding-agent", () => {
return {
AuthStorage: class AuthStorage {},
ModelRegistry: class ModelRegistry {},
createAgentSession: vi.fn(async () => {
const session = {
sessionId: "session-1",
@@ -324,6 +326,57 @@ import { getApiProvider, unregisterApiProviders } from "@mariozechner/pi-ai";
import { getCustomApiRegistrySourceId } from "../custom-api-registry.js";
import { compactEmbeddedPiSessionDirect, compactEmbeddedPiSession } from "./compact.js";
const TEST_SESSION_ID = "session-1";
const TEST_SESSION_KEY = "agent:main:session-1";
const TEST_SESSION_FILE = "/tmp/session.jsonl";
const TEST_WORKSPACE_DIR = "/tmp";
const TEST_CUSTOM_INSTRUCTIONS = "focus on decisions";
function mockResolvedModel() {
resolveModelMock.mockReset();
resolveModelMock.mockReturnValue({
model: { provider: "openai", api: "responses", id: "fake", input: [] },
error: null,
authStorage: { setRuntimeApiKey: vi.fn() },
modelRegistry: {},
});
}
function compactionConfig(mode: "await" | "off" | "async") {
return {
agents: {
defaults: {
compaction: {
postIndexSync: mode,
},
},
},
} as never;
}
function directCompactionArgs(overrides: Record<string, unknown> = {}) {
return {
sessionId: TEST_SESSION_ID,
sessionKey: TEST_SESSION_KEY,
sessionFile: TEST_SESSION_FILE,
workspaceDir: TEST_WORKSPACE_DIR,
customInstructions: TEST_CUSTOM_INSTRUCTIONS,
...overrides,
};
}
function wrappedCompactionArgs(overrides: Record<string, unknown> = {}) {
return {
sessionId: TEST_SESSION_ID,
sessionKey: TEST_SESSION_KEY,
sessionFile: TEST_SESSION_FILE,
workspaceDir: TEST_WORKSPACE_DIR,
customInstructions: TEST_CUSTOM_INSTRUCTIONS,
enqueue: (task: () => unknown) => task(),
...overrides,
};
}
const sessionHook = (action: string) =>
triggerInternalHook.mock.calls.find(
(call) => call[0]?.type === "session" && call[0]?.action === action,
@@ -336,13 +389,7 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
hookRunner.hasHooks.mockReset();
hookRunner.runBeforeCompaction.mockReset();
hookRunner.runAfterCompaction.mockReset();
resolveModelMock.mockReset();
resolveModelMock.mockReturnValue({
model: { provider: "openai", api: "responses", id: "fake", input: [] },
error: null,
authStorage: { setRuntimeApiKey: vi.fn() },
modelRegistry: {},
});
mockResolvedModel();
sessionCompactImpl.mockReset();
sessionCompactImpl.mockResolvedValue({
summary: "summary",
@@ -376,14 +423,12 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
unregisterApiProviders(getCustomApiRegistrySourceId("ollama"));
});
async function runDirectCompaction(customInstructions = "focus on decisions") {
return await compactEmbeddedPiSessionDirect({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
customInstructions,
});
async function runDirectCompaction(customInstructions = TEST_CUSTOM_INSTRUCTIONS) {
return await compactEmbeddedPiSessionDirect(
directCompactionArgs({
customInstructions,
}),
);
}
it("bootstraps runtime plugins with the resolved workspace", async () => {
@@ -594,26 +639,15 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
},
});
const result = await compactEmbeddedPiSessionDirect({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
customInstructions: "focus on decisions",
config: {
agents: {
defaults: {
compaction: {
postIndexSync: "await",
},
},
},
} as never,
});
const result = await compactEmbeddedPiSessionDirect(
directCompactionArgs({
config: compactionConfig("await"),
}),
);
expect(result.ok).toBe(true);
expect(resolveSessionAgentIdMock).toHaveBeenCalledWith({
sessionKey: "agent:main:session-1",
sessionKey: TEST_SESSION_KEY,
config: expect.any(Object),
});
expect(getMemorySearchManagerMock).not.toHaveBeenCalled();
@@ -629,22 +663,11 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
getMemorySearchManagerMock.mockResolvedValue({ manager: { sync } });
let settled = false;
const resultPromise = compactEmbeddedPiSessionDirect({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
customInstructions: "focus on decisions",
config: {
agents: {
defaults: {
compaction: {
postIndexSync: "await",
},
},
},
} as never,
});
const resultPromise = compactEmbeddedPiSessionDirect(
directCompactionArgs({
config: compactionConfig("await"),
}),
);
void resultPromise.then(() => {
settled = true;
@@ -652,7 +675,7 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
await vi.waitFor(() => {
expect(sync).toHaveBeenCalledWith({
reason: "post-compaction",
sessionFiles: ["/tmp/session.jsonl"],
sessionFiles: [TEST_SESSION_FILE],
});
});
expect(settled).toBe(false);
@@ -666,22 +689,11 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
const sync = vi.fn(async () => {});
getMemorySearchManagerMock.mockResolvedValue({ manager: { sync } });
const result = await compactEmbeddedPiSessionDirect({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
customInstructions: "focus on decisions",
config: {
agents: {
defaults: {
compaction: {
postIndexSync: "off",
},
},
},
} as never,
});
const result = await compactEmbeddedPiSessionDirect(
directCompactionArgs({
config: compactionConfig("off"),
}),
);
expect(result.ok).toBe(true);
expect(resolveSessionAgentIdMock).not.toHaveBeenCalled();
@@ -698,22 +710,11 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
getMemorySearchManagerMock.mockImplementation(() => managerGate);
let settled = false;
const resultPromise = compactEmbeddedPiSessionDirect({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
customInstructions: "focus on decisions",
config: {
agents: {
defaults: {
compaction: {
postIndexSync: "async",
},
},
},
} as never,
});
const resultPromise = compactEmbeddedPiSessionDirect(
directCompactionArgs({
config: compactionConfig("async"),
}),
);
await vi.waitFor(() => {
expect(getMemorySearchManagerMock).toHaveBeenCalledTimes(1);
@@ -730,7 +731,7 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
await vi.waitFor(() => {
expect(sync).toHaveBeenCalledWith({
reason: "post-compaction",
sessionFiles: ["/tmp/session.jsonl"],
sessionFiles: [TEST_SESSION_FILE],
});
});
const result = await resultPromise;
@@ -790,35 +791,25 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => {
reason: undefined,
result: { summary: "engine-summary", tokensAfter: 50 },
});
resolveModelMock.mockReset();
resolveModelMock.mockReturnValue({
model: { provider: "openai", api: "responses", id: "fake", input: [] },
error: null,
authStorage: { setRuntimeApiKey: vi.fn() },
modelRegistry: {},
});
mockResolvedModel();
});
it("fires before_compaction with sentinel -1 and after_compaction on success", async () => {
hookRunner.hasHooks.mockReturnValue(true);
const result = await compactEmbeddedPiSession({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
messageChannel: "telegram",
customInstructions: "focus on decisions",
enqueue: (task) => task(),
});
const result = await compactEmbeddedPiSession(
wrappedCompactionArgs({
messageChannel: "telegram",
}),
);
expect(result.ok).toBe(true);
expect(result.compacted).toBe(true);
expect(hookRunner.runBeforeCompaction).toHaveBeenCalledWith(
{ messageCount: -1, sessionFile: "/tmp/session.jsonl" },
{ messageCount: -1, sessionFile: TEST_SESSION_FILE },
expect.objectContaining({
sessionKey: "agent:main:session-1",
sessionKey: TEST_SESSION_KEY,
messageProvider: "telegram",
}),
);
@@ -827,10 +818,10 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => {
messageCount: -1,
compactedCount: -1,
tokenCount: 50,
sessionFile: "/tmp/session.jsonl",
sessionFile: TEST_SESSION_FILE,
},
expect.objectContaining({
sessionKey: "agent:main:session-1",
sessionKey: TEST_SESSION_KEY,
messageProvider: "telegram",
}),
);
@@ -843,30 +834,19 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => {
getMemorySearchManagerMock.mockResolvedValue({ manager: { sync } });
try {
const result = await compactEmbeddedPiSession({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: " /tmp/session.jsonl ",
workspaceDir: "/tmp",
customInstructions: "focus on decisions",
enqueue: (task) => task(),
config: {
agents: {
defaults: {
compaction: {
postIndexSync: "await",
},
},
},
} as never,
});
const result = await compactEmbeddedPiSession(
wrappedCompactionArgs({
sessionFile: ` ${TEST_SESSION_FILE} `,
config: compactionConfig("await"),
}),
);
expect(result.ok).toBe(true);
expect(listener).toHaveBeenCalledTimes(1);
expect(listener).toHaveBeenCalledWith({ sessionFile: "/tmp/session.jsonl" });
expect(listener).toHaveBeenCalledWith({ sessionFile: TEST_SESSION_FILE });
expect(sync).toHaveBeenCalledWith({
reason: "post-compaction",
sessionFiles: ["/tmp/session.jsonl"],
sessionFiles: [TEST_SESSION_FILE],
});
} finally {
cleanup();
@@ -884,14 +864,7 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => {
result: undefined,
});
const result = await compactEmbeddedPiSession({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
customInstructions: "focus on decisions",
enqueue: (task) => task(),
});
const result = await compactEmbeddedPiSession(wrappedCompactionArgs());
expect(result.ok).toBe(false);
expect(hookRunner.runBeforeCompaction).toHaveBeenCalled();
@@ -910,23 +883,11 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => {
});
try {
const result = await compactEmbeddedPiSession({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
customInstructions: "focus on decisions",
enqueue: (task) => task(),
config: {
agents: {
defaults: {
compaction: {
postIndexSync: "await",
},
},
},
} as never,
});
const result = await compactEmbeddedPiSession(
wrappedCompactionArgs({
config: compactionConfig("await"),
}),
);
expect(result.ok).toBe(true);
expect(listener).not.toHaveBeenCalled();
@@ -940,14 +901,7 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => {
hookRunner.hasHooks.mockReturnValue(true);
hookRunner.runBeforeCompaction.mockRejectedValue(new Error("hook boom"));
const result = await compactEmbeddedPiSession({
sessionId: "session-1",
sessionKey: "agent:main:session-1",
sessionFile: "/tmp/session.jsonl",
workspaceDir: "/tmp",
customInstructions: "focus on decisions",
enqueue: (task) => task(),
});
const result = await compactEmbeddedPiSession(wrappedCompactionArgs());
expect(result.ok).toBe(true);
expect(result.compacted).toBe(true);