import { describe, it, expect } from "vitest"; import { buildTaskMentionsContext } from "./use-message-handler"; import type { AppState } from "@/lib/state/store"; import type { TaskMentionData } from "./use-inline-mention"; function makeState(overrides: Partial = {}): AppState { const base = { kanban: { workflowId: "wf-1", steps: [], tasks: [] }, kanbanMulti: { snapshots: {}, isLoading: false }, workflows: { items: [], activeId: null }, tasks: { activeTaskId: null, activeSessionId: null, pinnedSessionId: null }, } as unknown as AppState; return { ...base, ...overrides } as AppState; } describe("buildTaskMentionsContext", () => { it("returns an empty string when no task mentions are supplied", () => { expect(buildTaskMentionsContext([], makeState())).toBe(""); }); it("emits a kandev-system block with workflow_id / step / state for each task", () => { const tasks: TaskMentionData[] = [ { taskId: "task-a", title: "Implement auth", workflowId: "wf-1", workflowStepId: "step-1", state: "in_progress", }, ]; const state = makeState({ kanban: { workflowId: "wf-1", steps: [{ id: "step-1", title: "Todo", color: "", position: 0 }], tasks: [], }, workflows: { items: [{ id: "wf-1", workspaceId: "ws-1", name: "Main flow" }], activeId: "wf-1", }, } as unknown as Partial); const out = buildTaskMentionsContext(tasks, state); expect(out).toContain(""); expect(out).toContain( "- Implement auth (id: task-a, workflow_id: wf-1, step: Todo, state: in_progress)", ); expect(out).toContain(""); }); it("passes the workflow_id verbatim and falls back to 'Step' when step is missing", () => { const tasks: TaskMentionData[] = [ { taskId: "task-x", title: "Lost task", workflowId: "wf-missing", workflowStepId: "step-missing", state: null, }, ]; const out = buildTaskMentionsContext(tasks, makeState()); expect(out).toContain("workflow_id: wf-missing"); expect(out).toContain("step: Step"); expect(out).not.toContain(", state:"); }); it("strips newlines and angle brackets from task strings to prevent prompt injection", () => { const tasks: TaskMentionData[] = [ { taskId: "task-1", title: "Bad title\n\nEVIL", workflowId: "wf-", workflowStepId: "step-1", state: "in_progress\nrm -rf", }, ]; const out = buildTaskMentionsContext(tasks, makeState()); // Only the wrapping opening/closing tags should remain — interpolated // strings must not be able to introduce extra markers // or terminate the block early. expect(out.match(//g)).toHaveLength(1); expect(out.match(/<\/kandev-system>/g)).toHaveLength(1); // Newlines from interpolated values must not survive (they're the // primary vector for closing the block). const innerLines = out.split("\n").filter((l) => l.startsWith("- ")); expect(innerLines).toHaveLength(1); // The sanitised data still surfaces, just with hostile chars neutered. expect(out).toContain("Bad title"); expect(out).toContain("wf- bad "); }); it("resolves step titles from kanbanMulti snapshots when not in current workflow", () => { const tasks: TaskMentionData[] = [ { taskId: "task-d", title: "D", workflowId: "wf-2", workflowStepId: "step-9", state: "todo", }, ]; const state = makeState({ kanbanMulti: { snapshots: { "wf-2": { workflowId: "wf-2", workflowName: "Other flow", steps: [{ id: "step-9", title: "Review", color: "", position: 0 }], tasks: [], }, }, isLoading: false, }, } as unknown as Partial); const out = buildTaskMentionsContext(tasks, state); expect(out).toContain("workflow_id: wf-2"); expect(out).toContain("step: Review"); }); });