@kapeta/local-cluster-service 0.49.0 → 0.51.0

package/src/storm/routes.ts

@@ -11,12 +11,17 @@ import { stringBody } from '../middleware/stringBody';
 import { KapetaBodyRequest } from '../types';
 import { StormContextRequest, StormCreateBlockRequest, StormStream } from './stream';
 import { ConversationIdHeader, stormClient } from './stormClient';
-import { StormEvent } from './events';
-import { resolveOptions, StormDefinitions, StormEventParser } from './event-parser';
+import { StormEvent, StormEventPhaseType } from './events';
+import {
+    createPhaseEndEvent,
+    createPhaseStartEvent,
+    resolveOptions,
+    StormDefinitions,
+    StormEventParser,
+} from './event-parser';
 import { StormCodegen } from './codegen';
 import { assetManager } from '../assetManager';
 import Path from 'path';
-import { normalizeKapetaUri } from '@kapeta/nodejs-utils';
 
 const router = Router();
 
@@ -36,18 +41,48 @@ router.post('/:handle/all', async (req: KapetaBodyRequest, res: Response) => {
         const aiRequest: StormContextRequest = JSON.parse(req.stringBody ?? '{}');
         const metaStream = await stormClient.createMetadata(aiRequest.prompt, conversationId);
 
+        onRequestAborted(req, res, () => {
+            metaStream.abort();
+        });
+
         res.set('Content-Type', 'application/x-ndjson');
         res.set('Access-Control-Expose-Headers', ConversationIdHeader);
         res.set(ConversationIdHeader, metaStream.getConversationId());
 
+        let currentPhase = StormEventPhaseType.META;
+
         metaStream.on('data', (data: StormEvent) => {
             const result = eventParser.processEvent(req.params.handle, data);
 
+            switch (data.type) {
+                case 'CREATE_API':
+                case 'CREATE_MODEL':
+                case 'CREATE_TYPE':
+                    if (currentPhase !== StormEventPhaseType.DEFINITIONS) {
+                        sendEvent(res, createPhaseEndEvent(StormEventPhaseType.META));
+                        currentPhase = StormEventPhaseType.DEFINITIONS;
+                        sendEvent(res, createPhaseStartEvent(StormEventPhaseType.DEFINITIONS));
+                    }
+                    break;
+            }
+
             sendEvent(res, data);
             sendDefinitions(res, result);
         });
 
-        await waitForStormStream(metaStream);
+        try {
+            sendEvent(res, createPhaseStartEvent(StormEventPhaseType.META));
+
+            await waitForStormStream(metaStream);
+        } finally {
+            if (!metaStream.isAborted()) {
+                sendEvent(res, createPhaseEndEvent(currentPhase));
+            }
+        }
+
+        if (metaStream.isAborted()) {
+            return;
+        }
 
         if (!eventParser.isValid()) {
             // We can't continue if the meta stream is invalid
@@ -63,27 +98,44 @@ router.post('/:handle/all', async (req: KapetaBodyRequest, res: Response) => {
 
         const result = eventParser.toResult(handle);
 
+        if (metaStream.isAborted()) {
+            return;
+        }
+
         sendDefinitions(res, result);
 
         if (!req.query.skipCodegen) {
-            const stormCodegen = new StormCodegen(
-                metaStream.getConversationId(),
-                aiRequest.prompt,
-                result.blocks,
-                eventParser.getEvents()
-            );
-
-            const codegenPromise = streamStormPartialResponse(stormCodegen.getStream(), res);
-
-            await stormCodegen.process();
-
-            await codegenPromise;
+            try {
+                sendEvent(res, createPhaseStartEvent(StormEventPhaseType.IMPLEMENTATION));
+                const stormCodegen = new StormCodegen(
+                    metaStream.getConversationId(),
+                    aiRequest.prompt,
+                    result.blocks,
+                    eventParser.getEvents()
+                );
+
+                onRequestAborted(req, res, () => {
+                    stormCodegen.abort();
+                });
+
+                const codegenPromise = streamStormPartialResponse(stormCodegen.getStream(), res);
+
+                await stormCodegen.process();
+
+                await codegenPromise;
+            } finally {
+                if (!metaStream.isAborted()) {
+                    sendEvent(res, createPhaseEndEvent(StormEventPhaseType.IMPLEMENTATION));
+                }
+            }
         }
 
         sendDone(res);
     } catch (err: any) {
         sendError(err, res);
-        res.end();
+        if (!res.closed) {
+            res.end();
+        }
     }
 });
 
@@ -93,19 +145,13 @@ router.post('/block/create', async (req: KapetaBodyRequest, res: Response) => {
     try {
         const ymlPath = Path.join(createRequest.newPath, 'kapeta.yml');
 
-        console.log('Creating block at', ymlPath);
-
         const [asset] = await assetManager.createAsset(ymlPath, createRequest.definition);
 
         if (await FS.pathExists(createRequest.tmpPath)) {
-            console.log('Moving block from', createRequest.tmpPath, 'to', createRequest.newPath);
-
             await FS.move(createRequest.tmpPath, createRequest.newPath, {
                 overwrite: true,
             });
 
-            console.log('Updating asset', asset.ref);
-
             res.send(await assetManager.updateAsset(asset.ref, createRequest.definition));
         } else {
             res.send(asset);
@@ -125,6 +171,9 @@ function sendDefinitions(res: Response, result: StormDefinitions) {
 }
 
 function sendDone(res: Response) {
+    if (res.closed) {
+        return;
+    }
     sendEvent(res, {
         type: 'DONE',
         created: Date.now(),
@@ -134,6 +183,9 @@ function sendDone(res: Response) {
 }
 
 function sendError(err: Error, res: Response) {
+    if (res.closed) {
+        return;
+    }
     console.error('Failed to send prompt', err);
     if (res.headersSent) {
         sendEvent(res, {
@@ -175,7 +227,19 @@ function streamStormPartialResponse(result: StormStream, res: Response) {
 }
 
 function sendEvent(res: Response, evt: StormEvent) {
+    if (res.closed) {
+        return;
+    }
     res.write(JSON.stringify(evt) + '\n');
 }
 
+function onRequestAborted(req: KapetaBodyRequest, res: Response, onAborted: () => void) {
+    req.on('close', () => {
+        onAborted();
+    });
+    res.on('close', () => {
+        onAborted();
+    });
+}
+
 export default router;

package/src/storm/stormClient.ts

@@ -13,6 +13,7 @@ import {
     StormStream,
     StormUIImplementationPrompt,
 } from './stream';
+import { getRawAsset } from 'node:sea';
 
 export const STORM_ID = 'storm';
 
@@ -60,6 +61,9 @@ class StormClient {
             conversationId: body.conversationId,
         });
 
+        const abort = new AbortController();
+        options.signal = abort.signal;
+
         const response = await fetch(options.url, options);
 
         if (response.status !== 200) {
@@ -69,7 +73,6 @@ class StormClient {
         }
 
         const conversationId = response.headers.get(ConversationIdHeader);
-        console.log('Received conversationId', conversationId);
 
         const out = new StormStream(stringPrompt, conversationId);
 
@@ -87,6 +90,10 @@ class StormClient {
             out.end();
         });
 
+        out.on('aborted', () => {
+            abort.abort();
+        });
+
         return out;
     }
 

package/src/storm/stream.ts

@@ -10,6 +10,7 @@ import { BlockDefinition } from '@kapeta/schemas';
 export class StormStream extends EventEmitter {
     private conversationId: string = '';
     private lines: string[] = [];
+    private aborted: boolean = false;
 
     constructor(prompt: string = '', conversationId?: string | null) {
         super();
@@ -20,6 +21,10 @@ export class StormStream extends EventEmitter {
         return this.conversationId;
     }
 
+    isAborted() {
+        return this.aborted;
+    }
+
     addJSONLine(line: string) {
         try {
             this.lines.push(line);
@@ -38,6 +43,7 @@ export class StormStream extends EventEmitter {
     }
 
     on(event: 'end', listener: () => void): this;
+    on(event: 'aborted', listener: () => void): this;
     on(event: 'error', listener: (e: Error) => void): this;
     on(event: 'data', listener: (data: StormEvent) => void): this;
     on(event: string, listener: (...args: any[]) => void): this {
@@ -45,6 +51,7 @@ export class StormStream extends EventEmitter {
     }
 
     emit(event: 'end'): boolean;
+    emit(event: 'aborted'): void;
     emit(event: 'error', e: Error): boolean;
     emit(event: 'data', data: StormEvent): boolean;
     emit(eventName: string | symbol, ...args: any[]): boolean {
@@ -62,6 +69,14 @@ export class StormStream extends EventEmitter {
             });
         });
     }
+
+    abort() {
+        if (this.aborted) {
+            return;
+        }
+        this.aborted = true;
+        this.emit('aborted');
+    }
 }
 
 export interface ConversationItem {