fix: copolit function call returning infinite line breaks until max_tokens limit

Author: caozhiyuan

src/lib/json-stream.ts

@@ -0,0 +1,136 @@
+export class JsonStreamValidationError extends Error {
+  constructor(message: string, options?: ErrorOptions) {
+    super(message, options)
+    this.name = "JsonStreamValidationError"
+  }
+}
+
+const FORBIDDEN_OBJECT_SEQUENCE = '","}'
+const FORBIDDEN_SEQUENCE_LOOKBACK = FORBIDDEN_OBJECT_SEQUENCE.length - 1
+
+export interface JsonStreamScanState {
+  braceBalance: number
+  bracketBalance: number
+  inString: boolean
+  isEscaped: boolean
+  recentChars: string
+}
+
+const cloneState = (state?: JsonStreamScanState): JsonStreamScanState => ({
+  braceBalance: state?.braceBalance ?? 0,
+  bracketBalance: state?.bracketBalance ?? 0,
+  inString: state?.inString ?? false,
+  isEscaped: state?.isEscaped ?? false,
+  recentChars: state?.recentChars ?? "",
+})
+
+export const validateJsonStreamChunk = (
+  previousState: JsonStreamScanState | undefined,
+  chunk: string,
+): JsonStreamScanState => {
+  const nextState = cloneState(previousState)
+
+  if (!chunk) {
+    return nextState
+  }
+
+  for (const char of chunk) {
+    checkForForbiddenObjectSequence(nextState, char)
+
+    if (nextState.inString) {
+      handleInsideStringChar(nextState, char)
+    } else {
+      handleOutsideStringChar(nextState, char)
+    }
+  }
+
+  return nextState
+}
+
+const checkForForbiddenObjectSequence = (
+  state: JsonStreamScanState,
+  char: string,
+): void => {
+  const previous = state.recentChars
+  const candidate = `${previous}${char}`
+
+  if (candidate.includes(FORBIDDEN_OBJECT_SEQUENCE)) {
+    throw new JsonStreamValidationError(
+      `JSON stream contains an unexpected sequence '${FORBIDDEN_OBJECT_SEQUENCE}' indicating malformed tool call arguments.`,
+    )
+  }
+
+  state.recentChars =
+    candidate.length > FORBIDDEN_SEQUENCE_LOOKBACK ?
+      candidate.slice(candidate.length - FORBIDDEN_SEQUENCE_LOOKBACK)
+    : candidate
+}
+
+const handleInsideStringChar = (
+  state: JsonStreamScanState,
+  char: string,
+): void => {
+  if (state.isEscaped) {
+    state.isEscaped = false
+    return
+  }
+
+  if (char === "\\") {
+    state.isEscaped = true
+    return
+  }
+
+  if (char === '"') {
+    state.inString = false
+    return
+  }
+
+  if (char === "\n" || char === "\r") {
+    throw new JsonStreamValidationError(
+      "JSON strings cannot contain unescaped newline characters.",
+    )
+  }
+}
+
+const handleOutsideStringChar = (
+  state: JsonStreamScanState,
+  char: string,
+): void => {
+  state.isEscaped = false
+
+  switch (char) {
+    case '"': {
+      state.inString = true
+      return
+    }
+    case "{": {
+      state.braceBalance += 1
+      return
+    }
+    case "}": {
+      state.braceBalance -= 1
+      if (state.braceBalance < 0) {
+        throw new JsonStreamValidationError(
+          "JSON stream has an unexpected closing '}' brace.",
+        )
+      }
+      return
+    }
+    case "[": {
+      state.bracketBalance += 1
+      return
+    }
+    case "]": {
+      state.bracketBalance -= 1
+      if (state.bracketBalance < 0) {
+        throw new JsonStreamValidationError(
+          "JSON stream has an unexpected closing ']' bracket.",
+        )
+      }
+      return
+    }
+    default: {
+      return
+    }
+  }
+}

src/routes/messages/handler.ts

@@ -8,6 +8,7 @@ import { createHandlerLogger } from "~/lib/logger"
 import { checkRateLimit } from "~/lib/rate-limit"
 import { state } from "~/lib/state"
 import {
+  buildErrorEvent,
   createResponsesStreamState,
   translateResponsesStreamEvent,
 } from "~/routes/messages/responses-stream-translation"
@@ -171,16 +172,23 @@ const handleWithResponsesApi = async (
             data: eventData,
           })
         }
+
+        if (streamState.messageCompleted) {
+          logger.debug("Message completed, ending stream")
+          break
+        }
       }
 
       if (!streamState.messageCompleted) {
         logger.warn(
-          "Responses stream ended without completion; sending fallback message_stop",
+          "Responses stream ended without completion; sending erorr event",
+        )
+        const errorEvent = buildErrorEvent(
+          "Responses stream ended without completion",
         )
-        const fallback = { type: "message_stop" as const }
         await stream.writeSSE({
-          event: fallback.type,
-          data: JSON.stringify(fallback),
+          event: errorEvent.type,
+          data: JSON.stringify(errorEvent),
         })
       }
     })

src/routes/messages/responses-stream-translation.ts

@@ -1,5 +1,10 @@
 import consola from "consola"
 
+import {
+  JsonStreamValidationError,
+  validateJsonStreamChunk,
+} from "~/lib/json-stream"
+import { type JsonStreamScanState } from "~/lib/json-stream"
 import {
   type ResponseCompletedEvent,
   type ResponseCreatedEvent,
@@ -35,6 +40,7 @@ type FunctionCallStreamState = {
   blockIndex: number
   toolCallId: string
   name: string
+  scanState?: JsonStreamScanState
 }
 
 export const createResponsesStreamState = (): ResponsesStreamState => ({
@@ -186,11 +192,39 @@ const handleFunctionCallArgumentsDelta = (
   const events = new Array<AnthropicStreamEventData>()
   const outputIndex = rawEvent.output_index
   const deltaText = rawEvent.delta
+
+  if (!deltaText) {
+    return events
+  }
+
   const blockIndex = openFunctionCallBlock(state, {
     outputIndex,
     events,
   })
 
+  const functionCallState =
+    state.functionCallStateByOutputIndex.get(outputIndex)
+  if (!functionCallState) {
+    return handleJsonValidationError(
+      new JsonStreamValidationError(
+        "Received function call arguments delta without an open tool call block.",
+      ),
+      state,
+      events,
+    )
+  }
+
+  // fix: copolit function call returning infinite line breaks until max_tokens limit
+  // "arguments": "{\"path\":\"xxx",\"pattern\":\"**/*.ts\",\"} }? Wait extra braces. Need correct. I should run? Wait overcame. Need proper JSON with pattern \"\n\n\n\n\n\n\n\n...
+  try {
+    functionCallState.scanState = validateJsonStreamChunk(
+      functionCallState.scanState,
+      deltaText,
+    )
+  } catch (error) {
+    return handleJsonValidationError(error, state, events)
+  }
+
   events.push({
     type: "content_block_delta",
     index: blockIndex,
@@ -394,6 +428,26 @@ const handleErrorEvent = (
   return [buildErrorEvent(message)]
 }
 
+const handleJsonValidationError = (
+  error: unknown,
+  state: ResponsesStreamState,
+  events: Array<AnthropicStreamEventData> = [],
+): Array<AnthropicStreamEventData> => {
+  const reason =
+    error instanceof JsonStreamValidationError ?
+      error.message
+    : "Received invalid JSON for function call arguments."
+
+  consola.error("Invalid function call arguments JSON:", error)
+
+  closeAllOpenBlocks(state, events)
+  state.messageCompleted = true
+
+  events.push(buildErrorEvent(reason))
+
+  return events
+}
+
 const messageStart = (
   state: ResponsesStreamState,
   response: ResponsesResult,
@@ -521,7 +575,7 @@ const closeAllOpenBlocks = (
   state.functionCallStateByOutputIndex.clear()
 }
 
-const buildErrorEvent = (message: string): AnthropicStreamEventData => ({
+export const buildErrorEvent = (message: string): AnthropicStreamEventData => ({
   type: "error",
   error: {
     type: "api_error",
@@ -556,6 +610,7 @@ const openFunctionCallBlock = (
       blockIndex,
       toolCallId: resolvedToolCallId,
       name: resolvedName,
+      scanState: undefined,
     }
 
     state.functionCallStateByOutputIndex.set(outputIndex, functionCallState)

tests/json-stream.test.ts

@@ -0,0 +1,32 @@
+import { describe, expect, it } from "bun:test"
+
+import {
+  JsonStreamValidationError,
+  validateJsonStreamChunk,
+} from "~/lib/json-stream"
+
+describe("validateJsonStreamChunk", () => {
+  it("throws when chunk contains unexpected sequence", () => {
+    expect(() =>
+      validateJsonStreamChunk(
+        undefined,
+        '{"pattern":"**/*.ts","path":".","}??',
+      ),
+    ).toThrow(JsonStreamValidationError)
+  })
+
+  it("throws when unexpected sequence spans multiple chunks", () => {
+    const state = validateJsonStreamChunk(
+      undefined,
+      '{"pattern":"**/*.ts","path":"."',
+    )
+
+    expect(() => validateJsonStreamChunk(state, ',"}')).toThrow(
+      JsonStreamValidationError,
+    )
+  })
+
+  it("accepts valid JSON with escaped quote sequence", () => {
+    validateJsonStreamChunk(undefined, String.raw`{"value":"\",\"}"}`)
+  })
+})