@kapeta/local-cluster-service 0.52.2 → 0.53.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -3,7 +3,7 @@
3
3
  * SPDX-License-Identifier: BUSL-1.1
4
4
  */
5
5
 
6
- import {Definition} from '@kapeta/local-cluster-config';
6
+ import { Definition } from '@kapeta/local-cluster-config';
7
7
  import {
8
8
  AIFileTypes,
9
9
  BlockCodeGenerator,
@@ -13,27 +13,27 @@ import {
13
13
  GeneratedResult,
14
14
  MODE_CREATE_ONLY,
15
15
  } from '@kapeta/codegen';
16
- import {BlockDefinition} from '@kapeta/schemas';
17
- import {codeGeneratorManager} from '../codeGeneratorManager';
18
- import {STORM_ID, stormClient} from './stormClient';
16
+ import { BlockDefinition } from '@kapeta/schemas';
17
+ import { codeGeneratorManager } from '../codeGeneratorManager';
18
+ import { STORM_ID, stormClient } from './stormClient';
19
19
  import {
20
20
  StormEvent,
21
+ StormEventBlockStatusType,
21
22
  StormEventErrorDetailsFile,
22
23
  StormEventFileChunk,
23
24
  StormEventFileDone,
24
- StormEventFileLogical
25
+ StormEventFileLogical,
25
26
  } from './events';
26
- import {BlockDefinitionInfo, StormEventParser} from './event-parser';
27
- import {ConversationItem, StormFileImplementationPrompt, StormFileInfo, StormStream} from './stream';
28
- import {KapetaURI, parseKapetaUri} from '@kapeta/nodejs-utils';
29
- import {writeFile} from 'fs/promises';
27
+ import { BlockDefinitionInfo, StormEventParser } from './event-parser';
28
+ import { ConversationItem, StormFileImplementationPrompt, StormFileInfo, StormStream } from './stream';
29
+ import { KapetaURI, parseKapetaUri } from '@kapeta/nodejs-utils';
30
+ import { writeFile } from 'fs/promises';
30
31
  import path from 'path';
31
- import Path, {join} from 'path';
32
+ import Path, { join } from 'path';
32
33
  import os from 'node:os';
33
- import {readFileSync, writeFileSync} from 'fs';
34
- import YAML from "yaml";
35
- import * as fs from "node:fs";
36
- import {v4 as uuidv4} from 'uuid';
34
+ import { readFileSync, writeFileSync } from 'fs';
35
+ import YAML from 'yaml';
36
+ import * as fs from 'node:fs';
37
37
 
38
38
  type ImplementationGenerator = (prompt: StormFileImplementationPrompt, conversationId?: string) => Promise<StormStream>;
39
39
 
@@ -44,7 +44,7 @@ interface ErrorClassification {
44
44
  }
45
45
 
46
46
  const SIMULATED_DELAY = 1000;
47
-
47
+ const ENABLE_SIMULATED_DELAY = false;
48
48
  class SimulatedFileDelay {
49
49
  private readonly file: StormEventFileDone;
50
50
  public readonly stream;
@@ -166,43 +166,8 @@ export class StormCodegen {
166
166
  },
167
167
  });
168
168
  break;
169
- case 'FILE_START':
170
- case 'FILE_CHUNK_RESET':
171
- this.out.emit('data', {
172
- ...data,
173
- payload: {
174
- ...data.payload,
175
- blockName,
176
- blockRef,
177
- instanceId,
178
- },
179
- });
180
- break;
181
- case 'FILE_CHUNK':
182
- this.out.emit('data', {
183
- ...data,
184
- payload: {
185
- ...data.payload,
186
- blockName,
187
- blockRef,
188
- instanceId,
189
- },
190
- });
191
- break;
192
- case 'FILE_STATE':
193
- this.out.emit('data', {
194
- ...data,
195
- payload: {
196
- ...data.payload,
197
- blockName,
198
- blockRef,
199
- instanceId,
200
- },
201
- });
202
- break;
203
169
  case 'FILE_DONE':
204
170
  return this.handleFileDoneOutput(blockUri, blockName, data);
205
- break;
206
171
  }
207
172
  }
208
173
 
@@ -335,11 +300,13 @@ export class StormCodegen {
335
300
  (file) => ![AIFileTypes.SERVICE, AIFileTypes.WEB_SCREEN].includes(file.type)
336
301
  );
337
302
 
303
+ const blockUri = parseKapetaUri(block.uri);
304
+
338
305
  // Send the service and UI templates to the AI. These will be sent one-by-one in addition to the context files
339
306
  const serviceFiles: StormFileInfo[] = allFiles.filter((file) => file.type === AIFileTypes.SERVICE);
340
307
  if (serviceFiles.length > 0) {
341
308
  await this.processTemplates(
342
- parseKapetaUri(block.uri),
309
+ blockUri,
343
310
  block.aiName,
344
311
  stormClient.createServiceImplementation.bind(stormClient),
345
312
  serviceFiles,
@@ -371,7 +338,7 @@ export class StormCodegen {
371
338
  await writeFile(filePath, screenFile.payload.content);
372
339
  }
373
340
 
374
- const screenFilesConverted = screenFiles.map(screenFile => {
341
+ const screenFilesConverted = screenFiles.map((screenFile) => {
375
342
  return {
376
343
  filename: screenFile.payload.filename,
377
344
  content: screenFile.payload.content,
@@ -381,12 +348,16 @@ export class StormCodegen {
381
348
  };
382
349
  });
383
350
  allFiles.push(...screenFilesConverted);
351
+ const blockRef = block.uri;
352
+
353
+ this.emitBlockStatus(blockUri, block.aiName, StormEventBlockStatusType.QA);
384
354
 
385
355
  const filesToBeFixed = serviceFiles.concat(contextFiles).concat(screenFilesConverted);
386
356
  const codeGenerator = new BlockCodeGenerator(block.content as BlockDefinition);
387
- await this.verifyAndFixCode(codeGenerator, basePath, filesToBeFixed, allFiles);
388
357
 
389
- const blockRef = block.uri;
358
+ this.emitBlockStatus(blockUri, block.aiName, StormEventBlockStatusType.BUILDING);
359
+ await this.verifyAndFixCode(blockUri, block.aiName, codeGenerator, basePath, filesToBeFixed, allFiles);
360
+
390
361
  this.out.emit('data', {
391
362
  type: 'BLOCK_READY',
392
363
  reason: 'Block ready',
@@ -400,7 +371,23 @@ export class StormCodegen {
400
371
  } satisfies StormEvent);
401
372
  }
402
373
 
374
+ private emitBlockStatus(blockUri: KapetaURI, blockName: string, status: StormEventBlockStatusType) {
375
+ this.out.emit('data', {
376
+ type: 'BLOCK_STATUS',
377
+ reason: status,
378
+ created: Date.now(),
379
+ payload: {
380
+ status,
381
+ blockName,
382
+ blockRef: blockUri.toNormalizedString(),
383
+ instanceId: StormEventParser.toInstanceIdFromRef(blockUri.toNormalizedString()),
384
+ },
385
+ } satisfies StormEvent);
386
+ }
387
+
403
388
  private async verifyAndFixCode(
389
+ blockUri: KapetaURI,
390
+ blockName: string,
404
391
  codeGenerator: CodeGenerator,
405
392
  basePath: string,
406
393
  filesToBeFixed: StormFileInfo[],
@@ -421,11 +408,30 @@ export class StormCodegen {
421
408
  if (result && !result.valid) {
422
409
  console.debug('Validation error:', result);
423
410
 
411
+ this.emitBlockStatus(blockUri, blockName, StormEventBlockStatusType.PLANNING_FIX);
412
+
424
413
  const errors = await this.classifyErrors(result.error, basePath);
425
- for (const [filename, fileErrors] of errors.entries()) {
426
- // todo: only try to fix file if it is part of filesToBeFixed
427
- await this.tryToFixFile(basePath, filename, fileErrors, allFiles, codeGenerator);
414
+
415
+ if (errors.size > 0) {
416
+ this.emitBlockStatus(blockUri, blockName, StormEventBlockStatusType.FIXING);
417
+
418
+ const promises = Array.from(errors.entries()).map(([filename, fileErrors]) => {
419
+ // todo: only try to fix file if it is part of filesToBeFixed
420
+ return this.tryToFixFile(
421
+ blockUri,
422
+ blockName,
423
+ basePath,
424
+ filename,
425
+ fileErrors,
426
+ allFiles,
427
+ codeGenerator
428
+ );
429
+ });
430
+
431
+ await Promise.all(promises);
428
432
  }
433
+
434
+ this.emitBlockStatus(blockUri, blockName, StormEventBlockStatusType.FIX_DONE);
429
435
  }
430
436
  } catch (e) {
431
437
  console.error('Error:', e);
@@ -439,33 +445,57 @@ export class StormCodegen {
439
445
  }
440
446
  }
441
447
 
442
- private async tryToFixFile(basePath: string, filename: string, fileErrors: ErrorClassification[], allFiles: StormFileInfo[], codeGenerator: CodeGenerator) {
448
+ private async tryToFixFile(
449
+ blockUri: KapetaURI,
450
+ blockName: string,
451
+ basePath: string,
452
+ filename: string,
453
+ fileErrors: ErrorClassification[],
454
+ allFiles: StormFileInfo[],
455
+ codeGenerator: CodeGenerator
456
+ ) {
443
457
  console.log(`Processing ${filename}`);
444
458
  const language = await codeGenerator.language();
445
- const relevantFiles = allFiles.filter(file => file.type != AIFileTypes.IGNORE);
459
+ const relevantFiles = allFiles.filter((file) => file.type != AIFileTypes.IGNORE);
446
460
 
447
- for (let attempts = 1; attempts <= 5; attempts++) {
461
+ for (let attempts = 1; attempts <= 5; attempts++) {
448
462
  if (fileErrors.length == 0) {
449
463
  console.log(`No more errors for ${filename}`);
450
464
  return;
451
465
  }
452
466
 
453
467
  console.log(`Errors in ${filename} - requesting error details`);
454
- const filesForContext = await this.getErrorDetailsForFile(basePath, filename, fileErrors[0], relevantFiles, language);
468
+ const filesForContext = await this.getErrorDetailsForFile(
469
+ basePath,
470
+ filename,
471
+ fileErrors[0],
472
+ relevantFiles,
473
+ language
474
+ );
455
475
  console.log(`Get error details for ${filename} requesting code fixes`);
456
476
 
457
- const fix = this.createFixRequestForFile(basePath, filename, fileErrors[0], filesForContext, relevantFiles, language);
458
- const codeFixFiles = await this.codeFix(fix);
477
+ const fix = this.createFixRequestForFile(
478
+ basePath,
479
+ filename,
480
+ fileErrors[0],
481
+ filesForContext,
482
+ relevantFiles,
483
+ language
484
+ );
485
+ const codeFixFile = await this.codeFix(blockUri, blockName, fix);
459
486
  console.log(`Got fixed code for ${filename}`);
460
- for (const codeFixFile of codeFixFiles) {
461
- const filePath = codeFixFile.filename.indexOf(basePath) > -1 ? codeFixFile.filename : join(basePath, codeFixFile.filename);
462
- const existing = readFileSync(filePath);
463
- if (existing.toString().replace(/(\r\n|\r|\n)+$/, '') == codeFixFile.content.replace(/(\r\n|\r|\n)+$/, '')) {
464
- console.log(`${filename} not changed by gemini`);
465
- continue;
466
- }
467
- writeFileSync(filePath, codeFixFile.content);
487
+ const filePath =
488
+ codeFixFile.filename.indexOf(basePath) > -1
489
+ ? codeFixFile.filename
490
+ : join(basePath, codeFixFile.filename);
491
+ const existing = readFileSync(filePath);
492
+ if (
493
+ existing.toString().replace(/(\r\n|\r|\n)+$/, '') == codeFixFile.content.replace(/(\r\n|\r|\n)+$/, '')
494
+ ) {
495
+ console.log(`${filename} not changed by gemini`);
496
+ continue;
468
497
  }
498
+ writeFileSync(filePath, codeFixFile.content);
469
499
 
470
500
  const result = await codeGenerator.validateForTarget(basePath);
471
501
  if (result && result.valid) {
@@ -488,7 +518,11 @@ export class StormCodegen {
488
518
  errorStream.on('data', (evt) => {
489
519
  if (evt.type === 'ERROR_CLASSIFIER') {
490
520
  const eventFileName = this.removePrefix(basePath + '/', evt.payload.filename);
491
- const fix = { error: evt.payload.error, lineNumber: evt.payload.lineNumber, column: evt.payload.column };
521
+ const fix = {
522
+ error: evt.payload.error,
523
+ lineNumber: evt.payload.lineNumber,
524
+ column: evt.payload.column,
525
+ };
492
526
 
493
527
  let existingFixes = fixes.get(eventFileName);
494
528
  if (existingFixes) {
@@ -504,17 +538,23 @@ export class StormCodegen {
504
538
  return fixes;
505
539
  }
506
540
 
507
- private async getErrorDetailsForFile(basePath: string, filename: string, error: ErrorClassification, allFiles: StormFileInfo[], language: string): Promise<string[]> {
541
+ private async getErrorDetailsForFile(
542
+ basePath: string,
543
+ filename: string,
544
+ error: ErrorClassification,
545
+ allFiles: StormFileInfo[],
546
+ language: string
547
+ ): Promise<string[]> {
508
548
  const filePath = filename.indexOf(basePath) > -1 ? filename : join(basePath, filename); // to compensate when compiler returns absolute path
509
549
  return new Promise<string[]>(async (resolve, reject) => {
510
550
  const request = {
511
- "language": language,
512
- "sourceFile": {
513
- "filename": filename,
514
- "content": readFileSync(filePath, 'utf8')
551
+ language: language,
552
+ sourceFile: {
553
+ filename: filename,
554
+ content: readFileSync(filePath, 'utf8'),
515
555
  },
516
- "error": error,
517
- "projectFiles": allFiles.map(f => f.filename)
556
+ error: error,
557
+ projectFiles: allFiles.map((f) => f.filename),
518
558
  };
519
559
 
520
560
  const detailsStream = await stormClient.createErrorDetails(JSON.stringify(request), []);
@@ -522,32 +562,38 @@ export class StormCodegen {
522
562
  if (evt.type === 'ERROR_DETAILS') {
523
563
  resolve(evt.payload.files);
524
564
  }
565
+ reject(new Error('Error details: Unexpected event [' + evt.type + ']'));
525
566
  });
526
567
  this.out.on('aborted', () => {
527
568
  detailsStream.abort();
569
+ reject(new Error('aborted'));
528
570
  });
529
571
  detailsStream.on('error', (err) => {
530
- console.log("error", err);
531
572
  reject(err);
532
573
  });
533
574
  await detailsStream.waitForDone();
534
575
  });
535
576
  }
536
577
 
537
- private createFixRequestForFile(basePath: string, filename: string, error: ErrorClassification, filesForContext: string[], allFiles: StormFileInfo[], language: string): string {
578
+ private createFixRequestForFile(
579
+ basePath: string,
580
+ filename: string,
581
+ error: ErrorClassification,
582
+ filesForContext: string[],
583
+ allFiles: StormFileInfo[],
584
+ language: string
585
+ ): string {
538
586
  const files = new Set(filesForContext);
539
587
  files.add(filename);
540
588
 
541
- const requestedFiles = Array.from(files).flatMap(file => {
589
+ const requestedFiles = Array.from(files).flatMap((file) => {
542
590
  if (fs.existsSync(file)) {
543
591
  return file;
544
592
  }
545
593
 
546
594
  // file does not exist - look for similar
547
595
  const candidateName = file.split('/').pop();
548
- return allFiles
549
- .filter(file => file.filename.split('/').pop() === candidateName)
550
- .map(f => f.filename);
596
+ return allFiles.filter((file) => file.filename.split('/').pop() === candidateName).map((f) => f.filename);
551
597
  });
552
598
 
553
599
  const filePath = filename.indexOf(basePath) > -1 ? filename : join(basePath, filename);
@@ -555,31 +601,26 @@ export class StormCodegen {
555
601
  const affectedLine = this.getErrorLine(error, content);
556
602
 
557
603
  const fixRequest = {
558
- "language": language,
559
- "filename": filename,
560
- "errors": error.error,
561
- "affectedLine": affectedLine,
562
- "projectFiles": requestedFiles.map(filename => {
604
+ language: language,
605
+ filename: filename,
606
+ error: error.error,
607
+ affectedLine: affectedLine,
608
+ projectFiles: requestedFiles.map((filename) => {
563
609
  const filePath = filename.indexOf(basePath) > -1 ? filename : join(basePath, filename);
564
610
  const content = readFileSync(filePath, 'utf8');
565
611
  return { filename: filename, content: content };
566
612
  }),
567
- "conversationId": this.conversationId,
568
- "sessionId": uuidv4()
569
613
  };
570
614
 
571
615
  return JSON.stringify(fixRequest);
572
616
  }
573
617
 
574
- private getErrorLine(
575
- errorDetails: ErrorClassification,
576
- sourceCode: string
577
- ): string {
618
+ private getErrorLine(errorDetails: ErrorClassification, sourceCode: string): string {
578
619
  const lines = sourceCode.split('\n');
579
620
  const errorLine = lines[errorDetails.lineNumber - 1];
580
621
 
581
622
  if (!errorLine) {
582
- return "Error: Line number out of range.";
623
+ return 'Error: Line number out of range.';
583
624
  }
584
625
 
585
626
  return errorLine;
@@ -595,22 +636,41 @@ export class StormCodegen {
595
636
  /**
596
637
  * Sends the code to the AI for a fix
597
638
  */
598
- private async codeFix(fix: string, history?: ConversationItem[]): Promise<StormEventErrorDetailsFile[]> {
599
- return new Promise<StormEventErrorDetailsFile[]>(async (resolve, reject) => {
600
- const fixStream = await stormClient.createCodeFix(fix, history);
601
- fixStream.on('data', (evt) => {
602
- if (evt.type === 'CODE_FIX') {
603
- resolve(evt.payload.content.files);
604
- }
605
- });
606
- this.out.on('aborted', () => {
607
- fixStream.abort();
608
- });
609
- fixStream.on('error', (err) => {
610
- console.log("error", err);
611
- reject(err);
612
- });
613
- await fixStream.waitForDone();
639
+ private async codeFix(
640
+ blockUri: KapetaURI,
641
+ blockName: string,
642
+ fix: string,
643
+ history?: ConversationItem[]
644
+ ): Promise<StormEventErrorDetailsFile> {
645
+ return new Promise<StormEventErrorDetailsFile>(async (resolve, reject) => {
646
+ try {
647
+ const fixStream = await stormClient.createCodeFix(fix, history, this.conversationId);
648
+ let resolved = false;
649
+ fixStream.on('data', (evt) => {
650
+ if (this.handleFileEvents(blockUri, blockName, evt)) {
651
+ return;
652
+ }
653
+
654
+ if (evt.type === 'CODE_FIX') {
655
+ resolved = true;
656
+ resolve(evt.payload);
657
+ }
658
+ });
659
+ this.out.on('aborted', () => {
660
+ fixStream.abort();
661
+ reject(new Error('aborted'));
662
+ });
663
+ fixStream.on('error', (err) => {
664
+ reject(err);
665
+ });
666
+ fixStream.on('end', () => {
667
+ if (!resolved) {
668
+ reject(new Error('Code fix never returned a valid event'));
669
+ }
670
+ });
671
+ } catch (e) {
672
+ reject(e);
673
+ }
614
674
  });
615
675
  }
616
676
 
@@ -650,7 +710,13 @@ export class StormCodegen {
650
710
  instanceId: StormEventParser.toInstanceIdFromRef(ref),
651
711
  },
652
712
  };
653
- return new SimulatedFileDelay(fileEvent, this.out).start();
713
+
714
+ if (ENABLE_SIMULATED_DELAY) {
715
+ // Simulate a delay when sending the file
716
+ return new SimulatedFileDelay(fileEvent, this.out).start();
717
+ } else {
718
+ this.out.emit('data', fileEvent);
719
+ }
654
720
  });
655
721
 
656
722
  return Promise.all(promises);
@@ -431,9 +431,12 @@ export class StormEventParser {
431
431
  }
432
432
 
433
433
  const clientTypes = DSLDataTypeParser.parse(
434
- clientConsumerBlock.content.spec.entities.source!.value, {ignoreSemantics: true}
434
+ clientConsumerBlock.content.spec.entities.source!.value,
435
+ { ignoreSemantics: true }
435
436
  );
436
- const apiTypes = DSLDataTypeParser.parse(apiProviderBlock.content.spec.entities?.source?.value, {ignoreSemantics: true});
437
+ const apiTypes = DSLDataTypeParser.parse(apiProviderBlock.content.spec.entities?.source?.value, {
438
+ ignoreSemantics: true,
439
+ });
437
440
 
438
441
  apiTypes.forEach((apiType) => {
439
442
  if (clientTypes.some((clientType) => clientType.name === apiType.name)) {
@@ -139,9 +139,8 @@ export interface StormEventCodeFix {
139
139
  reason: string;
140
140
  created: number;
141
141
  payload: {
142
- content: {
143
- files: StormEventErrorDetailsFile[]
144
- }
142
+ filename: string;
143
+ content: string;
145
144
  };
146
145
  }
147
146
 
@@ -162,7 +161,7 @@ export interface StormEventErrorDetails {
162
161
  reason: string;
163
162
  created: number;
164
163
  payload: {
165
- files: string[]
164
+ files: string[];
166
165
  };
167
166
  }
168
167
 
@@ -252,6 +251,26 @@ export interface StormEventBlockReady {
252
251
  };
253
252
  }
254
253
 
254
+ export enum StormEventBlockStatusType {
255
+ QA = 'QA',
256
+ FIXING = 'FIXING',
257
+ PLANNING_FIX = 'PLANNING_FIX',
258
+ FIX_DONE = 'FIX_DONE',
259
+ BUILDING = 'BUILDING',
260
+ }
261
+
262
+ export interface StormEventBlockStatus {
263
+ type: 'BLOCK_STATUS';
264
+ reason: string;
265
+ created: number;
266
+ payload: {
267
+ status: StormEventBlockStatusType;
268
+ blockName: string;
269
+ blockRef: string;
270
+ instanceId: string;
271
+ };
272
+ }
273
+
255
274
  export interface StormEventDone {
256
275
  type: 'DONE';
257
276
  created: number;
@@ -268,6 +287,7 @@ export enum StormEventPhaseType {
268
287
  META = 'META',
269
288
  DEFINITIONS = 'DEFINITIONS',
270
289
  IMPLEMENTATION = 'IMPLEMENTATION',
290
+ QA = 'QA',
271
291
  }
272
292
 
273
293
  export interface StormEventPhases {
@@ -299,4 +319,5 @@ export type StormEvent =
299
319
  | StormEventCodeFix
300
320
  | StormEventErrorDetails
301
321
  | StormEventBlockReady
302
- | StormEventPhases;
322
+ | StormEventPhases
323
+ | StormEventBlockStatus;
@@ -215,6 +215,7 @@ function sendEvent(res: Response, evt: StormEvent) {
215
215
  if (res.closed) {
216
216
  return;
217
217
  }
218
+ console.log('Sending event', evt.type);
218
219
  res.write(JSON.stringify(evt) + '\n');
219
220
  }
220
221
 
@@ -26,7 +26,11 @@ class StormClient {
26
26
  this._baseUrl = getRemoteUrl('ai-service', 'https://ai.kapeta.com');
27
27
  }
28
28
 
29
- private async createOptions(path: string, method: string, body: StormContextRequest): Promise<RequestInit & { url: string }> {
29
+ private async createOptions(
30
+ path: string,
31
+ method: string,
32
+ body: StormContextRequest
33
+ ): Promise<RequestInit & { url: string }> {
30
34
  const url = `${this._baseUrl}${path}`;
31
35
  const headers: { [k: string]: string } = {
32
36
  'Content-Type': 'application/json',
@@ -92,7 +96,11 @@ class StormClient {
92
96
  });
93
97
 
94
98
  out.on('aborted', () => {
95
- abort.abort();
99
+ try {
100
+ abort.abort();
101
+ } catch (e) {
102
+ console.warn('Error aborting stream', e);
103
+ }
96
104
  });
97
105
 
98
106
  return out;
@@ -11,6 +11,7 @@ export class StormStream extends EventEmitter {
11
11
  private conversationId: string = '';
12
12
  private lines: string[] = [];
13
13
  private aborted: boolean = false;
14
+ private done: boolean = false;
14
15
 
15
16
  constructor(prompt: string = '', conversationId?: string | null) {
16
17
  super();
@@ -39,6 +40,7 @@ export class StormStream extends EventEmitter {
39
40
  }
40
41
 
41
42
  end() {
43
+ this.done = true;
42
44
  this.emit('end');
43
45
  }
44
46
 
@@ -59,6 +61,9 @@ export class StormStream extends EventEmitter {
59
61
  }
60
62
 
61
63
  waitForDone() {
64
+ if (this.done) {
65
+ return Promise.resolve();
66
+ }
62
67
  return new Promise<void>((resolve, reject) => {
63
68
  const errorHandler = (err: any) => {
64
69
  this.removeListener('error', errorHandler);
@@ -70,8 +75,8 @@ export class StormStream extends EventEmitter {
70
75
  this.removeListener('end', endHandler);
71
76
  resolve();
72
77
  };
73
- this.once('error', errorHandler);
74
- this.once('end', endHandler);
78
+ this.on('error', errorHandler);
79
+ this.on('end', endHandler);
75
80
  });
76
81
  }
77
82
 
@@ -80,6 +85,7 @@ export class StormStream extends EventEmitter {
80
85
  return;
81
86
  }
82
87
  this.aborted = true;
88
+ this.done = true;
83
89
  this.emit('aborted');
84
90
  }
85
91
  }