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:
gsxdsm
2026-03-30 16:15:17 -07:00
parent cda352c428
commit 554f7d208a
3 changed files with 163 additions and 191 deletions

View File

@@ -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",

View File

@@ -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(() => {