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

package/spec/FlinkAgent.spec.ts

@@ -0,0 +1,681 @@
+import { z } from "zod";
+import { FlinkAgent } from "../src/ai/FlinkAgent";
+import { FlinkToolProps } from "../src/ai/FlinkTool";
+import { LLMAdapter } from "../src/ai/LLMAdapter";
+import { ToolExecutor } from "../src/ai/ToolExecutor";
+import { FlinkContext } from "../src/FlinkContext";
+import { createStreamingMock } from "./testHelpers";
+
+describe("FlinkAgent", () => {
+    let mockCtx: FlinkContext;
+    let mockLLMAdapter: LLMAdapter;
+
+    beforeEach(() => {
+        mockCtx = {
+            repos: {},
+            plugins: {},
+        };
+
+        // Mock LLM Adapter
+        mockLLMAdapter = createStreamingMock([
+            {
+                textContent: "Test response",
+                toolCalls: [],
+                usage: { inputTokens: 10, outputTokens: 20 },
+                stopReason: "end_turn" as const,
+            },
+        ]);
+    });
+
+    describe("withUser binding", () => {
+        class TestAgent extends FlinkAgent<FlinkContext> {
+            id = "test-agent";
+            description = "Test agent";
+            instructions() { return "Test instructions"; }
+            tools: string[] = [];
+
+            async query(message: string) {
+                const response = this.execute({ message });
+                return await response.result;
+            }
+
+            setContext(ctx: FlinkContext) {
+                (this as any).ctx = ctx;
+            }
+        }
+
+        it("should bind user to agent instance", () => {
+            const agent = new TestAgent();
+            agent.setContext(mockCtx);
+
+            const user = { id: "123", permissions: ["admin"] };
+            const boundAgent = agent.withUser(user);
+
+            // Should return a new instance
+            expect(boundAgent).not.toBe(agent);
+            // Should have bound user
+            expect((boundAgent as any)._boundUser).toBe(user);
+        });
+
+        it("should use bound user for agent permission checks", async () => {
+            class PermissionedAgent extends TestAgent {
+                permissions = "admin";
+            }
+
+            const agent = new PermissionedAgent();
+            agent.setContext(mockCtx);
+            agent.__init(new Map([["default", mockLLMAdapter]]), {});
+
+            const userWithoutPermission = { permissions: [] };
+            const userWithPermission = { permissions: ["admin"] };
+
+            // Should fail without permission
+            const boundAgentFail = agent.withUser(userWithoutPermission);
+            await expectAsync(boundAgentFail.query("test")).toBeRejected();
+
+            // Should succeed with permission
+            const boundAgentSuccess = agent.withUser(userWithPermission);
+            const result = await boundAgentSuccess.query("test");
+            expect(result.message).toBeDefined();
+        });
+
+        it("should pass bound user to tools", async () => {
+            const toolSpy = jasmine.createSpy("toolFn").and.returnValue(Promise.resolve({ success: true, data: { result: "ok" } }));
+
+            const toolProps: FlinkToolProps = {
+                id: "test_tool",
+                description: "Test tool",
+                inputSchema: z.object({}),
+            };
+
+            const toolExecutor = new ToolExecutor(toolProps, toolSpy as any, mockCtx);
+            const toolsMap = new Map([["test_tool", toolExecutor]]);
+
+            // Mock LLM to call tool
+            const toolMockAdapter = createStreamingMock([
+                {
+                    textContent: undefined,
+                    toolCalls: [{ id: "1", name: "test_tool", input: {} }],
+                    usage: { inputTokens: 10, outputTokens: 20 },
+                    stopReason: "tool_use" as const,
+                },
+                {
+                    textContent: "Done",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 20 },
+                    stopReason: "end_turn" as const,
+                },
+            ]);
+
+            class ToolAgent extends TestAgent {
+                tools: string[] = ["test_tool"];
+            }
+
+            const agent = new ToolAgent();
+            agent.setContext(mockCtx);
+            agent.__init(new Map([["default", toolMockAdapter]]), Object.fromEntries(toolsMap));
+
+            const user = { id: "123", permissions: ["admin"] };
+            const boundAgent = agent.withUser(user);
+
+            await boundAgent.query("test");
+
+            // Tool should have been called with the bound user
+            expect(toolSpy).toHaveBeenCalled();
+            const toolCall = (toolSpy as jasmine.Spy).calls.mostRecent();
+            expect(toolCall.args[0].user).toBe(user);
+        });
+    });
+
+    describe("Tool permission filtering", () => {
+        class TestAgent extends FlinkAgent<FlinkContext> {
+            id = "test-agent";
+            description = "Test agent";
+            instructions() { return "Test instructions"; }
+            tools: string[] = ["public_tool", "admin_tool"];
+
+            async query(message: string) {
+                const response = this.execute({ message });
+                return await response.result;
+            }
+
+            setContext(ctx: FlinkContext) {
+                (this as any).ctx = ctx;
+            }
+        }
+
+        it("should filter out tools user doesn't have permission for", async () => {
+            // Create public tool (no permissions)
+            const publicToolProps: FlinkToolProps = {
+                id: "public_tool",
+                description: "Public tool",
+                inputSchema: z.object({}),
+            };
+            const publicToolFn = async () => ({ success: true as const, data: {} });
+            const publicTool = new ToolExecutor(publicToolProps, publicToolFn, mockCtx);
+
+            // Create admin tool (requires admin permission)
+            const adminToolProps: FlinkToolProps = {
+                id: "admin_tool",
+                description: "Admin tool",
+                inputSchema: z.object({}),
+                permissions: "admin",
+            };
+            const adminToolFn = async () => ({ success: true as const, data: {} });
+            const adminTool = new ToolExecutor(adminToolProps, adminToolFn, mockCtx);
+
+            const toolsMap = new Map([
+                ["public_tool", publicTool],
+                ["admin_tool", adminTool],
+            ]);
+
+            const agent = new TestAgent();
+            agent.setContext(mockCtx);
+            agent.__init(new Map([["default", mockLLMAdapter]]), Object.fromEntries(toolsMap));
+
+            // User without admin permission
+            const regularUser = { permissions: [] };
+            const boundAgent = agent.withUser(regularUser);
+
+            await boundAgent.query("test");
+
+            // LLM should have been called with only the public tool via stream
+            const llmCall = (mockLLMAdapter.stream as jasmine.Spy).calls.mostRecent();
+            const toolsProvided = llmCall.args[0].tools;
+
+            expect(toolsProvided.length).toBe(1);
+            expect(toolsProvided[0].name).toBe("public_tool");
+        });
+
+        it("should show all tools to admin users", async () => {
+            // Create public and admin tools (same as above)
+            const publicToolProps: FlinkToolProps = {
+                id: "public_tool",
+                description: "Public tool",
+                inputSchema: z.object({}),
+            };
+            const publicToolFn = async () => ({ success: true as const, data: {} });
+            const publicTool = new ToolExecutor(publicToolProps, publicToolFn, mockCtx);
+
+            const adminToolProps: FlinkToolProps = {
+                id: "admin_tool",
+                description: "Admin tool",
+                inputSchema: z.object({}),
+                permissions: "admin",
+            };
+            const adminToolFn = async () => ({ success: true as const, data: {} });
+            const adminTool = new ToolExecutor(adminToolProps, adminToolFn, mockCtx);
+
+            const toolsMap = new Map([
+                ["public_tool", publicTool],
+                ["admin_tool", adminTool],
+            ]);
+
+            const agent = new TestAgent();
+            agent.setContext(mockCtx);
+            agent.__init(new Map([["default", mockLLMAdapter]]), Object.fromEntries(toolsMap));
+
+            // User with admin permission
+            const adminUser = { permissions: ["admin"] };
+            const boundAgent = agent.withUser(adminUser);
+
+            await boundAgent.query("test");
+
+            // LLM should have been called with both tools via stream
+            const llmCall = (mockLLMAdapter.stream as jasmine.Spy).calls.mostRecent();
+            const toolsProvided = llmCall.args[0].tools;
+
+            expect(toolsProvided.length).toBe(2);
+            expect(toolsProvided.map((t: any) => t.name).sort()).toEqual(["admin_tool", "public_tool"]);
+        });
+    });
+
+    describe("Tool references (FlinkToolFile imports)", () => {
+        it("should support tool file references in addition to string IDs", () => {
+            // Create mock tool file (same structure as real tool files)
+            const mockToolFile = {
+                Tool: {
+                    id: "mock-tool",
+                    description: "Mock tool for testing",
+                    inputSchema: z.object({ input: z.string() }),
+                },
+                default: async () => ({ success: true as const, data: "ok" }),
+            };
+
+            // Create tool executor from the mock file
+            const toolExecutor = new ToolExecutor(mockToolFile.Tool, mockToolFile.default as any, mockCtx);
+
+            // Set up tools in context (not as Map, but as plain object for _tools)
+            // Agent using tool reference instead of string ID
+            class TestAgentWithToolRef extends FlinkAgent<FlinkContext> {
+                id = "test-agent";
+                description = "Test agent";
+                instructions() { return "Test instructions"; }
+                tools = [mockToolFile]; // Direct tool file reference!
+
+                async query(message: string) {
+                    const response = this.execute({ message });
+                    return await response.result;
+                }
+            }
+
+            const agent = new TestAgentWithToolRef();
+            (agent as any).ctx = mockCtx;
+            agent.__init(new Map([["default", mockLLMAdapter]]), { "mock-tool": toolExecutor });
+
+            // Verify resolveTools() correctly extracts tool ID from FlinkToolFile
+            const resolvedTools = (agent as any).resolveTools();
+            expect(resolvedTools.has("mock-tool")).toBe(true);
+            expect(resolvedTools.get("mock-tool")).toBe(toolExecutor);
+        });
+
+        it("should support mixed string IDs and tool file references", () => {
+            // Create two tools: one referenced by string, one by file reference
+            const mockToolFile = {
+                Tool: {
+                    id: "tool-by-reference",
+                    description: "Tool referenced by import",
+                    inputSchema: z.object({}),
+                },
+                default: async () => ({ success: true as const, data: "ok" }),
+            };
+
+            const toolByRefExecutor = new ToolExecutor(mockToolFile.Tool, mockToolFile.default as any, mockCtx);
+            const toolByStringExecutor = new ToolExecutor(
+                { id: "tool-by-string", description: "Tool by string", inputSchema: z.object({}) },
+                async () => ({ success: true as const, data: "ok" }),
+                mockCtx
+            );
+
+            class MixedToolsAgent extends FlinkAgent<FlinkContext> {
+                id = "test-agent";
+                description = "Test agent";
+                instructions() { return "Test instructions"; }
+                tools = [
+                    mockToolFile, // File reference
+                    "tool-by-string", // String ID
+                ];
+            }
+
+            const agent = new MixedToolsAgent();
+            (agent as any).ctx = mockCtx;
+            agent.__init(new Map([["default", mockLLMAdapter]]), {
+                "tool-by-reference": toolByRefExecutor,
+                "tool-by-string": toolByStringExecutor,
+            });
+
+            // Both tools should be resolved correctly
+            const resolvedTools = (agent as any).resolveTools();
+            expect(resolvedTools.size).toBe(2);
+            expect(resolvedTools.has("tool-by-reference")).toBe(true);
+            expect(resolvedTools.has("tool-by-string")).toBe(true);
+        });
+    });
+
+    describe("Dynamic Instructions via Callback", () => {
+        it("should support static string instructions (backwards compatibility)", async () => {
+            class StaticAgent extends FlinkAgent<FlinkContext> {
+                id = "test-agent";
+                description = "Static agent";
+                instructions() { return "Static instructions"; }
+                tools: string[] = [];
+
+                async query(message: string) {
+                    const response = this.execute({ message });
+                    return await response.result;
+                }
+            }
+
+            const agent = new StaticAgent();
+            (agent as any).ctx = mockCtx;
+            agent.__init(new Map([["default", mockLLMAdapter]]), {});
+
+            const result = await agent.query("Hello");
+
+            expect(result.message).toBe("Test response");
+            expect(mockLLMAdapter.stream).toHaveBeenCalled();
+            const callArgs = (mockLLMAdapter.stream as jasmine.Spy).calls.first().args[0];
+            expect(callArgs.instructions).toBe("Static instructions");
+        });
+
+        it("should support synchronous callback", async () => {
+            class DynamicAgent extends FlinkAgent<FlinkContext> {
+                id = "test-agent";
+                description = "Dynamic agent";
+                instructions(ctx: FlinkContext, agentContext: any) {
+                    return `You are a support agent for user ${agentContext.user?.name || "unknown"}`;
+                }
+                tools: string[] = [];
+
+                async query(message: string) {
+                    const response = this.execute({ message });
+                    return await response.result;
+                }
+            }
+
+            const agent = new DynamicAgent();
+            (agent as any).ctx = mockCtx;
+            agent.__init(new Map([["default", mockLLMAdapter]]), {});
+
+            const boundAgent = agent.withUser({ id: "123", name: "John Doe" });
+            const result = await boundAgent.query("Hello");
+
+            expect(result.message).toBe("Test response");
+            expect(mockLLMAdapter.stream).toHaveBeenCalled();
+            const callArgs = (mockLLMAdapter.stream as jasmine.Spy).calls.first().args[0];
+            expect(callArgs.instructions).toBe("You are a support agent for user John Doe");
+        });
+
+        it("should support async callback with database lookup", async () => {
+            const mockUser = { id: "123", name: "Jane Smith", tier: "premium" };
+            const mockRepoCtx = {
+                repos: {
+                    userRepo: {
+                        getById: jasmine.createSpy("getById").and.returnValue(Promise.resolve(mockUser)),
+                    },
+                },
+                plugins: {},
+            };
+
+            class DatabaseAgent extends FlinkAgent<any> {
+                id = "test-agent";
+                description = "Database agent";
+                async instructions(ctx: any, agentContext: any) {
+                    if (!agentContext.user) {
+                        return "You are a support agent.";
+                    }
+                    const profile = await ctx.repos.userRepo.getById(agentContext.user.id);
+                    return `You are a support agent. Customer: ${profile.name}, Tier: ${profile.tier}`;
+                }
+                tools: string[] = [];
+
+                async query(message: string) {
+                    const response = this.execute({ message });
+                    return await response.result;
+                }
+            }
+
+            const agent = new DatabaseAgent();
+            (agent as any).ctx = mockRepoCtx;
+            agent.__init(new Map([["default", mockLLMAdapter]]), {});
+
+            const boundAgent = agent.withUser({ id: "123" });
+            const result = await boundAgent.query("Hello");
+
+            expect(mockRepoCtx.repos.userRepo.getById).toHaveBeenCalledWith("123");
+            expect(result.message).toBe("Test response");
+            expect(mockLLMAdapter.stream).toHaveBeenCalled();
+            const callArgs = (mockLLMAdapter.stream as jasmine.Spy).calls.first().args[0];
+            expect(callArgs.instructions).toBe("You are a support agent. Customer: Jane Smith, Tier: premium");
+        });
+
+        it("should work with runner caching", async () => {
+            let callCount = 0;
+
+            class CachedAgent extends FlinkAgent<FlinkContext> {
+                id = "test-agent";
+                description = "Cached agent";
+                instructions(ctx: FlinkContext, agentContext: any) {
+                    callCount++;
+                    return `Instructions for user ${agentContext.user?.name || "unknown"}`;
+                }
+                tools: string[] = [];
+
+                async query(message: string) {
+                    const response = this.execute({ message });
+                    return await response.result;
+                }
+            }
+
+            const agent = new CachedAgent();
+            (agent as any).ctx = mockCtx;
+            agent.__init(new Map([["default", mockLLMAdapter]]), {});
+
+            // First execution with user1
+            const user1 = { id: "1", name: "Alice" };
+            const boundAgent1 = agent.withUser(user1);
+            await boundAgent1.query("Hello");
+
+            expect(callCount).toBe(1);
+            expect(mockLLMAdapter.stream).toHaveBeenCalled();
+            let callArgs = (mockLLMAdapter.stream as jasmine.Spy).calls.first().args[0];
+            expect(callArgs.instructions).toBe("Instructions for user Alice");
+
+            // Reset mock
+            (mockLLMAdapter.stream as jasmine.Spy).calls.reset();
+
+            // Second execution with different user (new bound agent = new runner)
+            const user2 = { id: "2", name: "Bob" };
+            const boundAgent2 = agent.withUser(user2);
+            await boundAgent2.query("Hello");
+
+            expect(callCount).toBe(2); // Should increment because new runner instance
+            expect(mockLLMAdapter.stream).toHaveBeenCalled();
+            callArgs = (mockLLMAdapter.stream as jasmine.Spy).calls.first().args[0];
+            expect(callArgs.instructions).toBe("Instructions for user Bob");
+        });
+
+        it("should handle callback with conversationId", async () => {
+            let capturedContext: any;
+
+            class ConversationAgent extends FlinkAgent<FlinkContext> {
+                id = "test-agent";
+                description = "Conversation agent";
+                instructions(ctx: FlinkContext, agentContext: any) {
+                    capturedContext = agentContext;
+                    return `Conversation: ${agentContext.conversationId || "new"}`;
+                }
+                tools: string[] = [];
+
+                async query(message: string, conversationId?: string) {
+                    const response = this.execute({ message, conversationId });
+                    return await response.result;
+                }
+            }
+
+            const agent = new ConversationAgent();
+            (agent as any).ctx = mockCtx;
+            agent.__init(new Map([["default", mockLLMAdapter]]), {});
+
+            await agent.query("Hello", "conv-456");
+
+            expect(capturedContext.conversationId).toBe("conv-456");
+            expect(mockLLMAdapter.stream).toHaveBeenCalled();
+            const callArgs = (mockLLMAdapter.stream as jasmine.Spy).calls.first().args[0];
+            expect(callArgs.instructions).toBe("Conversation: conv-456");
+        });
+
+        it("should throw error if callback throws", async () => {
+            class ErrorAgent extends FlinkAgent<FlinkContext> {
+                id = "test-agent";
+                description = "Error agent";
+                instructions(): string {
+                    throw new Error("Database unavailable");
+                }
+                tools: string[] = [];
+
+                async query(message: string) {
+                    const response = this.execute({ message });
+                    return await response.result;
+                }
+            }
+
+            const agent = new ErrorAgent();
+            (agent as any).ctx = mockCtx;
+            agent.__init(new Map([["default", mockLLMAdapter]]), {});
+
+            await expectAsync(agent.query("Hello")).toBeRejectedWithError(/Failed to resolve instructions/);
+        });
+    });
+
+    describe("withConversationContext binding", () => {
+        interface TestConversationContext {
+            sessionId: string;
+            featureFlags: Record<string, boolean>;
+            requestId?: string;
+        }
+
+        class TestAgent extends FlinkAgent<FlinkContext, TestConversationContext> {
+            id = "test-agent";
+            description = "Test agent";
+            instructions() { return "Test instructions"; }
+            tools: string[] = [];
+
+            async query(message: string) {
+                const response = this.execute({ message });
+                return await response.result;
+            }
+
+            setContext(ctx: FlinkContext) {
+                (this as any).ctx = ctx;
+            }
+        }
+
+        it("should bind conversation context to agent instance", () => {
+            const agent = new TestAgent();
+            agent.setContext(mockCtx);
+
+            const conversationCtx: TestConversationContext = {
+                sessionId: "session-123",
+                featureFlags: { enableNewUI: true },
+                requestId: "req-456",
+            };
+            const boundAgent = agent.withConversationContext(conversationCtx);
+
+            // Should return a new instance
+            expect(boundAgent).not.toBe(agent);
+            // Should have bound conversation context
+            expect((boundAgent as any)._boundConversationContext).toBe(conversationCtx);
+        });
+
+        it("should pass bound conversation context to tools", async () => {
+            const toolSpy = jasmine.createSpy("toolFn").and.returnValue(Promise.resolve({ success: true, data: { result: "ok" } }));
+
+            const toolProps: FlinkToolProps = {
+                id: "test_tool",
+                description: "Test tool",
+                inputSchema: z.object({}),
+            };
+
+            const toolExecutor = new ToolExecutor(toolProps, toolSpy as any, mockCtx);
+            const toolsMap = new Map([["test_tool", toolExecutor]]);
+
+            // Mock LLM to call tool
+            const toolMockAdapter = createStreamingMock([
+                {
+                    textContent: undefined,
+                    toolCalls: [{ id: "1", name: "test_tool", input: {} }],
+                    usage: { inputTokens: 10, outputTokens: 20 },
+                    stopReason: "tool_use" as const,
+                },
+                {
+                    textContent: "Done",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 20 },
+                    stopReason: "end_turn" as const,
+                },
+            ]);
+
+            class ToolAgent extends TestAgent {
+                tools = ["test_tool"];
+            }
+
+            const agent = new ToolAgent();
+            agent.setContext(mockCtx);
+            agent.__init(new Map([["default", toolMockAdapter]]), Object.fromEntries(toolsMap));
+
+            const conversationCtx: TestConversationContext = {
+                sessionId: "session-123",
+                featureFlags: { enableNewUI: true },
+            };
+
+            const boundAgent = agent.withConversationContext(conversationCtx);
+            await boundAgent.query("test");
+
+            // Verify tool was called with conversation context
+            expect(toolSpy).toHaveBeenCalled();
+            const toolCallArgs = toolSpy.calls.first().args[0];
+            expect(toolCallArgs.conversationCtx).toBe(conversationCtx);
+            expect(toolCallArgs.conversationCtx.sessionId).toBe("session-123");
+            expect(toolCallArgs.conversationCtx.featureFlags.enableNewUI).toBe(true);
+        });
+
+        it("should support fluent chaining with other binding methods", () => {
+            const agent = new TestAgent();
+            agent.setContext(mockCtx);
+            agent.__init(new Map([["default", mockLLMAdapter]]), {});
+
+            const user = { id: "123", permissions: ["admin"] };
+            const conversationCtx: TestConversationContext = {
+                sessionId: "session-123",
+                featureFlags: { enableNewUI: true },
+            };
+
+            // Chain withUser and withConversationContext
+            const boundAgent = agent.withUser(user).withConversationContext(conversationCtx);
+
+            expect((boundAgent as any)._boundUser).toBe(user);
+            expect((boundAgent as any)._boundConversationContext).toBe(conversationCtx);
+        });
+
+        it("should allow explicit conversationContext in execute to override bound", async () => {
+            const toolSpy = jasmine.createSpy("toolFn").and.returnValue(Promise.resolve({ success: true, data: { result: "ok" } }));
+
+            const toolProps: FlinkToolProps = {
+                id: "test_tool",
+                description: "Test tool",
+                inputSchema: z.object({}),
+            };
+
+            const toolExecutor = new ToolExecutor(toolProps, toolSpy as any, mockCtx);
+            const toolsMap = new Map([["test_tool", toolExecutor]]);
+
+            const toolMockAdapter = createStreamingMock([
+                {
+                    textContent: undefined,
+                    toolCalls: [{ id: "1", name: "test_tool", input: {} }],
+                    usage: { inputTokens: 10, outputTokens: 20 },
+                    stopReason: "tool_use" as const,
+                },
+                {
+                    textContent: "Done",
+                    toolCalls: [],
+                    usage: { inputTokens: 10, outputTokens: 20 },
+                    stopReason: "end_turn" as const,
+                },
+            ]);
+
+            class ToolAgent extends TestAgent {
+                tools = ["test_tool"];
+
+                async queryWithContext(message: string, ctx: TestConversationContext) {
+                    const response = this.execute({ message, conversationContext: ctx });
+                    return await response.result;
+                }
+            }
+
+            const agent = new ToolAgent();
+            agent.setContext(mockCtx);
+            agent.__init(new Map([["default", toolMockAdapter]]), Object.fromEntries(toolsMap));
+
+            const boundCtx: TestConversationContext = {
+                sessionId: "bound-session",
+                featureFlags: {},
+            };
+
+            const explicitCtx: TestConversationContext = {
+                sessionId: "explicit-session",
+                featureFlags: { newFeature: true },
+            };
+
+            const boundAgent = agent.withConversationContext(boundCtx);
+            await boundAgent.queryWithContext("test", explicitCtx);
+
+            // Verify tool was called with explicit context (override)
+            expect(toolSpy).toHaveBeenCalled();
+            const toolCallArgs = toolSpy.calls.first().args[0];
+            expect(toolCallArgs.conversationCtx.sessionId).toBe("explicit-session");
+            expect(toolCallArgs.conversationCtx.featureFlags.newFeature).toBe(true);
+        });
+    });
+});