1// Complete implementation of condition-based waiting utilities
2// From: Lace test infrastructure improvements (2025-10-03)
3// Context: Fixed 15 flaky tests by replacing arbitrary timeouts
4 
5import type { ThreadManager } from '~/threads/thread-manager';
6import type { LaceEvent, LaceEventType } from '~/threads/types';
7 
8/**
9 * Wait for a specific event type to appear in thread
10 *
11 * @param threadManager - The thread manager to query
12 * @param threadId - Thread to check for events
13 * @param eventType - Type of event to wait for
14 * @param timeoutMs - Maximum time to wait (default 5000ms)
15 * @returns Promise resolving to the first matching event
16 *
17 * Example:
18 *   await waitForEvent(threadManager, agentThreadId, 'TOOL_RESULT');
19 */
20export function waitForEvent(
21  threadManager: ThreadManager,
22  threadId: string,
23  eventType: LaceEventType,
24  timeoutMs = 5000
25): Promise<LaceEvent> {
26  return new Promise((resolve, reject) => {
27    const startTime = Date.now();
28 
29    const check = () => {
30      const events = threadManager.getEvents(threadId);
31      const event = events.find((e) => e.type === eventType);
32 
33      if (event) {
34        resolve(event);
35      } else if (Date.now() - startTime > timeoutMs) {
36        reject(new Error(`Timeout waiting for ${eventType} event after ${timeoutMs}ms`));
37      } else {
38        setTimeout(check, 10); // Poll every 10ms for efficiency
39      }
40    };
41 
42    check();
43  });
44}
45 
46/**
47 * Wait for a specific number of events of a given type
48 *
49 * @param threadManager - The thread manager to query
50 * @param threadId - Thread to check for events
51 * @param eventType - Type of event to wait for
52 * @param count - Number of events to wait for
53 * @param timeoutMs - Maximum time to wait (default 5000ms)
54 * @returns Promise resolving to all matching events once count is reached
55 *
56 * Example:
57 *   // Wait for 2 AGENT_MESSAGE events (initial response + continuation)
58 *   await waitForEventCount(threadManager, agentThreadId, 'AGENT_MESSAGE', 2);
59 */
60export function waitForEventCount(
61  threadManager: ThreadManager,
62  threadId: string,
63  eventType: LaceEventType,
64  count: number,
65  timeoutMs = 5000
66): Promise<LaceEvent[]> {
67  return new Promise((resolve, reject) => {
68    const startTime = Date.now();
69 
70    const check = () => {
71      const events = threadManager.getEvents(threadId);
72      const matchingEvents = events.filter((e) => e.type === eventType);
73 
74      if (matchingEvents.length >= count) {
75        resolve(matchingEvents);
76      } else if (Date.now() - startTime > timeoutMs) {
77        reject(
78          new Error(
79            `Timeout waiting for ${count} ${eventType} events after ${timeoutMs}ms (got ${matchingEvents.length})`
80          )
81        );
82      } else {
83        setTimeout(check, 10);
84      }
85    };
86 
87    check();
88  });
89}
90 
91/**
92 * Wait for an event matching a custom predicate
93 * Useful when you need to check event data, not just type
94 *
95 * @param threadManager - The thread manager to query
96 * @param threadId - Thread to check for events
97 * @param predicate - Function that returns true when event matches
98 * @param description - Human-readable description for error messages
99 * @param timeoutMs - Maximum time to wait (default 5000ms)
100 * @returns Promise resolving to the first matching event
101 *
102 * Example:
103 *   // Wait for TOOL_RESULT with specific ID
104 *   await waitForEventMatch(
105 *     threadManager,
106 *     agentThreadId,
107 *     (e) => e.type === 'TOOL_RESULT' && e.data.id === 'call_123',
108 *     'TOOL_RESULT with id=call_123'
109 *   );
110 */
111export function waitForEventMatch(
112  threadManager: ThreadManager,
113  threadId: string,
114  predicate: (event: LaceEvent) => boolean,
115  description: string,
116  timeoutMs = 5000
117): Promise<LaceEvent> {
118  return new Promise((resolve, reject) => {
119    const startTime = Date.now();
120 
121    const check = () => {
122      const events = threadManager.getEvents(threadId);
123      const event = events.find(predicate);
124 
125      if (event) {
126        resolve(event);
127      } else if (Date.now() - startTime > timeoutMs) {
128        reject(new Error(`Timeout waiting for ${description} after ${timeoutMs}ms`));
129      } else {
130        setTimeout(check, 10);
131      }
132    };
133 
134    check();
135  });
136}
137 
138// Usage example from actual debugging session:
139//
140// BEFORE (flaky):
141// ---------------
142// const messagePromise = agent.sendMessage('Execute tools');
143// await new Promise(r => setTimeout(r, 300)); // Hope tools start in 300ms
144// agent.abort();
145// await messagePromise;
146// await new Promise(r => setTimeout(r, 50));  // Hope results arrive in 50ms
147// expect(toolResults.length).toBe(2);         // Fails randomly
148//
149// AFTER (reliable):
150// ----------------
151// const messagePromise = agent.sendMessage('Execute tools');
152// await waitForEventCount(threadManager, threadId, 'TOOL_CALL', 2); // Wait for tools to start
153// agent.abort();
154// await messagePromise;
155// await waitForEventCount(threadManager, threadId, 'TOOL_RESULT', 2); // Wait for results
156// expect(toolResults.length).toBe(2); // Always succeeds
157//
158// Result: 60% pass rate → 100%, 40% faster execution
159