keystone-cli 0.5.1 → 0.6.1

package/src/runner/stream-utils.test.ts

@@ -0,0 +1,171 @@
+import { describe, expect, it, mock } from 'bun:test';
+import { processOpenAIStream } from './stream-utils';
+
+const encoder = new TextEncoder();
+
+function responseFromChunks(chunks: string[]): Response {
+  let index = 0;
+  const reader = {
+    async read(): Promise<{ done: boolean; value?: Uint8Array }> {
+      if (index >= chunks.length) {
+        return { done: true, value: undefined };
+      }
+      const value = encoder.encode(chunks[index]);
+      index += 1;
+      return { done: false, value };
+    },
+    async cancel(): Promise<void> {},
+  };
+
+  return {
+    body: {
+      getReader: () => reader,
+    },
+  } as Response;
+}
+
+describe('processOpenAIStream', () => {
+  it('accumulates content and tool calls across chunks', async () => {
+    const onStream = mock(() => {});
+    const response = responseFromChunks([
+      'data: {"choices":[{"delta":{"content":"hello "}}]}\n',
+      'data: {"choices":[{"delta":{"content":"world","tool_calls":[{"index":0,"id":"call_1","function":{"name":"my_tool","arguments":"{\\"arg\\":"}}]}}]}\n',
+      'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"1}"}}]}}]}\n',
+      'data: [DONE]\n',
+    ]);
+
+    const result = await processOpenAIStream(response, { onStream });
+
+    expect(result.message.content).toBe('hello world');
+    expect(onStream).toHaveBeenCalledTimes(2);
+    expect(result.message.tool_calls?.[0]?.function?.name).toBe('my_tool');
+    expect(result.message.tool_calls?.[0]?.function?.arguments).toBe('{"arg":1}');
+  });
+
+  it('parses a final line without a newline', async () => {
+    const onStream = mock(() => {});
+    const response = responseFromChunks(['data: {"choices":[{"delta":{"content":"tail"}}]}']);
+
+    const result = await processOpenAIStream(response, { onStream });
+
+    expect(result.message.content).toBe('tail');
+    expect(onStream).toHaveBeenCalledTimes(1);
+  });
+
+  it('logs malformed JSON and continues processing', async () => {
+    const logger = {
+      log: mock(() => {}),
+      error: mock(() => {}),
+      warn: mock(() => {}),
+      info: mock(() => {}),
+    };
+    const response = responseFromChunks([
+      'data: {bad json}\n',
+      'data: {"choices":[{"delta":{"content":"ok"}}]}\n',
+      'data: [DONE]\n',
+    ]);
+
+    const result = await processOpenAIStream(response, { logger });
+
+    expect(result.message.content).toBe('ok');
+    expect(logger.warn).toHaveBeenCalledTimes(1);
+    expect(logger.warn.mock.calls[0][0]).toContain('Malformed JSON line');
+  });
+
+  it('throws error when buffer size is exceeded', async () => {
+    const response = responseFromChunks(['a'.repeat(1024 * 1024 + 1)]);
+    await expect(processOpenAIStream(response)).rejects.toThrow(
+      'LLM stream line exceed maximum size'
+    );
+  });
+
+  it('throws error when response size limit is exceeded', async () => {
+    const response = responseFromChunks([
+      `data: {"choices":[{"delta":{"content":"${'a'.repeat(600 * 1024)}"}}]}\n`,
+      `data: {"choices":[{"delta":{"content":"${'a'.repeat(500 * 1024)}"}}]}\n`,
+    ]);
+    await expect(processOpenAIStream(response)).rejects.toThrow(
+      'LLM response exceeds maximum size'
+    );
+  });
+
+  it('throws error when tool call arguments size limit is exceeded', async () => {
+    const response = responseFromChunks([
+      `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"${'a'.repeat(600 * 1024)}"}}]}}]}\n`,
+      `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"${'a'.repeat(500 * 1024)}"}}]}}]}\n`,
+    ]);
+    await expect(processOpenAIStream(response)).rejects.toThrow(
+      'LLM tool call arguments exceed maximum size'
+    );
+  });
+
+  it('handles and logs generic errors during chunk processing', async () => {
+    const logger = {
+      log: mock(() => {}),
+      error: mock(() => {}),
+      warn: mock(() => {}),
+      info: mock(() => {}),
+    };
+    // Mocking JSON.parse to throw a non-SyntaxError
+    const originalParse = JSON.parse;
+    JSON.parse = (str: string) => {
+      if (str === '{"trigger_error":true}') throw new Error('Generic error');
+      return originalParse(str);
+    };
+
+    try {
+      const response = responseFromChunks(['data: {"trigger_error":true}\n']);
+      await processOpenAIStream(response, { logger });
+      expect(logger.warn).toHaveBeenCalledTimes(1);
+      expect(logger.warn.mock.calls[0][0]).toContain(
+        'Error processing chunk: Error: Generic error'
+      );
+    } finally {
+      JSON.parse = originalParse;
+    }
+  });
+
+  it('handles errors in the final line processing', async () => {
+    const logger = {
+      log: mock(() => {}),
+      error: mock(() => {}),
+      warn: mock(() => {}),
+      info: mock(() => {}),
+    };
+    const response = responseFromChunks(['data: {bad json}']); // No newline, triggers buffer processing
+
+    await processOpenAIStream(response, { logger });
+
+    expect(logger.warn).toHaveBeenCalledTimes(1);
+    expect(logger.warn.mock.calls[0][0]).toContain('Malformed JSON line');
+  });
+
+  it('throws size limit error in final line processing', async () => {
+    const response = responseFromChunks([
+      `data: {"choices":[{"delta":{"content":"${'a'.repeat(600 * 1024)}"}}]}\n`,
+      `data: {"choices":[{"delta":{"content":"${'a'.repeat(500 * 1024)}"}}]}`,
+    ]);
+    // The first line is ok, the second line is in the final buffer and exceeds size
+    await expect(processOpenAIStream(response)).rejects.toThrow(
+      'LLM response exceeds maximum size'
+    );
+  });
+
+  it('bubbles up reader cancel errors', async () => {
+    const reader = {
+      read: async () => {
+        throw new Error('Read error');
+      },
+      cancel: async () => {
+        throw new Error('Cancel error');
+      },
+    };
+    const response = {
+      body: {
+        getReader: () => reader,
+      },
+    } as unknown as Response;
+
+    await expect(processOpenAIStream(response)).rejects.toThrow('Read error');
+  });
+});

package/src/runner/stream-utils.ts

@@ -0,0 +1,186 @@
+import { ConsoleLogger, type Logger } from '../utils/logger.ts';
+import type { LLMResponse, LLMToolCall } from './llm-adapter.ts';
+
+// Maximum response size to prevent memory exhaustion (1MB)
+const MAX_RESPONSE_SIZE = 1024 * 1024;
+const MAX_BUFFER_SIZE = MAX_RESPONSE_SIZE;
+
+type ToolCallDelta = {
+  index: number;
+  id?: string;
+  function?: {
+    name?: string;
+    arguments?: string;
+  };
+};
+
+export async function processOpenAIStream(
+  response: Response,
+  options?: { onStream?: (chunk: string) => void; logger?: Logger },
+  streamLabel = 'OpenAI'
+): Promise<LLMResponse> {
+  if (!response.body) throw new Error('Response body is null');
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  let fullContent = '';
+  const toolCalls: LLMToolCall[] = [];
+  let buffer = '';
+
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+
+      const chunk = decoder.decode(value, { stream: true });
+      buffer += chunk;
+      if (buffer.length > MAX_BUFFER_SIZE) {
+        throw new Error(`LLM stream line exceed maximum size of ${MAX_BUFFER_SIZE} bytes`);
+      }
+      const lines = buffer.split('\n');
+      // Keep the last partial line in the buffer
+      buffer = lines.pop() || '';
+
+      for (const line of lines) {
+        const trimmedLine = line.trim();
+        if (trimmedLine === '' || trimmedLine === 'data: [DONE]') continue;
+        if (!trimmedLine.startsWith('data: ')) continue;
+
+        try {
+          const data = JSON.parse(trimmedLine.slice(6));
+
+          // Handle Copilot's occasional 'choices' missing or different structure if needed,
+          // but generally they match OpenAI.
+          // Some proxies might return null delta.
+          const delta = data.choices?.[0]?.delta;
+          if (!delta) continue;
+
+          if (delta.content) {
+            if (fullContent.length + delta.content.length > MAX_RESPONSE_SIZE) {
+              throw new Error(`LLM response exceeds maximum size of ${MAX_RESPONSE_SIZE} bytes`);
+            }
+            fullContent += delta.content;
+            options?.onStream?.(delta.content);
+          }
+
+          if (delta.tool_calls) {
+            for (const tc of delta.tool_calls) {
+              const toolCall = tc as ToolCallDelta;
+              if (!toolCalls[toolCall.index]) {
+                toolCalls[toolCall.index] = {
+                  id: toolCall.id || '',
+                  type: 'function',
+                  function: { name: '', arguments: '' },
+                };
+              }
+              const existing = toolCalls[toolCall.index];
+              if (toolCall.function?.name) existing.function.name += toolCall.function.name;
+              if (toolCall.function?.arguments) {
+                if (
+                  fullContent.length +
+                    toolCalls.reduce((acc, t) => acc + (t?.function?.arguments?.length || 0), 0) +
+                    toolCall.function.arguments.length >
+                  MAX_RESPONSE_SIZE
+                ) {
+                  throw new Error(
+                    `LLM tool call arguments exceed maximum size of ${MAX_RESPONSE_SIZE} bytes`
+                  );
+                }
+                existing.function.arguments += toolCall.function.arguments;
+              }
+            }
+          }
+        } catch (e) {
+          const activeLogger = options?.logger || new ConsoleLogger();
+
+          // Rethrow size limit errors so they bubble up
+          if (e instanceof Error && e.message.toLowerCase().includes('maximum size')) {
+            throw e;
+          }
+
+          if (e instanceof SyntaxError) {
+            activeLogger.warn(
+              `[${streamLabel} Stream] Malformed JSON line: ${line.slice(0, 80)}...`
+            );
+          } else {
+            activeLogger.warn(`[${streamLabel} Stream] Error processing chunk: ${e}`);
+          }
+        }
+      }
+    }
+  } catch (error) {
+    try {
+      await reader.cancel();
+    } catch {
+      // Ignore cancel errors while bubbling up the original issue.
+    }
+    throw error;
+  }
+
+  // Final check for any remaining data in the buffer (in case of no final newline)
+  if (buffer.trim()) {
+    const trimmedLine = buffer.trim();
+    if (trimmedLine.startsWith('data: ') && trimmedLine !== 'data: [DONE]') {
+      try {
+        const data = JSON.parse(trimmedLine.slice(6));
+        const delta = data.choices?.[0]?.delta;
+        if (delta) {
+          if (delta.content) {
+            if (fullContent.length + delta.content.length > MAX_RESPONSE_SIZE) {
+              throw new Error(`LLM response exceeds maximum size of ${MAX_RESPONSE_SIZE} bytes`);
+            }
+            fullContent += delta.content;
+            options?.onStream?.(delta.content);
+          }
+          if (delta.tool_calls) {
+            // Tool calls in the very last chunk are unlikely but possible
+            for (const tc of delta.tool_calls) {
+              const toolCall = tc as ToolCallDelta;
+              if (!toolCalls[toolCall.index]) {
+                toolCalls[toolCall.index] = {
+                  id: toolCall.id || '',
+                  type: 'function',
+                  function: { name: '', arguments: '' },
+                };
+              }
+              const existing = toolCalls[toolCall.index];
+              if (toolCall.function?.name) existing.function.name += toolCall.function.name;
+              if (toolCall.function?.arguments) {
+                if (
+                  fullContent.length +
+                    toolCalls.reduce((acc, t) => acc + (t?.function?.arguments?.length || 0), 0) +
+                    toolCall.function.arguments.length >
+                  MAX_RESPONSE_SIZE
+                ) {
+                  throw new Error(
+                    `LLM tool call arguments exceed maximum size of ${MAX_RESPONSE_SIZE} bytes`
+                  );
+                }
+                existing.function.arguments += toolCall.function.arguments;
+              }
+            }
+          }
+        }
+      } catch (e) {
+        if (e instanceof Error && e.message.toLowerCase().includes('maximum size')) {
+          throw e;
+        }
+        const activeLogger = options?.logger || new ConsoleLogger();
+        if (e instanceof SyntaxError) {
+          activeLogger.warn(
+            `[${streamLabel} Stream] Malformed JSON line: ${trimmedLine.slice(0, 80)}...`
+          );
+        } else {
+          activeLogger.warn(`[${streamLabel} Stream] Error processing final line: ${e}`);
+        }
+      }
+    }
+  }
+
+  return {
+    message: {
+      role: 'assistant',
+      content: fullContent || null,
+      tool_calls: toolCalls.length > 0 ? toolCalls.filter(Boolean) : undefined,
+    },
+  };
+}

package/src/runner/workflow-runner.test.ts

@@ -457,10 +457,10 @@ describe('WorkflowRunner', () => {
 
     // Check DB status - parent should be 'paused' and step should be 'suspended'
     const db = new WorkflowDb(resumeDbPath);
-    const run = db.getRun(runId);
+    const run = await db.getRun(runId);
     expect(run?.status).toBe('paused');
 
-    const steps = db.getStepsByRun(runId);
+    const steps = await db.getStepsByRun(runId);
     const parentStep = steps.find(
       (s: { step_id: string; iteration_index: number | null }) =>
         s.step_id === 'process' && s.iteration_index === null
@@ -481,8 +481,8 @@ describe('WorkflowRunner', () => {
     expect(outputs.results).toEqual(['ok', 'ok']);
 
     const finalDb = new WorkflowDb(resumeDbPath);
-    const finalRun = finalDb.getRun(runId);
-    expect(finalRun?.status).toBe('completed');
+    const finalRun = await finalDb.getRun(runId);
+    expect(finalRun?.status).toBe('success');
     finalDb.close();
 
     if (existsSync(resumeDbPath)) rmSync(resumeDbPath);