Skip to content

[rig-sampler] fix: extract balanced JSON arrays from prose responses #31

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pelikhan merged 1 commit into from
Jul 23, 2026
+26 −7
Merged

[rig-sampler] fix: extract balanced JSON arrays from prose responses #31

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 15 additions & 7 deletions skills/rig/rig.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1284,7 +1284,7 @@ function parseJson(text: string): { ok: true; value: unknown } | { ok: false; er
}
}

const objectText = extractBalancedJsonObject(text);
const objectText = extractBalancedJson(text);
if (objectText === undefined) {
return { ok: false, error: "No JSON value found." };
}
Expand Down Expand Up @@ -1329,16 +1329,24 @@ function extractFencedJson(text: string): string | undefined {
return text.slice(cursor, fenceEnd);
}

function extractBalancedJsonObject(text: string): string | undefined {
function extractBalancedJson(text: string): string | undefined {
const objectStart = text.indexOf("{");
if (objectStart === -1) {
const arrayStart = text.indexOf("[");
const start =
objectStart === -1 ? arrayStart :
arrayStart === -1 ? objectStart :
Math.min(objectStart, arrayStart);

@github-actions github-actions Bot Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[/diagnosing-bugs] Picking the earlier of [ and { can silently break object extraction when prose contains a bracket-delimited span (e.g. "Use tags [optional]: {\"key\":1}") — [optional] would be extracted, fail JSON.parse, and return an error instead of finding the valid object.

💡 Safer approach: try both, return the one that parses

One robust fix is to try extracting from both candidate positions and return whichever one successfully parses:

function extractBalancedJson(text: string): string | undefined {
  const candidates = [text.indexOf("{"), text.indexOf("[")]
    .filter(i => i !== -1)
    .sort((a, b) => a - b);
  for (const start of candidates) {
    const result = extractBalancedJsonFrom(text, start);
    if (result !== undefined) {
      try { JSON.parse(result); return result; } catch { /* try next */ }
    }
  }
  return undefined;
}

Alternatively, always prefer the earlier candidate but fall through to the other on parse failure — the current structure in parseJson already catches parse errors, so the fallback could live there.

The current heuristic works for the new test case but introduces a new failure mode for any prose that uses [...] notation before a JSON object.

if (start === -1) {
return undefined;
}

const openChar = text[start] as "{" | "[";
const closeChar = openChar === "{" ? "}" : "]";

let depth = 0;
let inString = false;
let escaping = false;
for (let index = objectStart; index < text.length; index += 1) {
for (let index = start; index < text.length; index += 1) {
const char = text[index]!;
if (inString) {
if (escaping) {
Expand All @@ -1356,12 +1364,12 @@ function extractBalancedJsonObject(text: string): string | undefined {
continue;
}

if (char === "{") {
if (char === openChar) {
depth += 1;
} else if (char === "}") {
} else if (char === closeChar) {
depth -= 1;
if (depth === 0) {
return text.slice(objectStart, index + 1);
return text.slice(start, index + 1);
}
}
}
Expand Down
11 changes: 11 additions & 0 deletions src/rig.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -422,6 +422,17 @@ describe("agent invocation", () => {
await expect(reviewer("go")).resolves.toEqual({ text: "hello" });
});

it("parses a JSON array embedded in surrounding text", async () => {
mocks.setSendAndWaitImpl(async () => 'Here are the results:\n["alpha","beta"]\nDone.');

const lister = agent({
name: "lister",
output: s.array(s.string),
});

await expect(lister("go")).resolves.toEqual(["alpha", "beta"]);

@github-actions github-actions Bot Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[/tdd] The new test covers the happy path but misses the case where prose contains [markdown link] or [note] text before a JSON object — exactly the regression the Math.min heuristic can introduce.

💡 Suggested regression test 
it("parses a JSON object when prose contains a bracket span before it", async () => {
  mocks.setSendAndWaitImpl(async () => "Tags [optional]: {\"key\":\"value\"}");

  const tagger = agent({
    name: "tagger",
    output: s.object({ key: s.string }),
  });

  await expect(tagger("go")).resolves.toEqual({ key: "value" });
});

This test would currently fail, making the regression detectable before merging. Adding it alongside the implementation fix would lock in correctness for both types.

});

it("retries validation failures with addon-customized repair prompts", async () => {
const prompts: string[] = [];
let calls = 0;
Expand Down
Loading