diff --git a/src/agents/cli-output.ts b/src/agents/cli-output.ts index 05724873983..25bfd7dbdb3 100644 --- a/src/agents/cli-output.ts +++ b/src/agents/cli-output.ts @@ -1,4 +1,5 @@ import type { CliBackendConfig } from "../config/types.js"; +import { extractBalancedJsonFragments } from "../shared/balanced-json.js"; import { normalizeLowercaseStringOrEmpty } from "../shared/string-coerce.js"; import { isRecord } from "../utils.js"; @@ -39,48 +40,7 @@ function usesClaudeStreamJsonDialect(params: { } function extractJsonObjectCandidates(raw: string): string[] { - const candidates: string[] = []; - let depth = 0; - let start = -1; - let inString = false; - let escaped = false; - - for (let index = 0; index < raw.length; index += 1) { - const char = raw[index] ?? ""; - if (escaped) { - escaped = false; - continue; - } - if (char === "\\") { - if (inString) { - escaped = true; - } - continue; - } - if (char === '"') { - inString = !inString; - continue; - } - if (inString) { - continue; - } - if (char === "{") { - if (depth === 0) { - start = index; - } - depth += 1; - continue; - } - if (char === "}" && depth > 0) { - depth -= 1; - if (depth === 0 && start >= 0) { - candidates.push(raw.slice(start, index + 1)); - start = -1; - } - } - } - - return candidates; + return extractBalancedJsonFragments(raw, { openers: ["{"] }).map((fragment) => fragment.json); } function parseJsonRecordCandidates(raw: string): Record[] { diff --git a/src/agents/pi-embedded-runner/run/attempt.tool-call-argument-repair.ts b/src/agents/pi-embedded-runner/run/attempt.tool-call-argument-repair.ts index d74f6ac333a..8c4f1119d61 100644 --- a/src/agents/pi-embedded-runner/run/attempt.tool-call-argument-repair.ts +++ b/src/agents/pi-embedded-runner/run/attempt.tool-call-argument-repair.ts @@ -4,6 +4,7 @@ import { createHtmlEntityToolCallArgumentDecodingWrapper, decodeHtmlEntitiesInObject, } from "../../../plugin-sdk/provider-stream-shared.js"; +import { extractBalancedJsonPrefix } from "../../../shared/balanced-json.js"; import { normalizeProviderId } from "../../model-selection.js"; import { log } from "../logger.js"; import { wrapStreamObjectEvents } from "./stream-wrapper.js"; @@ -12,60 +13,6 @@ function isToolCallBlockType(type: unknown): boolean { return type === "toolCall" || type === "toolUse" || type === "functionCall"; } -type BalancedJsonPrefix = { - json: string; - startIndex: number; -}; - -function extractBalancedJsonPrefix(raw: string): BalancedJsonPrefix | null { - let start = 0; - while (start < raw.length) { - const char = raw[start]; - if (char === "{" || char === "[") { - break; - } - start += 1; - } - if (start >= raw.length) { - return null; - } - - let depth = 0; - let inString = false; - let escaped = false; - for (let i = start; i < raw.length; i += 1) { - const char = raw[i]; - if (char === undefined) { - break; - } - if (inString) { - if (escaped) { - escaped = false; - } else if (char === "\\") { - escaped = true; - } else if (char === '"') { - inString = false; - } - continue; - } - if (char === '"') { - inString = true; - continue; - } - if (char === "{" || char === "[") { - depth += 1; - continue; - } - if (char === "}" || char === "]") { - depth -= 1; - if (depth === 0) { - return { json: raw.slice(start, i + 1), startIndex: start }; - } - } - } - return null; -} - const MAX_TOOLCALL_REPAIR_BUFFER_CHARS = 64_000; const MAX_TOOLCALL_REPAIR_LEADING_CHARS = 96; const MAX_TOOLCALL_REPAIR_TRAILING_CHARS = 3; diff --git a/src/memory-host-sdk/host/qmd-query-parser.ts b/src/memory-host-sdk/host/qmd-query-parser.ts index dbc415eafbe..a6ec8bf049f 100644 --- a/src/memory-host-sdk/host/qmd-query-parser.ts +++ b/src/memory-host-sdk/host/qmd-query-parser.ts @@ -1,5 +1,6 @@ import { formatErrorMessage } from "../../infra/errors.js"; import { createSubsystemLogger } from "../../logging/subsystem.js"; +import { extractBalancedJsonPrefix } from "../../shared/balanced-json.js"; import { normalizeLowercaseStringOrEmpty } from "../../shared/string-coerce.js"; const log = createSubsystemLogger("memory"); @@ -34,7 +35,7 @@ export function parseQmdQueryJson(stdout: string, stderr: string): QmdQueryResul if (parsed !== null) { return parsed; } - const noisyPayload = extractFirstJsonArray(trimmedStdout); + const noisyPayload = extractBalancedJsonPrefix(trimmedStdout, { openers: ["["] })?.json; if (!noisyPayload) { throw new Error("qmd query JSON response was not an array"); } @@ -113,44 +114,3 @@ function parseQmdQueryResultArray(raw: string): QmdQueryResult[] | null { function parseQmdLineNumber(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) && value > 0 ? value : undefined; } - -function extractFirstJsonArray(raw: string): string | null { - const start = raw.indexOf("["); - if (start < 0) { - return null; - } - let depth = 0; - let inString = false; - let escaped = false; - for (let i = start; i < raw.length; i += 1) { - const char = raw[i]; - if (char === undefined) { - break; - } - if (inString) { - if (escaped) { - escaped = false; - continue; - } - if (char === "\\") { - escaped = true; - } else if (char === '"') { - inString = false; - } - continue; - } - if (char === '"') { - inString = true; - continue; - } - if (char === "[") { - depth += 1; - } else if (char === "]") { - depth -= 1; - if (depth === 0) { - return raw.slice(start, i + 1); - } - } - } - return null; -} diff --git a/src/shared/balanced-json.ts b/src/shared/balanced-json.ts new file mode 100644 index 00000000000..cfab52c4887 --- /dev/null +++ b/src/shared/balanced-json.ts @@ -0,0 +1,90 @@ +export type JsonOpeningDelimiter = "{" | "["; + +export type BalancedJsonFragment = { + json: string; + startIndex: number; + endIndex: number; +}; + +const CLOSING_DELIMITER: Record = { + "{": "}", + "[": "]", +}; + +function isJsonOpeningDelimiter( + char: string | undefined, + openers: readonly JsonOpeningDelimiter[], +): char is JsonOpeningDelimiter { + return char === "{" ? openers.includes("{") : char === "[" && openers.includes("["); +} + +export function extractBalancedJsonPrefix( + raw: string, + opts: { openers?: readonly JsonOpeningDelimiter[] } = {}, +): BalancedJsonFragment | null { + const openers = opts.openers ?? (["{", "["] as const); + let start = 0; + while (start < raw.length && !isJsonOpeningDelimiter(raw[start], openers)) { + start += 1; + } + if (start >= raw.length) { + return null; + } + + const stack: JsonOpeningDelimiter[] = []; + let inString = false; + let escaped = false; + for (let i = start; i < raw.length; i += 1) { + const char = raw[i]; + if (char === undefined) { + break; + } + if (inString) { + if (escaped) { + escaped = false; + } else if (char === "\\") { + escaped = true; + } else if (char === '"') { + inString = false; + } + continue; + } + if (char === '"') { + inString = true; + continue; + } + if (isJsonOpeningDelimiter(char, openers)) { + stack.push(char); + continue; + } + const opener = stack.at(-1); + if (opener && char === CLOSING_DELIMITER[opener]) { + stack.pop(); + if (stack.length === 0) { + return { json: raw.slice(start, i + 1), startIndex: start, endIndex: i }; + } + } + } + return null; +} + +export function extractBalancedJsonFragments( + raw: string, + opts: { openers?: readonly JsonOpeningDelimiter[] } = {}, +): BalancedJsonFragment[] { + const fragments: BalancedJsonFragment[] = []; + let offset = 0; + while (offset < raw.length) { + const fragment = extractBalancedJsonPrefix(raw.slice(offset), opts); + if (!fragment) { + break; + } + fragments.push({ + json: fragment.json, + startIndex: offset + fragment.startIndex, + endIndex: offset + fragment.endIndex, + }); + offset += fragment.endIndex + 1; + } + return fragments; +}