@flink-app/flink 0.14.3 → 2.0.0-alpha.100

package/spec/ai/ContextCompaction.spec.ts

@@ -0,0 +1,405 @@
+import { FlinkAgent } from "../../src/ai/FlinkAgent";
+import { FlinkContext } from "../../src/FlinkContext";
+import { createStreamingMock } from "../testHelpers";
+import { LLMAdapter, LLMMessage } from "../../src/ai/LLMAdapter";
+
+// Test agent with context compaction
+class CompactingAgent extends FlinkAgent<FlinkContext> {
+    id = "compacting-agent";
+    description = "Test agent with context compaction";
+    instructions() { return "You are a test assistant"; }
+
+    setCallbacks(
+        shouldCompact?: (messages: LLMMessage[], step: number) => boolean | Promise<boolean>,
+        compactHistory?: (messages: LLMMessage[], step: number) => Promise<LLMMessage[]> | LLMMessage[]
+    ) {
+        (this as any).shouldCompact = shouldCompact;
+        (this as any).compactHistory = compactHistory;
+        // Clear runner cache to pick up new callbacks
+        (this as any).runner = undefined;
+    }
+
+    async query(message: string) {
+        const response = (this as any).execute({ message });
+        return await response.result;
+    }
+}
+
+describe("Context Compaction", () => {
+    let ctx: FlinkContext;
+    let agent: CompactingAgent;
+    let llmAdapter: LLMAdapter;
+    let llmAdapters: Map<string, LLMAdapter>;
+
+    beforeEach(() => {
+        ctx = {} as FlinkContext;
+        agent = new CompactingAgent();
+        agent.ctx = ctx;
+    });
+
+    describe("Basic compaction trigger", () => {
+        it("should not compact when shouldCompact is not provided", async () => {
+            // Create agent without compaction callbacks
+            agent.setCallbacks(undefined, undefined);
+
+            // Mock LLM to return a response
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Hello!",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const result = await agent.query("Hi");
+
+            expect(result.message).toBe("Hello!");
+            // Should only be called once (no compaction)
+            expect((llmAdapter.stream as jasmine.Spy).calls.count()).toBe(1);
+        });
+
+        it("should compact when shouldCompact returns true", async () => {
+            let compactCalled = false;
+            let messagesBeforeCompact: LLMMessage[] = [];
+
+            agent.setCallbacks(
+                (messages, step) => {
+                    messagesBeforeCompact = [...messages];
+                    return messages.length > 5;
+                },
+                (messages, step) => {
+                    compactCalled = true;
+                    return messages.slice(-3); // Keep last 3 messages
+                }
+            );
+
+            // Mock LLM to return a response
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            // Create history with 5 messages
+            const history = [
+                { role: "user" as const, content: "Message 1" },
+                { role: "assistant" as const, content: "Response 1" },
+                { role: "user" as const, content: "Message 2" },
+                { role: "assistant" as const, content: "Response 2" },
+                { role: "user" as const, content: "Message 3" },
+            ];
+
+            const result = await (agent as any).execute({ message: "New message", history }).result;
+
+            expect(result.message).toBe("Response");
+            expect(compactCalled).toBe(true);
+            expect(messagesBeforeCompact.length).toBe(6); // 5 history + 1 new message
+        });
+
+        it("should not compact when shouldCompact returns false", async () => {
+            let compactCalled = false;
+
+            agent.setCallbacks(
+                (messages) => false,
+                (messages) => {
+                    compactCalled = true;
+                    return messages.slice(-3);
+                }
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const result = await agent.query("Message");
+
+            expect(result.message).toBe("Response");
+            expect(compactCalled).toBe(false);
+        });
+    });
+
+    describe("Custom compactHistory strategies", () => {
+        it("should support sliding window strategy", async () => {
+            agent.setCallbacks(
+                (messages) => messages.length > 5,
+                (messages) => messages.slice(-3) // Keep last 3
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const history: LLMMessage[] = [
+                { role: "user", content: "1" },
+                { role: "assistant", content: "2" },
+                { role: "user", content: "3" },
+                { role: "assistant", content: "4" },
+                { role: "user", content: "5" },
+            ];
+
+            await (agent as any).execute({ message: "6", history }).result;
+
+            // Verify LLM received compacted messages
+            const llmCall = (llmAdapter.stream as jasmine.Spy).calls.mostRecent().args[0];
+            // Should have compacted to ~3 messages (may be +1 due to internal processing)
+            expect(llmCall.messages.length).toBeLessThanOrEqual(4);
+        });
+
+        it("should support keep first + last N strategy", async () => {
+            agent.setCallbacks(
+                (messages) => messages.length > 6,
+                (messages) => {
+                    if (messages.length <= 1) return messages;
+                    return [messages[0], ...messages.slice(-3)];
+                }
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const history: LLMMessage[] = [
+                { role: "user", content: "FIRST" },
+                { role: "assistant", content: "2" },
+                { role: "user", content: "3" },
+                { role: "assistant", content: "4" },
+                { role: "user", content: "5" },
+                { role: "assistant", content: "6" },
+            ];
+
+            await (agent as any).execute({ message: "LAST", history }).result;
+
+            const llmCall = (llmAdapter.stream as jasmine.Spy).calls.mostRecent().args[0];
+            // Should have first message + last N
+            expect(llmCall.messages.length).toBeLessThanOrEqual(5);
+            expect(llmCall.messages[0].content).toBe("FIRST");
+            // Verify compaction preserved first message
+            const userMessages = llmCall.messages.filter((m: any) => m.role === "user");
+            expect(userMessages.length).toBeGreaterThan(0);
+            expect(userMessages[0].content).toBe("FIRST");
+        });
+    });
+
+    describe("Default behavior", () => {
+        it("should use default strategy when shouldCompact is true but compactHistory is not provided", async () => {
+            agent.setCallbacks(
+                (messages) => messages.length > 5,
+                undefined // No custom strategy
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const history: LLMMessage[] = [];
+            for (let i = 0; i < 15; i++) {
+                history.push({ role: i % 2 === 0 ? "user" : "assistant", content: `Message ${i}` });
+            }
+
+            await (agent as any).execute({ message: "New", history }).result;
+
+            const llmCall = (llmAdapter.stream as jasmine.Spy).calls.mostRecent().args[0];
+            // Default strategy keeps last 10 messages (may be +1 due to internal processing)
+            expect(llmCall.messages.length).toBeLessThanOrEqual(11);
+            expect(llmCall.messages.length).toBeGreaterThan(5);
+        });
+    });
+
+    describe("Async callbacks", () => {
+        it("should support async shouldCompact", async () => {
+            agent.setCallbacks(
+                async (messages) => {
+                    await new Promise((resolve) => setTimeout(resolve, 10));
+                    return messages.length > 5;
+                },
+                (messages) => messages.slice(-3)
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const history: LLMMessage[] = [];
+            for (let i = 0; i < 6; i++) {
+                history.push({ role: i % 2 === 0 ? "user" : "assistant", content: `Message ${i}` });
+            }
+
+            await (agent as any).execute({ message: "New", history }).result;
+
+            const llmCall = (llmAdapter.stream as jasmine.Spy).calls.mostRecent().args[0];
+            // Should be compacted to ~3 messages
+            expect(llmCall.messages.length).toBeLessThanOrEqual(4);
+        });
+
+        it("should support async compactHistory", async () => {
+            agent.setCallbacks(
+                (messages) => messages.length > 5,
+                async (messages) => {
+                    await new Promise((resolve) => setTimeout(resolve, 10));
+                    return messages.slice(-3);
+                }
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const history: LLMMessage[] = [];
+            for (let i = 0; i < 6; i++) {
+                history.push({ role: i % 2 === 0 ? "user" : "assistant", content: `Message ${i}` });
+            }
+
+            await (agent as any).execute({ message: "New", history }).result;
+
+            const llmCall = (llmAdapter.stream as jasmine.Spy).calls.mostRecent().args[0];
+            // Should be compacted to ~3 messages
+            expect(llmCall.messages.length).toBeLessThanOrEqual(4);
+        });
+    });
+
+    describe("Error handling", () => {
+        it("should log error and continue when compactHistory returns empty array", async () => {
+            spyOn(console, "error"); // Spy on console.error to verify logging
+
+            agent.setCallbacks(
+                (messages) => messages.length > 5,
+                (messages) => [] // Invalid - empty array
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const history: LLMMessage[] = [];
+            for (let i = 0; i < 6; i++) {
+                history.push({ role: i % 2 === 0 ? "user" : "assistant", content: `Message ${i}` });
+            }
+
+            // Should not throw, but continue with original messages
+            const result = await (agent as any).execute({ message: "New", history }).result;
+
+            expect(result.message).toBe("Response");
+
+            // LLM should receive original (uncompacted) messages due to error
+            const llmCall = (llmAdapter.stream as jasmine.Spy).calls.mostRecent().args[0];
+            // Should have all original messages (6 history + 1 new = 7, or may be +1)
+            expect(llmCall.messages.length).toBeGreaterThanOrEqual(7);
+        });
+
+        it("should log error and continue when shouldCompact throws exception", async () => {
+            spyOn(console, "error");
+
+            agent.setCallbacks(
+                (messages) => {
+                    throw new Error("Test error in shouldCompact");
+                },
+                (messages) => messages.slice(-3)
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            // Should not throw, execution continues
+            const result = await agent.query("Message");
+
+            expect(result.message).toBe("Response");
+        });
+
+        it("should log error and continue when compactHistory throws exception", async () => {
+            spyOn(console, "error");
+
+            agent.setCallbacks(
+                (messages) => messages.length > 5,
+                (messages) => {
+                    throw new Error("Test error in compactHistory");
+                }
+            );
+
+            llmAdapter = createStreamingMock([
+                {
+                    textContent: "Response",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 5 },
+                    stopReason: "end_turn",
+                },
+            ]);
+            llmAdapters = new Map([["default", llmAdapter]]);
+            agent.__init(llmAdapters, {});
+
+            const history: LLMMessage[] = [];
+            for (let i = 0; i < 6; i++) {
+                history.push({ role: i % 2 === 0 ? "user" : "assistant", content: `Message ${i}` });
+            }
+
+            // Should not throw, execution continues with original messages
+            const result = await (agent as any).execute({ message: "New", history }).result;
+
+            expect(result.message).toBe("Response");
+        });
+    });
+});