From 52e039ce71ab1045bbf22110b556801c993a0f85 Mon Sep 17 00:00:00 2001 From: Aishanee Shah Date: Sat, 11 Apr 2026 03:30:02 +0000 Subject: [PATCH] fix(watcher): robustly extract JSON from sub-agent output --- packages/core/src/core/client.ts | 5 +- packages/core/src/core/client_watcher.test.ts | 88 +++++++++++++++++++ packages/core/src/utils/jsonUtils.test.ts | 57 ++++++++++++ packages/core/src/utils/jsonUtils.ts | 48 ++++++++++ 4 files changed, 197 insertions(+), 1 deletion(-) create mode 100644 packages/core/src/utils/jsonUtils.test.ts create mode 100644 packages/core/src/utils/jsonUtils.ts diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 1471d8f6eb..d7a5ba3855 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -48,6 +48,7 @@ import { LoopDetectionService } from '../services/loopDetectionService.js'; import { ChatCompressionService } from '../context/chatCompressionService.js'; import { AgentHistoryProvider } from '../context/agentHistoryProvider.js'; import { isSubagentProgress, type WatcherProgress } from '../agents/types.js'; +import { extractAndParseJson } from '../utils/jsonUtils.js'; import { WatcherReportSchema } from '../agents/watcher-agent.js'; import { ideContextStore } from '../ide/ideContext.js'; import { @@ -1409,7 +1410,9 @@ export class GeminiClient { try { const rawOutput = result.returnDisplay.result; debugLogger.log(`[Watcher] Raw content response: ${rawOutput}`); - const parsed = WatcherReportSchema.parse(JSON.parse(rawOutput)); + const parsed = WatcherReportSchema.parse( + extractAndParseJson(rawOutput), + ); // Internally write the status report to avoid requiring user permission const projectTempDir = this.config.storage.getProjectTempDir(); diff --git a/packages/core/src/core/client_watcher.test.ts b/packages/core/src/core/client_watcher.test.ts index a26ee9de6a..ac3f577985 100644 --- a/packages/core/src/core/client_watcher.test.ts +++ b/packages/core/src/core/client_watcher.test.ts @@ -271,4 +271,92 @@ describe('GeminiClient Watcher Integration', () => { client.dispose(); expect(fs.existsSync(statusFilePath)).toBe(false); }); + + it('should robustly handle messy subagent output with conversational filler and markdown', async () => { + vi.spyOn(config, 'isExperimentalWatcherEnabled').mockReturnValue(true); + vi.spyOn(config, 'getExperimentalWatcherInterval').mockReturnValue(1); + vi.spyOn(config, 'getApprovalMode').mockReturnValue(ApprovalMode.DEFAULT); + + const reportData = { + userDirections: 'Messy test direction', + progressSummary: 'Messy progress', + evaluation: 'ON_TRACK', + }; + + const messyOutput = ` +Subagent "watcher" finished with result: +\`\`\`json +${JSON.stringify(reportData, null, 2)} +\`\`\` +I hope this status update is helpful! + `; + + const mockWatcherTool = { + build: vi.fn().mockReturnValue({ + execute: vi.fn().mockResolvedValue({ + llmContent: [{ text: 'Subagent finished' }], + returnDisplay: { + isSubagentProgress: true, + agentName: 'watcher', + recentActivity: [], + state: 'completed', + result: messyOutput, + }, + }), + }), + name: 'watcher', + displayName: 'Watcher', + description: 'Watcher tool', + inputConfig: { inputSchema: {} }, + outputConfig: { outputName: 'report', schema: {} }, + }; + + const mockToolRegistry = { + getFunctionDeclarations: vi.fn().mockReturnValue([]), + getTool: vi.fn().mockImplementation((name) => { + if (name === 'watcher') return mockWatcherTool; + return undefined; + }), + getAllToolNames: vi.fn().mockReturnValue(['watcher']), + sortTools: vi.fn(), + discoverAllTools: vi.fn(), + }; + + const clientAccess = client as unknown as { + context: AgentLoopContext; + }; + + Object.defineProperty(clientAccess.context, 'toolRegistry', { + get: () => mockToolRegistry, + configurable: true, + }); + + ( + clientAccess.context as unknown as { agentRegistry: unknown } + ).agentRegistry = { + getAllDefinitions: vi.fn().mockReturnValue([]), + getDefinition: vi.fn().mockReturnValue(undefined), + initialize: vi.fn().mockResolvedValue(undefined), + }; + + await config.storage.initialize(); + await client.initialize(); + + const signal = new AbortController().signal; + const generator = client.sendMessageStream( + [{ text: 'test turn' }], + signal, + 'test-prompt', + ); + for await (const _ of generator) { + /* consume */ + } + + const projectTempDir = config.storage.getProjectTempDir(); + const statusFilePath = path.join(projectTempDir, 'watcher_status.md'); + expect(fs.existsSync(statusFilePath)).toBe(true); + const content = fs.readFileSync(statusFilePath, 'utf-8'); + expect(content).toContain('Messy test direction'); + expect(content).toContain('Messy progress'); + }); }); diff --git a/packages/core/src/utils/jsonUtils.test.ts b/packages/core/src/utils/jsonUtils.test.ts new file mode 100644 index 0000000000..77d928f6e0 --- /dev/null +++ b/packages/core/src/utils/jsonUtils.test.ts @@ -0,0 +1,57 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { extractAndParseJson } from './jsonUtils.js'; + +describe('extractAndParseJson', () => { + it('should parse pure JSON objects', () => { + const input = '{"key": "value"}'; + expect(extractAndParseJson(input)).toEqual({ key: 'value' }); + }); + + it('should parse pure JSON arrays', () => { + const input = '[1, 2, 3]'; + expect(extractAndParseJson(input)).toEqual([1, 2, 3]); + }); + + it('should extract JSON from conversational filler (The Watcher Bug)', () => { + const input = + 'Subagent "watcher" finished with result: {"userDirections": "Keep going", "progressSummary": "Done", "evaluation": "ON_TRACK"}'; + expect(extractAndParseJson(input)).toEqual({ + userDirections: 'Keep going', + progressSummary: 'Done', + evaluation: 'ON_TRACK', + }); + }); + + it('should extract JSON from markdown code blocks', () => { + const input = + 'Here is the report:\n```json\n{"status": "ok"}\n```\nHope this helps!'; + expect(extractAndParseJson(input)).toEqual({ status: 'ok' }); + }); + + it('should handle leading and trailing filler simultaneously', () => { + const input = 'PREFIX {"a": 1} SUFFIX'; + expect(extractAndParseJson(input)).toEqual({ a: 1 }); + }); + + it('should handle nested braces correctly', () => { + const input = 'Result: {"outer": {"inner": 42}} - end'; + expect(extractAndParseJson(input)).toEqual({ outer: { inner: 42 } }); + }); + + it('should fallback to full string if no braces found (for numbers/booleans/strings)', () => { + expect(extractAndParseJson('true')).toBe(true); + expect(extractAndParseJson('123')).toBe(123); + expect(extractAndParseJson('"just a string"')).toBe('just a string'); + }); + + it('should throw SyntaxError for truly invalid JSON', () => { + expect(() => extractAndParseJson('not json at all')).toThrow(SyntaxError); + expect(() => extractAndParseJson('{"unfinished": ')).toThrow(SyntaxError); + }); +}); diff --git a/packages/core/src/utils/jsonUtils.ts b/packages/core/src/utils/jsonUtils.ts new file mode 100644 index 0000000000..12b31f651e --- /dev/null +++ b/packages/core/src/utils/jsonUtils.ts @@ -0,0 +1,48 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Attempts to extract and parse a JSON object or array from a string that may + * contain conversational filler or markdown code blocks. + * + * @param text The text to extract JSON from. + * @returns The parsed JSON object or array as unknown. + * @throws SyntaxError if no valid JSON is found. + */ +export function extractAndParseJson(text: string): unknown { + const firstBrace = text.indexOf('{'); + const firstBracket = text.indexOf('['); + + let start = -1; + let end = -1; + + // Determine if we should look for an object or an array first based on which starts earlier. + if (firstBrace !== -1 && (firstBracket === -1 || firstBrace < firstBracket)) { + start = firstBrace; + end = text.lastIndexOf('}'); + } else if (firstBracket !== -1) { + start = firstBracket; + end = text.lastIndexOf(']'); + } + + if (start === -1 || end === -1 || end <= start) { + // Fallback: try parsing the whole trimmed text. + return JSON.parse(text.trim()); + } + + const cleanJson = text.substring(start, end + 1); + try { + return JSON.parse(cleanJson); + } catch (e) { + // If extraction failed to produce valid JSON (e.g. mismatched braces), + // try the whole text as a last resort. + try { + return JSON.parse(text.trim()); + } catch { + throw e; + } + } +}