File size: 4,760 Bytes
bc575bc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | /**
* Verify the documentation example from extensions.md compiles and works.
*/
import { describe, expect, it, vi } from "vitest";
import type { ExtensionAPI, SessionBeforeCompactEvent, SessionCompactEvent } from "../src/core/extensions/index.ts";
vi.mock("@earendil-works/pi-coding-agent", () => ({
convertToLlm: (messages: unknown) => messages,
serializeConversation: () => "conversation",
}));
const { default: customCompactionExtension } = await import("../examples/extensions/custom-compaction.ts");
describe("Documentation example", () => {
it("custom compaction example should type-check correctly", () => {
// This is the example from extensions.md - verify it compiles
const exampleExtension = (pi: ExtensionAPI) => {
pi.on("session_before_compact", async (event: SessionBeforeCompactEvent, ctx) => {
// All these should be accessible on the event
const { preparation, branchEntries } = event;
// sessionManager, modelRegistry, and model come from ctx
const { sessionManager, modelRegistry } = ctx;
const { messagesToSummarize, turnPrefixMessages, tokensBefore, firstKeptEntryId, isSplitTurn } =
preparation;
// Verify types
expect(Array.isArray(messagesToSummarize)).toBe(true);
expect(Array.isArray(turnPrefixMessages)).toBe(true);
expect(typeof isSplitTurn).toBe("boolean");
expect(typeof tokensBefore).toBe("number");
expect(typeof sessionManager.getEntries).toBe("function");
expect(typeof modelRegistry.getApiKeyAndHeaders).toBe("function");
expect(typeof firstKeptEntryId).toBe("string");
expect(Array.isArray(branchEntries)).toBe(true);
const summary = messagesToSummarize
.filter((m) => m.role === "user")
.map((m) => `- ${typeof m.content === "string" ? m.content.slice(0, 100) : "[complex]"}`)
.join("\n");
// Extensions return compaction content - SessionManager adds id/parentId
return {
compaction: {
summary: `User requests:\n${summary}`,
firstKeptEntryId,
tokensBefore,
},
};
});
};
// Just verify the function exists and is callable
expect(typeof exampleExtension).toBe("function");
});
it("custom compaction example dispatches through modelRegistry.complete", async () => {
let handler: ((event: any, ctx: any) => Promise<any>) | undefined;
customCompactionExtension({
on(event, fn) {
if (event === "session_before_compact") handler = fn as typeof handler;
},
} as ExtensionAPI);
expect(handler).toBeDefined();
const complete = vi.fn(async () => ({
role: "assistant",
content: [{ type: "text", text: "custom provider summary" }],
provider: "example-custom",
api: "example-custom-api",
model: "summary-model",
stopReason: "stop",
usage: {
input: 1,
output: 2,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 3,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
timestamp: Date.now(),
}));
const model = {
provider: "example-custom",
api: "example-custom-api",
id: "summary-model",
name: "Summary Model",
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 1000,
maxTokens: 100,
};
const result = await handler!(
{
preparation: {
messagesToSummarize: [
{ role: "user", content: [{ type: "text", text: "please remember this" }], timestamp: Date.now() },
],
turnPrefixMessages: [],
tokensBefore: 42,
firstKeptEntryId: "entry-1",
},
branchEntries: [],
signal: new AbortController().signal,
},
{
ui: { notify: vi.fn() },
modelRegistry: {
find: vi.fn(() => model),
complete,
},
},
);
expect(complete).toHaveBeenCalledWith(
model,
expect.objectContaining({ messages: expect.any(Array) }),
expect.objectContaining({ maxTokens: 8192 }),
);
expect(complete).not.toHaveBeenCalledWith(
expect.anything(),
expect.anything(),
expect.objectContaining({ apiKey: expect.anything() }),
);
expect(result).toMatchObject({
compaction: {
summary: "custom provider summary",
firstKeptEntryId: "entry-1",
tokensBefore: 42,
},
});
});
it("compact event should have correct fields", () => {
const checkCompactEvent = (pi: ExtensionAPI) => {
pi.on("session_compact", async (event: SessionCompactEvent) => {
// These should all be accessible
const entry = event.compactionEntry;
const fromExtension = event.fromExtension;
expect(entry.type).toBe("compaction");
expect(typeof entry.summary).toBe("string");
expect(typeof entry.tokensBefore).toBe("number");
expect(typeof fromExtension).toBe("boolean");
});
};
expect(typeof checkCompactEvent).toBe("function");
});
});
|