Skip to content

Commit f006678

Browse files
steipetesteipete
committed
refactor: share balanced json extraction
1 parent 655e0be commit f006678

4 files changed

Lines changed: 95 additions & 138 deletions

File tree

‎src/agents/cli-output.ts‎

Lines changed: 2 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { CliBackendConfig } from "../config/types.js";
2+
import { extractBalancedJsonFragments } from "../shared/balanced-json.js";
23
import { normalizeLowercaseStringOrEmpty } from "../shared/string-coerce.js";
34
import { isRecord } from "../utils.js";
45

@@ -39,48 +40,7 @@ function usesClaudeStreamJsonDialect(params: {
3940
}
4041

4142
function extractJsonObjectCandidates(raw: string): string[] {
42-
const candidates: string[] = [];
43-
let depth = 0;
44-
let start = -1;
45-
let inString = false;
46-
let escaped = false;
47-
48-
for (let index = 0; index < raw.length; index += 1) {
49-
const char = raw[index] ?? "";
50-
if (escaped) {
51-
escaped = false;
52-
continue;
53-
}
54-
if (char === "\\") {
55-
if (inString) {
56-
escaped = true;
57-
}
58-
continue;
59-
}
60-
if (char === '"') {
61-
inString = !inString;
62-
continue;
63-
}
64-
if (inString) {
65-
continue;
66-
}
67-
if (char === "{") {
68-
if (depth === 0) {
69-
start = index;
70-
}
71-
depth += 1;
72-
continue;
73-
}
74-
if (char === "}" && depth > 0) {
75-
depth -= 1;
76-
if (depth === 0 && start >= 0) {
77-
candidates.push(raw.slice(start, index + 1));
78-
start = -1;
79-
}
80-
}
81-
}
82-
83-
return candidates;
43+
return extractBalancedJsonFragments(raw, { openers: ["{"] }).map((fragment) => fragment.json);
8444
}
8545

8646
function parseJsonRecordCandidates(raw: string): Record<string, unknown>[] {

‎src/agents/pi-embedded-runner/run/attempt.tool-call-argument-repair.ts‎

Lines changed: 1 addition & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import {
44
createHtmlEntityToolCallArgumentDecodingWrapper,
55
decodeHtmlEntitiesInObject,
66
} from "../../../plugin-sdk/provider-stream-shared.js";
7+
import { extractBalancedJsonPrefix } from "../../../shared/balanced-json.js";
78
import { normalizeProviderId } from "../../model-selection.js";
89
import { log } from "../logger.js";
910
import { wrapStreamObjectEvents } from "./stream-wrapper.js";
@@ -12,60 +13,6 @@ function isToolCallBlockType(type: unknown): boolean {
1213
return type === "toolCall" || type === "toolUse" || type === "functionCall";
1314
}
1415

15-
type BalancedJsonPrefix = {
16-
json: string;
17-
startIndex: number;
18-
};
19-
20-
function extractBalancedJsonPrefix(raw: string): BalancedJsonPrefix | null {
21-
let start = 0;
22-
while (start < raw.length) {
23-
const char = raw[start];
24-
if (char === "{" || char === "[") {
25-
break;
26-
}
27-
start += 1;
28-
}
29-
if (start >= raw.length) {
30-
return null;
31-
}
32-
33-
let depth = 0;
34-
let inString = false;
35-
let escaped = false;
36-
for (let i = start; i < raw.length; i += 1) {
37-
const char = raw[i];
38-
if (char === undefined) {
39-
break;
40-
}
41-
if (inString) {
42-
if (escaped) {
43-
escaped = false;
44-
} else if (char === "\\") {
45-
escaped = true;
46-
} else if (char === '"') {
47-
inString = false;
48-
}
49-
continue;
50-
}
51-
if (char === '"') {
52-
inString = true;
53-
continue;
54-
}
55-
if (char === "{" || char === "[") {
56-
depth += 1;
57-
continue;
58-
}
59-
if (char === "}" || char === "]") {
60-
depth -= 1;
61-
if (depth === 0) {
62-
return { json: raw.slice(start, i + 1), startIndex: start };
63-
}
64-
}
65-
}
66-
return null;
67-
}
68-
6916
const MAX_TOOLCALL_REPAIR_BUFFER_CHARS = 64_000;
7017
const MAX_TOOLCALL_REPAIR_LEADING_CHARS = 96;
7118
const MAX_TOOLCALL_REPAIR_TRAILING_CHARS = 3;

‎src/memory-host-sdk/host/qmd-query-parser.ts‎

Lines changed: 2 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { formatErrorMessage } from "../../infra/errors.js";
22
import { createSubsystemLogger } from "../../logging/subsystem.js";
3+
import { extractBalancedJsonPrefix } from "../../shared/balanced-json.js";
34
import { normalizeLowercaseStringOrEmpty } from "../../shared/string-coerce.js";
45

56
const log = createSubsystemLogger("memory");
@@ -34,7 +35,7 @@ export function parseQmdQueryJson(stdout: string, stderr: string): QmdQueryResul
3435
if (parsed !== null) {
3536
return parsed;
3637
}
37-
const noisyPayload = extractFirstJsonArray(trimmedStdout);
38+
const noisyPayload = extractBalancedJsonPrefix(trimmedStdout, { openers: ["["] })?.json;
3839
if (!noisyPayload) {
3940
throw new Error("qmd query JSON response was not an array");
4041
}
@@ -113,44 +114,3 @@ function parseQmdQueryResultArray(raw: string): QmdQueryResult[] | null {
113114
function parseQmdLineNumber(value: unknown): number | undefined {
114115
return typeof value === "number" && Number.isFinite(value) && value > 0 ? value : undefined;
115116
}
116-
117-
function extractFirstJsonArray(raw: string): string | null {
118-
const start = raw.indexOf("[");
119-
if (start < 0) {
120-
return null;
121-
}
122-
let depth = 0;
123-
let inString = false;
124-
let escaped = false;
125-
for (let i = start; i < raw.length; i += 1) {
126-
const char = raw[i];
127-
if (char === undefined) {
128-
break;
129-
}
130-
if (inString) {
131-
if (escaped) {
132-
escaped = false;
133-
continue;
134-
}
135-
if (char === "\\") {
136-
escaped = true;
137-
} else if (char === '"') {
138-
inString = false;
139-
}
140-
continue;
141-
}
142-
if (char === '"') {
143-
inString = true;
144-
continue;
145-
}
146-
if (char === "[") {
147-
depth += 1;
148-
} else if (char === "]") {
149-
depth -= 1;
150-
if (depth === 0) {
151-
return raw.slice(start, i + 1);
152-
}
153-
}
154-
}
155-
return null;
156-
}

‎src/shared/balanced-json.ts‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
export type JsonOpeningDelimiter = "{" | "[";
2+
3+
export type BalancedJsonFragment = {
4+
json: string;
5+
startIndex: number;
6+
endIndex: number;
7+
};
8+
9+
const CLOSING_DELIMITER: Record<JsonOpeningDelimiter, "}" | "]"> = {
10+
"{": "}",
11+
"[": "]",
12+
};
13+
14+
function isJsonOpeningDelimiter(
15+
char: string | undefined,
16+
openers: readonly JsonOpeningDelimiter[],
17+
): char is JsonOpeningDelimiter {
18+
return char === "{" ? openers.includes("{") : char === "[" && openers.includes("[");
19+
}
20+
21+
export function extractBalancedJsonPrefix(
22+
raw: string,
23+
opts: { openers?: readonly JsonOpeningDelimiter[] } = {},
24+
): BalancedJsonFragment | null {
25+
const openers = opts.openers ?? (["{", "["] as const);
26+
let start = 0;
27+
while (start < raw.length && !isJsonOpeningDelimiter(raw[start], openers)) {
28+
start += 1;
29+
}
30+
if (start >= raw.length) {
31+
return null;
32+
}
33+
34+
const stack: JsonOpeningDelimiter[] = [];
35+
let inString = false;
36+
let escaped = false;
37+
for (let i = start; i < raw.length; i += 1) {
38+
const char = raw[i];
39+
if (char === undefined) {
40+
break;
41+
}
42+
if (inString) {
43+
if (escaped) {
44+
escaped = false;
45+
} else if (char === "\\") {
46+
escaped = true;
47+
} else if (char === '"') {
48+
inString = false;
49+
}
50+
continue;
51+
}
52+
if (char === '"') {
53+
inString = true;
54+
continue;
55+
}
56+
if (isJsonOpeningDelimiter(char, openers)) {
57+
stack.push(char);
58+
continue;
59+
}
60+
const opener = stack.at(-1);
61+
if (opener && char === CLOSING_DELIMITER[opener]) {
62+
stack.pop();
63+
if (stack.length === 0) {
64+
return { json: raw.slice(start, i + 1), startIndex: start, endIndex: i };
65+
}
66+
}
67+
}
68+
return null;
69+
}
70+
71+
export function extractBalancedJsonFragments(
72+
raw: string,
73+
opts: { openers?: readonly JsonOpeningDelimiter[] } = {},
74+
): BalancedJsonFragment[] {
75+
const fragments: BalancedJsonFragment[] = [];
76+
let offset = 0;
77+
while (offset < raw.length) {
78+
const fragment = extractBalancedJsonPrefix(raw.slice(offset), opts);
79+
if (!fragment) {
80+
break;
81+
}
82+
fragments.push({
83+
json: fragment.json,
84+
startIndex: offset + fragment.startIndex,
85+
endIndex: offset + fragment.endIndex,
86+
});
87+
offset += fragment.endIndex + 1;
88+
}
89+
return fragments;
90+
}

0 commit comments

Comments
 (0)