Fix Kiro

open-sse/executors/kiro.js

@@ -58,9 +58,9 @@ export class KiroExecutor extends BaseExecutor {
 
   /**
    * Transform AWS EventStream binary response to SSE text stream
+   * Using TransformStream instead of ReadableStream.pull() to avoid Workers timeout
    */
   transformEventStreamToSSE(response, model) {
-    const reader = response.body.getReader();
     let buffer = new Uint8Array(0);
     let chunkIndex = 0;
     const responseId = `chatcmpl-${Date.now()}`;
@@ -70,56 +70,30 @@ export class KiroExecutor extends BaseExecutor {
       finishEmitted: false,
       hasToolCalls: false,
       toolCallIndex: 0,
-      seenToolIds: new Map() // Map toolUseId -> index
+      seenToolIds: new Map()
     };
 
-    const stream = new ReadableStream({
-      async pull(controller) {
-        try {
-          const { done, value } = await reader.read();
-          
-          if (done) {
-            // Emit finish_reason chunk if not already sent
-            if (!state.finishEmitted) {
-              state.finishEmitted = true;
-              const finishChunk = {
-                id: responseId,
-                object: "chat.completion.chunk",
-                created,
-                model,
-                choices: [{
-                  index: 0,
-                  delta: {},
-                  finish_reason: state.hasToolCalls ? "tool_calls" : "stop"
-                }]
-              };
-              controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(finishChunk)}\n\n`));
-            }
-            
-            // Send final done message
-            controller.enqueue(new TextEncoder().encode("data: [DONE]\n\n"));
-            controller.close();
-            return;
-          }
-
+    const transformStream = new TransformStream({
+      async transform(chunk, controller) {
         // Append to buffer
-          const newBuffer = new Uint8Array(buffer.length + value.length);
+        const newBuffer = new Uint8Array(buffer.length + chunk.length);
         newBuffer.set(buffer);
-          newBuffer.set(value, buffer.length);
+        newBuffer.set(chunk, buffer.length);
         buffer = newBuffer;
 
         // Parse events from buffer
-          while (buffer.length >= 16) {
+        let iterations = 0;
+        const maxIterations = 1000;
+        while (buffer.length >= 16 && iterations < maxIterations) {
+          iterations++;
           const view = new DataView(buffer.buffer, buffer.byteOffset);
           const totalLength = view.getUint32(0, false);
 
-            if (totalLength < 16 || buffer.length < totalLength) break;
+          if (totalLength < 16 || totalLength > buffer.length || buffer.length < totalLength) break;
 
-            // Extract event
           const eventData = buffer.slice(0, totalLength);
           buffer = buffer.slice(totalLength);
 
-            // Parse event headers and payload
           const event = parseEventFrame(eventData);
           if (!event) continue;
 
@@ -163,25 +137,19 @@ export class KiroExecutor extends BaseExecutor {
 
           // Handle toolUseEvent
           if (eventType === "toolUseEvent" && event.payload) {
-              
-              state.hasToolCalls = true; // Track that we have tool calls
+            state.hasToolCalls = true;
             const toolUse = event.payload;
-              
-              // AWS Kiro sends toolUse as object or array
-              // If it's an array, process each tool separately
             const toolUses = Array.isArray(toolUse) ? toolUse : [toolUse];
 
             for (const singleToolUse of toolUses) {
               const toolCallId = singleToolUse.toolUseId || `call_${Date.now()}`;
               const toolName = singleToolUse.name || "";
-                const toolInput = singleToolUse.input; // Can be undefined, string, or object
+              const toolInput = singleToolUse.input;
 
-                // Get or assign tool call index
               let toolIndex;
               const isNewTool = !state.seenToolIds.has(toolCallId);
 
               if (isNewTool) {
-                  // NEW TOOL: Create start chunk
                 toolIndex = state.toolCallIndex++;
                 state.seenToolIds.set(toolCallId, toolIndex);
 
@@ -210,23 +178,17 @@ export class KiroExecutor extends BaseExecutor {
                 chunkIndex++;
                 controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(startChunk)}\n\n`));
               } else {
-                  // EXISTING TOOL: Get its index
                 toolIndex = state.seenToolIds.get(toolCallId);
               }
 
-                // Emit arguments chunk if input exists
-                // AWS Kiro streams input as: undefined (first event) → string chunks
               if (toolInput !== undefined) {
                 let argumentsStr;
 
                 if (typeof toolInput === 'string') {
-                    // AWS Kiro sends partial JSON as STRING
                   argumentsStr = toolInput;
                 } else if (typeof toolInput === 'object') {
-                    // Fallback: if it's an object, stringify it
                   argumentsStr = JSON.stringify(toolInput);
                 } else {
-                    // Skip if not string or object
                   continue;
                 }
 
@@ -271,12 +233,9 @@ export class KiroExecutor extends BaseExecutor {
             controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(chunk)}\n\n`));
           }
 
-            // Detect end of stream: meteringEvent + contextUsageEvent usually come last
-            // Kiro doesn't always send messageStopEvent, so we need to detect completion
+          // Detect end of stream
           if ((eventType === "meteringEvent" || eventType === "contextUsageEvent") && !state.endDetected) {
             state.endDetected = true;
-              // Schedule finish chunk emission after a short delay
-              setTimeout(() => {
             if (!state.finishEmitted) {
               state.finishEmitted = true;
               const finishChunk = {
@@ -292,20 +251,41 @@ export class KiroExecutor extends BaseExecutor {
               };
               controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(finishChunk)}\n\n`));
             }
-              }, 100); // 100ms delay to check for more events
           }
         }
-        } catch (error) {
-          controller.error(error);
+
+        if (iterations >= maxIterations) {
+          console.warn("[Kiro] Max iterations reached in event parsing");
         }
       },
-      cancel() {
-        reader.cancel();
+
+      flush(controller) {
+        // Emit finish chunk if not already sent
+        if (!state.finishEmitted) {
+          state.finishEmitted = true;
+          const finishChunk = {
+            id: responseId,
+            object: "chat.completion.chunk",
+            created,
+            model,
+            choices: [{
+              index: 0,
+              delta: {},
+              finish_reason: state.hasToolCalls ? "tool_calls" : "stop"
+            }]
+          };
+          controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(finishChunk)}\n\n`));
+        }
+
+        // Send final done message
+        controller.enqueue(new TextEncoder().encode("data: [DONE]\n\n"));
       }
     });
 
-    // Create new response with SSE headers
-    return new Response(stream, {
+    // Pipe response body through transform stream
+    const transformedStream = response.body.pipeThrough(transformStream);
+
+    return new Response(transformedStream, {
       status: response.status,
       statusText: response.statusText,
       headers: {