Merge pull request #3 from watercrawl/feature/fix-async

amirasaran · web-flow · commit 0c8e69d22ea3 · 2025-05-03T11:53:23.000+02:00
fix(stream): implement improved async event streaming mechanism
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.1.1] - 2025-05-03
+
+### Fixed
+- Improved async streaming implementation for more reliable event handling
+- Fixed search event monitoring to correctly process event status
+- Refactored stream processing for better error handling and resource management
+- Removed debug console.log statements from tests
+
 ## [1.1.0] - 2025-04-30
 
 ### Added
diff --git a/src/base.ts b/src/base.ts
@@ -4,9 +4,12 @@ import urlJoin from 'url-join';
 export class BaseAPIClient {
     protected client: AxiosInstance;
     protected baseUrl: string;
+    protected apiKey: string;
 
     constructor(apiKey: string, baseUrl: string = 'https://app.watercrawl.dev') {
+        this.apiKey = apiKey;
         this.baseUrl = baseUrl;
+
         this.client = axios.create({
             baseURL: this.baseUrl,
             headers: {
@@ -16,21 +19,19 @@ export class BaseAPIClient {
             }
         });
 
-        // Add response interceptor for error handling
         this.client.interceptors.response.use(
-            response => response,
-            error => {
-                if (error.response) {
-                    // Log API errors
-                    console.error('API Error:', {
-                        url: error.config.url,
-                        status: error.response.status,
-                        data: error.response.data,
-                        headers: error.response.headers
-                    });
-                }
-                throw error;
-            }
+          response => response,
+          error => {
+              if (error.response) {
+                  console.error('API Error:', {
+                      url: error.config.url,
+                      status: error.response.status,
+                      data: error.response.data,
+                      headers: error.response.headers
+                  });
+              }
+              throw error;
+          }
         );
     }
 
@@ -63,30 +64,54 @@ export class BaseAPIClient {
         return urlJoin(this.baseUrl, ...parts);
     }
 
-    protected async streamEvents(endpoint: string, onEvent: (event: any) => void, config: AxiosRequestConfig = {}): Promise<void> {
-        const response = await this.client.get(endpoint, {
-            responseType: 'stream',
-            ...config
+
+    /**
+     * Async generator that streams SSE data using fetch and yields parsed JSON
+     */
+    protected async *fetchStream<T>(endpoint: string, config: AxiosRequestConfig = {}): AsyncGenerator<T> {
+        const url = new URL(this.buildUrl(endpoint));
+        if(config.params) {
+            Object.keys(config.params).forEach(key => {
+                url.searchParams.append(key, config.params[key]);
+            })
+        }
+        const response = await fetch(url.toString(), {
+            method: 'GET',
+            headers: {
+                // 'Accept': 'text/event-stream',
+                'X-API-KEY': this.apiKey
+            }
         });
+
+        if (!response.body) {
+            throw new Error("No response body");
+        }
+
+        const reader = response.body.getReader();
+        const decoder = new TextDecoder("utf-8");
         let buffer = '';
 
-        const stream = response.data;
-        stream.on('data', (chunk: string) => {
-            buffer += chunk.toString();
+        while (true) {
+            const { done, value } = await reader.read();
+            if (done) break;
+
+            buffer += decoder.decode(value, { stream: true });
+
             const lines = buffer.split('\n');
             buffer = lines.pop() || '';
+
             for (const line of lines) {
-                if (line.startsWith('data: ')) {
+                const trimmed = line.trim();
+                if (trimmed.startsWith('data:')) {
+                    const dataStr = trimmed.slice(5).trim();
+                    if (dataStr === '[DONE]') return;
                     try {
-                        const data = JSON.parse(line.slice(6));
-                        onEvent(data);
-                    } catch (error) {
-                        console.error('Error parsing event data:', line, error);
+                        yield JSON.parse(dataStr) as T;
+                    } catch (err) {
+                        console.error('Failed to parse JSON from stream:', dataStr, err);
                     }
                 }
             }
-        });
-
-        return response.data;
+        }
     }
 }
diff --git a/src/index.ts b/src/index.ts
@@ -53,56 +53,11 @@ export class WaterCrawlAPIClient extends BaseAPIClient {
     }
 
     async *monitorCrawlRequest(itemId: string, download: boolean = true): AsyncGenerator<CrawlEvent, void, unknown> {
-        const events: CrawlEvent[] = [];
-        let resolveNext: ((value: IteratorResult<CrawlEvent, void>) => void) | null = null;
-        let isDone = false;
-        let streamError: Error | null = null;
 
-        const processEvent = (event: CrawlEvent) => {
-            if (resolveNext) {
-                resolveNext({ value: event, done: false });
-                resolveNext = null;
-            } else {
-                events.push(event);
-            }
-        };
-
-        const streamPromise = this.streamEvents(
+        yield* this.fetchStream<CrawlEvent>(
             `/api/v1/core/crawl-requests/${itemId}/status/`,
-            processEvent,
             { params: { prefetched: download } }
-        ).catch((error) => {
-            streamError = error;
-            isDone = true;
-        }).finally(() => {
-            isDone = true;
-        });
-    
-
-        try {
-            while (!isDone || events.length > 0) {
-                if (events.length > 0) {
-                    yield events.shift()!;
-                } else if (!isDone) {
-                    await new Promise<IteratorResult<CrawlEvent, void>>((resolve) => {
-                        resolveNext = resolve;
-                    });
-                }
-            }
-    
-            // Check for any remaining events after the stream ends
-            while (events.length > 0) {
-                yield events.shift()!;
-            }
-    
-            // If the stream failed, propagate the error
-            if (streamError) {
-                throw streamError;
-            }
-        } finally {
-            // Ensure the stream is awaited and cleaned up properly
-            await streamPromise;
-        }
+        );
     }
 
     async getCrawlRequestResults(itemId: string): Promise<{ results: CrawlResult[] }> {
@@ -148,14 +103,14 @@ export class WaterCrawlAPIClient extends BaseAPIClient {
      * @throws Error if the sitemap is not available
      */
     private async getCrawlRequestForSitemap(crawlRequest: string | CrawlRequest): Promise<CrawlRequest> {
-        const request = typeof crawlRequest === 'string' 
-            ? await this.getCrawlRequest(crawlRequest) 
+        const request = typeof crawlRequest === 'string'
+            ? await this.getCrawlRequest(crawlRequest)
             : crawlRequest;
-        
+
         if (!request.sitemap) {
             throw new Error('Sitemap not found in crawl request');
         }
-        
+
         return request;
     }
 
@@ -166,11 +121,11 @@ export class WaterCrawlAPIClient extends BaseAPIClient {
      */
     async downloadSitemap(crawlRequest: string | CrawlRequest): Promise<SitemapNode[]> {
         const request = await this.getCrawlRequestForSitemap(crawlRequest);
-        
+
         if (!request.sitemap) {
             throw new Error('Sitemap URL is missing or undefined');
         }
-        
+
         const response = await axios.get(request.sitemap);
         return response.data;
     }
@@ -237,23 +192,23 @@ export class WaterCrawlAPIClient extends BaseAPIClient {
             search_options: searchOptions,
             result_limit: resultLimit
         };
-        
+
         const response = await this.post<SearchRequest>('/api/v1/core/search/', request);
-        
+
         if (!sync) {
             return response;
         }
-        
+
         // Monitor the search request until completion
         for await (const event of this.monitorSearchRequest(response.uuid, download)) {
-            if (event.type === 'state' && ['finished', 'failed'].includes(event.status)) {
+            if (event.type === 'state' && ['finished', 'failed'].includes(event.data.status)) {
                 if (download && Array.isArray(event.data.result)) {
                     return event.data.result as SearchResult[];
                 }
                 return event.data;
             }
         }
-        
+
         throw new Error('Search request failed or timed out');
     }
 
@@ -264,55 +219,10 @@ export class WaterCrawlAPIClient extends BaseAPIClient {
      * @returns AsyncGenerator yielding search events
      */
     async *monitorSearchRequest(itemId: string, download: boolean = true): AsyncGenerator<SearchEvent, void, unknown> {
-        const events: SearchEvent[] = [];
-        let resolveNext: ((value: IteratorResult<SearchEvent, void>) => void) | null = null;
-        let isDone = false;
-        let streamError: Error | null = null;
-
-        const processEvent = (event: SearchEvent) => {
-            if (resolveNext) {
-                resolveNext({ value: event, done: false });
-                resolveNext = null;
-            } else {
-                events.push(event);
-            }
-        };
-
-        const streamPromise = this.streamEvents(
-            `/api/v1/core/search/${itemId}/status/`,
-            processEvent,
-            { params: { prefetched: download } }
-        ).catch((error) => {
-            streamError = error;
-            isDone = true;
-        }).finally(() => {
-            isDone = true;
-        });
-
-        try {
-            while (!isDone || events.length > 0) {
-                if (events.length > 0) {
-                    yield events.shift()!;
-                } else if (!isDone) {
-                    await new Promise<IteratorResult<SearchEvent, void>>((resolve) => {
-                        resolveNext = resolve;
-                    });
-                }
-            }
-    
-            // Check for any remaining events after the stream ends
-            while (events.length > 0) {
-                yield events.shift()!;
-            }
-    
-            // If the stream failed, propagate the error
-            if (streamError) {
-                throw streamError;
-            }
-        } finally {
-            // Ensure the stream is awaited and cleaned up properly
-            await streamPromise;
-        }
+        yield* this.fetchStream<SearchEvent>(
+          `/api/v1/core/search/${itemId}/status/`,
+          { params: { prefetched: download } }
+        );
     }
 
     /**
diff --git a/test/api.test.ts b/test/api.test.ts
@@ -49,7 +49,6 @@ describe('WaterCrawlAPI', () => {
 
         let count = 0;
         for await (const item of generator) {
-            console.log("TEST", item);
             expect(typeof item).toBe('object');
             count++;
             if (count >= 2) break; // Only test first few events to keep test duration reasonable
