fix(KB-209): fix infinite loop in useMultiAgentLogs hook
- Fix infinite loop in useMultiAgentLogs hook when handling task updates - Correct cleanup logic to only close connections for removed tasks - Add global EventSource mock infrastructure to vitest.setup.ts - Refactor useMultiAgentLogs tests to use shared MockEventSource pattern
This commit is contained in:
@@ -2,6 +2,7 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||
import { renderHook, act, waitFor } from "@testing-library/react";
|
||||
import { MAX_LOG_ENTRIES, useMultiAgentLogs } from "../useMultiAgentLogs";
|
||||
import { fetchAgentLogs } from "../../api";
|
||||
import { MockEventSource } from "../../../vitest.setup";
|
||||
|
||||
// Mock the api module
|
||||
vi.mock("../../api", () => ({
|
||||
@@ -10,68 +11,27 @@ vi.mock("../../api", () => ({
|
||||
|
||||
const mockFetchAgentLogs = vi.mocked(fetchAgentLogs);
|
||||
|
||||
// Mock EventSource - track instances per hook render, not globally
|
||||
class MockEventSource {
|
||||
url: string;
|
||||
listeners: Record<string, ((e: { data: string }) => void)[]> = {};
|
||||
readyState = 0;
|
||||
close = vi.fn(() => {
|
||||
this.readyState = 2;
|
||||
});
|
||||
|
||||
constructor(url: string) {
|
||||
this.url = url;
|
||||
this.readyState = 1;
|
||||
}
|
||||
|
||||
addEventListener(event: string, fn: (e: { data: string }) => void) {
|
||||
if (!this.listeners[event]) this.listeners[event] = [];
|
||||
this.listeners[event].push(fn);
|
||||
}
|
||||
|
||||
removeEventListener(event: string, fn: (e: { data: string }) => void) {
|
||||
this.listeners[event] = (this.listeners[event] || []).filter((listener) => listener !== fn);
|
||||
}
|
||||
|
||||
// Helper to simulate a server event
|
||||
_emit(event: string, data?: unknown) {
|
||||
for (const fn of this.listeners[event] || []) {
|
||||
fn(data === undefined ? ({ } as { data: string }) : { data: JSON.stringify(data) });
|
||||
}
|
||||
}
|
||||
// Helper to get the last connection for a specific task ID
|
||||
function getConnection(taskId: string): MockEventSource | undefined {
|
||||
const url = `/api/tasks/${taskId}/logs/stream`;
|
||||
const matching = MockEventSource.instances.filter((e) => e.url === url);
|
||||
return matching[matching.length - 1];
|
||||
}
|
||||
|
||||
const originalEventSource = globalThis.EventSource;
|
||||
// Helper to get all connections for a task ID
|
||||
function getConnections(taskId: string): MockEventSource[] {
|
||||
const url = `/api/tasks/${taskId}/logs/stream`;
|
||||
return MockEventSource.instances.filter((e) => e.url === url);
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = MockEventSource;
|
||||
mockFetchAgentLogs.mockReset().mockResolvedValue([]);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = originalEventSource;
|
||||
// Clean up is handled by global afterEach in vitest.setup.ts
|
||||
});
|
||||
|
||||
function getActiveConnections(): MockEventSource[] {
|
||||
// Get all MockEventSource instances that haven't been closed
|
||||
// We need to track this ourselves since the mock is recreated each time
|
||||
const allSources: MockEventSource[] = [];
|
||||
|
||||
// Hook into the constructor to track instances
|
||||
const OriginalMock = MockEventSource;
|
||||
const instances: MockEventSource[] = [];
|
||||
|
||||
// Override to capture instances
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
return instances;
|
||||
}
|
||||
|
||||
describe("useMultiAgentLogs", () => {
|
||||
it("initializes with empty entries for all provided task IDs", () => {
|
||||
const { result } = renderHook(() => useMultiAgentLogs(["KB-001", "KB-002"]));
|
||||
@@ -119,21 +79,11 @@ describe("useMultiAgentLogs", () => {
|
||||
it("opens SSE EventSource for each task ID", async () => {
|
||||
mockFetchAgentLogs.mockResolvedValue([]);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
|
||||
// Override to capture instances
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
renderHook(() => useMultiAgentLogs(["KB-001", "KB-002"]));
|
||||
|
||||
await waitFor(() => {
|
||||
// Filter to unique URLs (Strict Mode may create duplicates)
|
||||
const urls = [...new Set(instances.map((es) => es.url))];
|
||||
const urls = [...new Set(MockEventSource.instances.map((es) => es.url))];
|
||||
expect(urls).toContain("/api/tasks/KB-001/logs/stream");
|
||||
expect(urls).toContain("/api/tasks/KB-002/logs/stream");
|
||||
});
|
||||
@@ -146,21 +96,13 @@ describe("useMultiAgentLogs", () => {
|
||||
// Use mockResolvedValue (not Once) to handle Strict Mode double-run
|
||||
mockFetchAgentLogs.mockResolvedValue(historical);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
const { result } = renderHook(() => useMultiAgentLogs(["KB-001"]));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(result.current["KB-001"].entries).toHaveLength(1);
|
||||
});
|
||||
|
||||
const es = instances.find((e) => e.url.includes("KB-001"));
|
||||
const es = getConnection("KB-001");
|
||||
expect(es).toBeDefined();
|
||||
|
||||
act(() => {
|
||||
@@ -182,24 +124,16 @@ describe("useMultiAgentLogs", () => {
|
||||
it("closes all SSE connections on unmount (memory leak prevention)", async () => {
|
||||
mockFetchAgentLogs.mockResolvedValue([]);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
const { unmount } = renderHook(() => useMultiAgentLogs(["KB-001", "KB-002"]));
|
||||
|
||||
// Wait for connections to be established
|
||||
await waitFor(() => {
|
||||
expect(instances.length).toBeGreaterThanOrEqual(2);
|
||||
expect(MockEventSource.instances.length).toBeGreaterThanOrEqual(2);
|
||||
});
|
||||
|
||||
// Get unique instances by URL (handling Strict Mode duplicates)
|
||||
const uniqueByUrl = new Map<string, MockEventSource>();
|
||||
for (const es of instances) {
|
||||
for (const es of MockEventSource.instances) {
|
||||
if (!uniqueByUrl.has(es.url) || !es.close.mock?.calls?.length) {
|
||||
uniqueByUrl.set(es.url, es);
|
||||
}
|
||||
@@ -217,66 +151,43 @@ describe("useMultiAgentLogs", () => {
|
||||
it("closes specific connection when task ID removed from array", async () => {
|
||||
mockFetchAgentLogs.mockResolvedValue([]);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ taskIds }: { taskIds: string[] }) => useMultiAgentLogs(taskIds),
|
||||
{ initialProps: { taskIds: ["KB-001", "KB-002"] } },
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(instances.length).toBeGreaterThanOrEqual(2);
|
||||
expect(MockEventSource.instances.length).toBeGreaterThanOrEqual(2);
|
||||
});
|
||||
|
||||
// Get the last connection for each URL
|
||||
const getConnection = (taskId: string) => {
|
||||
const url = `/api/tasks/${taskId}/logs/stream`;
|
||||
const matching = instances.filter((e) => e.url === url);
|
||||
return matching[matching.length - 1];
|
||||
};
|
||||
|
||||
const es1 = getConnection("KB-001");
|
||||
const es2 = getConnection("KB-002");
|
||||
|
||||
rerender({ taskIds: ["KB-001"] });
|
||||
|
||||
await waitFor(() => {
|
||||
expect(es2.close).toHaveBeenCalled();
|
||||
expect(es2!.close).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
expect(es1.close).not.toHaveBeenCalled();
|
||||
expect(es1!.close).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("opens new connection when task ID added to array", async () => {
|
||||
mockFetchAgentLogs.mockResolvedValue([]);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ taskIds }: { taskIds: string[] }) => useMultiAgentLogs(taskIds),
|
||||
{ initialProps: { taskIds: ["KB-001"] } },
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(instances.length).toBeGreaterThanOrEqual(1);
|
||||
expect(MockEventSource.instances.length).toBeGreaterThanOrEqual(1);
|
||||
});
|
||||
|
||||
rerender({ taskIds: ["KB-001", "KB-002"] });
|
||||
|
||||
await waitFor(() => {
|
||||
const urls = [...new Set(instances.map((es) => es.url))];
|
||||
const urls = [...new Set(MockEventSource.instances.map((es) => es.url))];
|
||||
expect(urls).toContain("/api/tasks/KB-002/logs/stream");
|
||||
});
|
||||
});
|
||||
@@ -325,28 +236,24 @@ describe("useMultiAgentLogs", () => {
|
||||
}),
|
||||
);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ taskIds }: { taskIds: string[] }) => useMultiAgentLogs(taskIds),
|
||||
{ initialProps: { taskIds: ["KB-001"] } },
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(instances.filter((es) => es.url === "/api/tasks/KB-001/logs/stream")).toHaveLength(1);
|
||||
// Allow for Strict Mode double-rendering
|
||||
expect(getConnections("KB-001").length).toBeGreaterThanOrEqual(1);
|
||||
});
|
||||
|
||||
const initialCount = getConnections("KB-001").length;
|
||||
|
||||
rerender({ taskIds: ["KB-001"] });
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
expect(instances.filter((es) => es.url === "/api/tasks/KB-001/logs/stream")).toHaveLength(1);
|
||||
// Should not create additional connections on rerender with same IDs
|
||||
expect(getConnections("KB-001").length).toBe(initialCount);
|
||||
|
||||
resolveFetch?.([]);
|
||||
|
||||
@@ -358,27 +265,20 @@ describe("useMultiAgentLogs", () => {
|
||||
it("closes a task connection when its stream emits an error", async () => {
|
||||
mockFetchAgentLogs.mockResolvedValue([]);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
renderHook(() => useMultiAgentLogs(["KB-001"]));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(instances).toHaveLength(1);
|
||||
expect(MockEventSource.instances.length).toBeGreaterThanOrEqual(1);
|
||||
});
|
||||
|
||||
const es = instances[0];
|
||||
const es = getConnection("KB-001");
|
||||
expect(es).toBeDefined();
|
||||
|
||||
act(() => {
|
||||
es._emit("error");
|
||||
es!._emit("error");
|
||||
});
|
||||
|
||||
expect(es.close).toHaveBeenCalledTimes(1);
|
||||
expect(es!.close).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("truncates oversized historical logs per task to the most recent entries", async () => {
|
||||
@@ -409,22 +309,17 @@ describe("useMultiAgentLogs", () => {
|
||||
}),
|
||||
);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
const { result } = renderHook(() => useMultiAgentLogs(["KB-001"]));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(instances).toHaveLength(1);
|
||||
expect(MockEventSource.instances.length).toBeGreaterThanOrEqual(1);
|
||||
});
|
||||
|
||||
const es = getConnection("KB-001");
|
||||
expect(es).toBeDefined();
|
||||
|
||||
act(() => {
|
||||
instances[0]._emit("agent:log", {
|
||||
es!._emit("agent:log", {
|
||||
timestamp: "2026-01-01T00:01:00Z",
|
||||
taskId: "KB-001",
|
||||
text: "live-before-history",
|
||||
@@ -454,23 +349,18 @@ describe("useMultiAgentLogs", () => {
|
||||
it("truncates live SSE entries per task to the most recent entries", async () => {
|
||||
mockFetchAgentLogs.mockResolvedValue([]);
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
const { result } = renderHook(() => useMultiAgentLogs(["KB-001"]));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(instances).toHaveLength(1);
|
||||
expect(MockEventSource.instances.length).toBeGreaterThanOrEqual(1);
|
||||
});
|
||||
|
||||
const es = getConnection("KB-001");
|
||||
expect(es).toBeDefined();
|
||||
|
||||
act(() => {
|
||||
for (let index = 0; index < MAX_LOG_ENTRIES + 15; index++) {
|
||||
instances[0]._emit("agent:log", {
|
||||
es!._emit("agent:log", {
|
||||
timestamp: `2026-01-01T00:${String(index).padStart(2, "0")}:00Z`,
|
||||
taskId: "KB-001",
|
||||
text: `live-${index}`,
|
||||
@@ -501,14 +391,6 @@ describe("useMultiAgentLogs", () => {
|
||||
return Promise.resolve([]);
|
||||
});
|
||||
|
||||
const instances: MockEventSource[] = [];
|
||||
(globalThis as unknown as Record<string, unknown>).EventSource = class extends MockEventSource {
|
||||
constructor(url: string) {
|
||||
super(url);
|
||||
instances.push(this);
|
||||
}
|
||||
};
|
||||
|
||||
const { result } = renderHook(() => useMultiAgentLogs(["KB-001", "KB-002"]));
|
||||
|
||||
await waitFor(() => {
|
||||
@@ -516,20 +398,13 @@ describe("useMultiAgentLogs", () => {
|
||||
expect(result.current["KB-002"].entries).toHaveLength(1);
|
||||
});
|
||||
|
||||
// Get the last connection for each URL
|
||||
const getConnection = (taskId: string) => {
|
||||
const url = `/api/tasks/${taskId}/logs/stream`;
|
||||
const matching = instances.filter((e) => e.url === url);
|
||||
return matching[matching.length - 1];
|
||||
};
|
||||
|
||||
const es1 = getConnection("KB-001");
|
||||
const es2 = getConnection("KB-002");
|
||||
expect(es1).toBeDefined();
|
||||
expect(es2).toBeDefined();
|
||||
|
||||
act(() => {
|
||||
es1._emit("agent:log", {
|
||||
es1!._emit("agent:log", {
|
||||
timestamp: "2026-01-01T00:01:00Z",
|
||||
taskId: "KB-001",
|
||||
text: "task1-new",
|
||||
@@ -543,7 +418,7 @@ describe("useMultiAgentLogs", () => {
|
||||
});
|
||||
|
||||
act(() => {
|
||||
es2._emit("agent:log", {
|
||||
es2!._emit("agent:log", {
|
||||
timestamp: "2026-01-01T00:01:00Z",
|
||||
taskId: "KB-002",
|
||||
text: "task2-new",
|
||||
|
||||
@@ -59,6 +59,9 @@ export function useMultiAgentLogs(taskIds: string[]): LogStateMap {
|
||||
};
|
||||
}, []);
|
||||
|
||||
// Stable comparison of task IDs to prevent effect re-runs on every render
|
||||
const taskIdsKey = taskIds.join(",");
|
||||
|
||||
// Main effect to manage connections
|
||||
useEffect(() => {
|
||||
const currentIds = new Set(taskIds);
|
||||
@@ -66,7 +69,30 @@ export function useMultiAgentLogs(taskIds: string[]): LogStateMap {
|
||||
const initializing = initializingRef.current;
|
||||
const cancelled = cancelledRef.current;
|
||||
|
||||
// Track which task IDs need state initialization (not already in stateMap)
|
||||
const newTaskIds: string[] = [];
|
||||
for (const taskId of taskIds) {
|
||||
if (!stateMap[taskId]) {
|
||||
newTaskIds.push(taskId);
|
||||
}
|
||||
}
|
||||
|
||||
// Only initialize state for new tasks that aren't already in stateMap
|
||||
if (newTaskIds.length > 0) {
|
||||
setStateMap((prev) => {
|
||||
const updates: Record<string, InitState> = {};
|
||||
for (const taskId of newTaskIds) {
|
||||
if (!prev[taskId]) {
|
||||
updates[taskId] = { entries: [], loading: true };
|
||||
}
|
||||
}
|
||||
if (Object.keys(updates).length === 0) return prev;
|
||||
return { ...prev, ...updates };
|
||||
});
|
||||
}
|
||||
|
||||
// Close connections for tasks no longer in the list
|
||||
const removedTaskIds: string[] = [];
|
||||
for (const [taskId, es] of Object.entries(sources)) {
|
||||
if (!currentIds.has(taskId)) {
|
||||
cancelled[taskId] = true;
|
||||
@@ -75,14 +101,30 @@ export function useMultiAgentLogs(taskIds: string[]): LogStateMap {
|
||||
initializing.delete(taskId);
|
||||
delete cancelled[taskId];
|
||||
delete pendingLiveEntriesRef.current[taskId];
|
||||
|
||||
// Remove state for disconnected task
|
||||
setStateMap((prev) => {
|
||||
const { [taskId]: _, ...rest } = prev;
|
||||
return rest;
|
||||
});
|
||||
removedTaskIds.push(taskId);
|
||||
}
|
||||
}
|
||||
|
||||
// Only remove state for disconnected tasks if there are any
|
||||
if (removedTaskIds.length > 0) {
|
||||
setStateMap((prev) => {
|
||||
let hasChanges = false;
|
||||
for (const taskId of removedTaskIds) {
|
||||
if (taskId in prev) {
|
||||
hasChanges = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!hasChanges) return prev;
|
||||
const newState: Record<string, InitState> = {};
|
||||
for (const [id, state] of Object.entries(prev)) {
|
||||
if (!removedTaskIds.includes(id)) {
|
||||
newState[id] = state;
|
||||
}
|
||||
}
|
||||
return newState;
|
||||
});
|
||||
}
|
||||
|
||||
// Mark removed pending initializations as cancelled even if EventSource not created yet
|
||||
for (const taskId of Object.keys(cancelled)) {
|
||||
@@ -93,14 +135,8 @@ export function useMultiAgentLogs(taskIds: string[]): LogStateMap {
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize state and connections for current tasks
|
||||
// Initialize connections for current tasks
|
||||
for (const taskId of taskIds) {
|
||||
// Initialize state if not present
|
||||
setStateMap((prev) => {
|
||||
if (prev[taskId]) return prev;
|
||||
return { ...prev, [taskId]: { entries: [], loading: true } };
|
||||
});
|
||||
|
||||
// Skip if already connected or currently initializing
|
||||
if (sources[taskId] || initializing.has(taskId)) continue;
|
||||
|
||||
@@ -179,22 +215,28 @@ export function useMultiAgentLogs(taskIds: string[]): LogStateMap {
|
||||
initializingRef.current.delete(taskId);
|
||||
});
|
||||
}
|
||||
|
||||
// Update previous task IDs ref for cleanup comparison
|
||||
const initialTaskIds = [...taskIds];
|
||||
|
||||
// Cleanup on effect re-run or unmount
|
||||
return () => {
|
||||
for (const taskId of taskIds) {
|
||||
cancelledRef.current[taskId] = true;
|
||||
// Only close connections for tasks that were removed (not in current taskIds)
|
||||
for (const taskId of initialTaskIds) {
|
||||
if (!currentIds.has(taskId)) {
|
||||
cancelledRef.current[taskId] = true;
|
||||
|
||||
const es = sourcesRef.current[taskId];
|
||||
if (es) {
|
||||
es.close();
|
||||
delete sourcesRef.current[taskId];
|
||||
const es = sourcesRef.current[taskId];
|
||||
if (es) {
|
||||
es.close();
|
||||
delete sourcesRef.current[taskId];
|
||||
}
|
||||
|
||||
initializingRef.current.delete(taskId);
|
||||
}
|
||||
|
||||
initializingRef.current.delete(taskId);
|
||||
}
|
||||
};
|
||||
}, [taskIds]);
|
||||
}, [taskIdsKey]); // Use stable string key instead of array reference
|
||||
|
||||
// Close all connections on unmount
|
||||
useEffect(() => {
|
||||
|
||||
Reference in New Issue
Block a user