From 904b81676c6ccaaa1fa4ca239656d547d559a159 Mon Sep 17 00:00:00 2001 From: Vincent Driessen Date: Mon, 1 Sep 2025 15:47:17 +0200 Subject: [PATCH] Optimize AI tool call streaming with incremental JSON parser (#2617) --- CHANGELOG.md | 5 + .../app/toolcall-streaming/[chatId]/page.tsx | 403 +++++++++++++++++- .../liveblocks-core/src/__tests__/ai.test.ts | 89 ++-- .../src/lib/IncrementalJsonParser.ts | 251 +++++++++++ .../__tests__/IncrementalJsonParser.test.ts | 403 ++++++++++++++++++ .../__tests__/parsePartialJsonObject.test.ts | 154 ------- packages/liveblocks-core/src/types/ai.ts | 58 ++- 7 files changed, 1145 insertions(+), 218 deletions(-) create mode 100644 packages/liveblocks-core/src/lib/IncrementalJsonParser.ts create mode 100644 packages/liveblocks-core/src/lib/__tests__/IncrementalJsonParser.test.ts delete mode 100644 packages/liveblocks-core/src/lib/__tests__/parsePartialJsonObject.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ae7f6982be..bd59091f752 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ ## vNEXT (not yet published) +### `@liveblocks/core` + +- Optimized partial JSON parser for improved tool invocation streaming + performance + ## v3.5.1 ### `@liveblocks/react-tiptap` diff --git a/e2e/next-ai-kitchen-sink/app/toolcall-streaming/[chatId]/page.tsx b/e2e/next-ai-kitchen-sink/app/toolcall-streaming/[chatId]/page.tsx index 67318705528..e0e83db53d5 100644 --- a/e2e/next-ai-kitchen-sink/app/toolcall-streaming/[chatId]/page.tsx +++ b/e2e/next-ai-kitchen-sink/app/toolcall-streaming/[chatId]/page.tsx @@ -1,6 +1,6 @@ "use client"; -import { use } from "react"; +import { use, useRef } from "react"; import { defineAiTool } from "@liveblocks/core"; import { ClientSideSuspense, @@ -13,7 +13,68 @@ import { AiTool, } from "@liveblocks/react-ui"; -export default function HtmlStreamingPage({ params }: { params: Promise<{ chatId: string }> }) { +function useRenderCount() { + const ref = useRef(0); + return ++ref.current; +} + +function useStageTimer(invocationId: string) { + const timersRef = useRef( + new Map< + string, + { + startTime: number | null; + executingTime: number | null; + receivingToExecutingTime: number | null; + executingToExecutedTime: number | null; + } + >() + ); + + const getTimer = (id: string) => { + if (!timersRef.current.has(id)) { + timersRef.current.set(id, { + startTime: null, + executingTime: null, + receivingToExecutingTime: null, + executingToExecutedTime: null, + }); + } + return timersRef.current.get(id)!; + }; + + const timer = getTimer(invocationId); + + return { + markReceivingStart: () => { + // Only start timing on the very first time we see receiving + if (timer.startTime === null) { + timer.startTime = Date.now(); + timer.receivingToExecutingTime = null; + timer.executingToExecutedTime = null; + } + }, + markExecuting: () => { + if (timer.startTime && timer.receivingToExecutingTime === null) { + timer.receivingToExecutingTime = Date.now() - timer.startTime; + timer.executingTime = Date.now(); + } + }, + markExecuted: () => { + if (timer.executingTime && timer.executingToExecutedTime === null) { + timer.executingToExecutedTime = Date.now() - timer.executingTime; + } + }, + getReceivingToExecutingTime: () => timer.receivingToExecutingTime, + getExecutingToExecutedTime: () => timer.executingToExecutedTime, + }; +} + +export default function HtmlStreamingPage({ + params, +}: { + params: Promise<{ chatId: string }>; +}) { const { chatId } = use(params); return (
@@ -53,9 +114,340 @@ export default function HtmlStreamingPage({ params }: { params: Promise<{ chatId execute: () => { return { data: { success: true } }; }, - render: () => { + render: (props) => { + /* eslint-disable react-hooks/rules-of-hooks */ + const renderCount = useRenderCount(); + const stageTimer = useStageTimer(props.invocationId); + /* eslint-enable react-hooks/rules-of-hooks */ + + // Track stage transitions + if (props.stage === "receiving") { + stageTimer.markReceivingStart(); + } else if (props.stage === "executing") { + stageTimer.markExecuting(); + } else if (props.stage === "executed") { + stageTimer.markExecuted(); + } + + const receivingToExecutingTime = + stageTimer.getReceivingToExecutingTime(); + const executingToExecutedTime = + stageTimer.getExecutingToExecutedTime(); + + return ( + +
+ Stage: {props.stage} +
+
+ Render count: {renderCount} +
+ {receivingToExecutingTime !== null && ( +
+ Receiving → Executing: {receivingToExecutingTime}ms +
+ )} + {executingToExecutedTime !== null && ( +
+ Executing → Executed: {executingToExecutedTime}ms +
+ )} + + + ); + }, + }), + largeDataProcessor: defineAiTool()({ + description: + "Process large amounts of data with many parameters for stress testing", + parameters: { + type: "object", + properties: { + config: { + type: "object", + description: "Configuration object", + properties: { + environment: { + type: "string", + description: "Environment setting", + }, + debug: { type: "boolean", description: "Debug mode" }, + timeout: { + type: "number", + description: "Timeout in milliseconds", + }, + retries: { + type: "number", + description: "Number of retries", + }, + }, + }, + metadata: { + type: "object", + description: "Metadata information", + properties: { + version: { + type: "string", + description: "Version number", + }, + author: { type: "string", description: "Author name" }, + timestamp: { + type: "string", + description: "Creation timestamp", + }, + tags: { + type: "array", + items: { type: "string" }, + description: "Tags list", + }, + }, + }, + dataPoints: { + type: "array", + description: "Array of data points to process", + items: { + type: "object", + properties: { + id: { + type: "string", + description: "Unique identifier", + }, + value: { + type: "number", + description: "Numeric value", + }, + label: { + type: "string", + description: "Human readable label", + }, + category: { + type: "string", + description: "Category classification", + }, + properties: { + type: "object", + description: "Additional properties", + additionalProperties: true, + }, + }, + }, + }, + filters: { + type: "array", + description: "Filtering criteria", + items: { + type: "object", + properties: { + field: { + type: "string", + description: "Field to filter on", + }, + operator: { + type: "string", + enum: [ + "equals", + "contains", + "greater_than", + "less_than", + ], + }, + value: { + type: "string", + description: "Filter value", + }, + }, + }, + }, + transformations: { + type: "array", + description: "Data transformations to apply", + items: { + type: "object", + properties: { + type: { + type: "string", + enum: ["map", "filter", "reduce", "sort"], + }, + field: { + type: "string", + description: "Target field", + }, + operation: { + type: "string", + description: "Operation to perform", + }, + parameters: { + type: "object", + additionalProperties: true, + }, + }, + }, + }, + outputFormat: { + type: "object", + description: "Output formatting options", + properties: { + format: { + type: "string", + enum: ["json", "csv", "xml", "yaml"], + }, + compression: { + type: "boolean", + description: "Enable compression", + }, + encryption: { + type: "boolean", + description: "Enable encryption", + }, + headers: { + type: "array", + items: { type: "string" }, + description: "Custom headers", + }, + }, + }, + performance: { + type: "object", + description: "Performance tuning options", + properties: { + batchSize: { + type: "number", + description: "Processing batch size", + }, + parallelism: { + type: "number", + description: "Parallel processing threads", + }, + caching: { + type: "boolean", + description: "Enable result caching", + }, + optimization: { + type: "string", + enum: ["speed", "memory", "balanced"], + }, + }, + }, + validation: { + type: "object", + description: "Data validation rules", + properties: { + strict: { + type: "boolean", + description: "Enable strict validation", + }, + rules: { + type: "array", + items: { + type: "object", + properties: { + field: { type: "string" }, + type: { type: "string" }, + required: { type: "boolean" }, + pattern: { type: "string" }, + }, + }, + }, + }, + }, + additionalContext: { + type: "string", + description: + "Any additional context or instructions for processing the data", + }, + }, + required: ["config", "dataPoints", "outputFormat"], + additionalProperties: false, + }, + execute: () => { + return { data: { processed: true, recordCount: 0 } }; + }, + render: (props) => { + /* eslint-disable react-hooks/rules-of-hooks */ + const renderCount = useRenderCount(); + const stageTimer = useStageTimer(props.invocationId); + /* eslint-enable react-hooks/rules-of-hooks */ + + // Track stage transitions + if (props.stage === "receiving") { + stageTimer.markReceivingStart(); + } else if (props.stage === "executing") { + stageTimer.markExecuting(); + } else if (props.stage === "executed") { + stageTimer.markExecuted(); + } + + const receivingToExecutingTime = + stageTimer.getReceivingToExecutingTime(); + const executingToExecutedTime = + stageTimer.getExecutingToExecutedTime(); + return ( +
+ Stage: {props.stage} +
+
+ Render count: {renderCount} +
+ {receivingToExecutingTime !== null && ( +
+ Receiving → Executing: {receivingToExecutingTime}ms +
+ )} + {executingToExecutedTime !== null && ( +
+ Executing → Executed: {executingToExecutedTime}ms +
+ )} ); @@ -89,6 +481,11 @@ const CHAT_SUGGESTIONS = [ label: "Contact page", message: "Build a contact page with a form and company information", }, + { + label: "🔥 Stress Test (Large Tool Call)", + message: + "Process a dataset with 500 records, apply complex transformations, and generate a comprehensive analytics report with detailed filtering, validation rules, and performance optimization settings. Include metadata tracking, custom headers, batch processing configuration, and extensive data point analysis.", + }, ]; function AiChatEmptyComponent({ chatId }: AiChatComponentsEmptyProps) { diff --git a/packages/liveblocks-core/src/__tests__/ai.test.ts b/packages/liveblocks-core/src/__tests__/ai.test.ts index 485003d2858..e6d0ad4bb20 100644 --- a/packages/liveblocks-core/src/__tests__/ai.test.ts +++ b/packages/liveblocks-core/src/__tests__/ai.test.ts @@ -11,7 +11,10 @@ import type { AiToolInvocationDelta, AiToolInvocationStreamStart, } from "../types/ai"; -import { patchContentWithDelta } from "../types/ai"; +import { + createReceivingToolInvocation, + patchContentWithDelta, +} from "../types/ai"; describe("KnowledgeStack", () => { test("should be empty by default", () => { @@ -157,6 +160,45 @@ describe("KnowledgeStack", () => { }); }); +describe("createReceivingToolInvocation", () => { + test("creates receiving tool invocation with empty args", () => { + const tool = createReceivingToolInvocation("inv-test", "testTool"); + + expect(tool.type).toBe("tool-invocation"); + expect(tool.stage).toBe("receiving"); + expect(tool.invocationId).toBe("inv-test"); + expect(tool.name).toBe("testTool"); + expect(tool.partialArgsText).toBe(""); + expect(tool.partialArgs).toEqual({}); + }); + + test("creates receiving tool invocation with partial args", () => { + const tool = createReceivingToolInvocation( + "inv-123", + "search", + '{"query": "test"}' + ); + + expect(tool.partialArgsText).toBe('{"query": "test"}'); + expect(tool.partialArgs).toEqual({ query: "test" }); + }); + + test("allows appending deltas via __appendDelta", () => { + const tool = createReceivingToolInvocation( + "inv-456", + "calculator", + '{"expr": "2+' + ); + + expect(tool.partialArgs).toEqual({ expr: "2+" }); + + tool.__appendDelta?.('2"}'); + + expect(tool.partialArgsText).toBe('{"expr": "2+2"}'); + expect(tool.partialArgs).toEqual({ expr: "2+2" }); + }); +}); + describe("patchContentWithDelta", () => { describe("text-delta", () => { test("appends to existing text part", () => { @@ -623,14 +665,7 @@ describe("patchContentWithDelta", () => { test("replaces receiving tool with executing tool (same invocationId)", () => { const content: AiAssistantContentPart[] = [ - { - type: "tool-invocation", - stage: "receiving", - invocationId: "inv-123", - name: "search", - partialArgsText: '{"query": "par"}', - partialArgs: { dummy: "not used in this test" }, - }, + createReceivingToolInvocation("inv-123", "search", '{"query": "par"}'), ]; const delta: AiExecutingToolInvocationPart = { @@ -673,14 +708,7 @@ describe("patchContentWithDelta", () => { test("replaces tool in middle of content array", () => { const content: AiAssistantContentPart[] = [ { type: "text", text: "Before" }, - { - type: "tool-invocation", - stage: "receiving", - invocationId: "inv-123", - name: "search", - partialArgsText: '{"q": "test"}', - partialArgs: { dummy: "not used in this test" }, - }, + createReceivingToolInvocation("inv-123", "search", '{"q": "test"}'), { type: "text", text: "After" }, ]; @@ -703,23 +731,9 @@ describe("patchContentWithDelta", () => { test("replaces the LAST matching tool when multiple have same invocationId", () => { const content: AiAssistantContentPart[] = [ - { - type: "tool-invocation", - stage: "receiving", - invocationId: "inv-dup", - name: "first", - partialArgsText: "a", - partialArgs: { dummy: "not used in this test" }, - }, + createReceivingToolInvocation("inv-dup", "first", "a"), { type: "text", text: "Middle" }, - { - type: "tool-invocation", - stage: "receiving", - invocationId: "inv-dup", - name: "second", - partialArgsText: "b", - partialArgs: { dummy: "not used in this test" }, - }, + createReceivingToolInvocation("inv-dup", "second", "b"), ]; const delta: AiExecutingToolInvocationPart = { @@ -740,14 +754,7 @@ describe("patchContentWithDelta", () => { test("does not affect tools with different invocationIds", () => { const content: AiAssistantContentPart[] = [ - { - type: "tool-invocation", - stage: "receiving", - invocationId: "inv-1", - name: "tool1", - partialArgsText: "", - partialArgs: { dummy: "not used in this test" }, - }, + createReceivingToolInvocation("inv-1", "tool1"), { type: "tool-invocation", stage: "executing", diff --git a/packages/liveblocks-core/src/lib/IncrementalJsonParser.ts b/packages/liveblocks-core/src/lib/IncrementalJsonParser.ts new file mode 100644 index 00000000000..6b1a7bcb248 --- /dev/null +++ b/packages/liveblocks-core/src/lib/IncrementalJsonParser.ts @@ -0,0 +1,251 @@ +import type { JsonObject } from "./Json"; +import { tryParseJson } from "./utils"; + +const EMPTY_OBJECT = Object.freeze({}) as JsonObject; + +// Characters that can end partial keywords: n, u, l, t, r, e, f, a, s +const NULL_KEYWORD_CHARS = Array.from(new Set("null")); +const TRUE_KEYWORD_CHARS = Array.from(new Set("true")); +const FALSE_KEYWORD_CHARS = Array.from(new Set("false")); +const ALL_KEYWORD_CHARS = Array.from(new Set("nulltruefalse")); + +/** + * Strips the last character from `str` if it is one of the chars in the given + * `chars` string. + */ +function stripChar(str: string, chars: string): string { + const lastChar = str[str.length - 1]; + if (chars.includes(lastChar)) { + return str.slice(0, -1); + } + return str; +} + +export class IncrementalJsonParser { + // Input + #sourceText: string = ""; + + // Output + #cachedJson?: JsonObject; + + /** How much we've already parsed */ + #scanIndex: number = 0; + /** Whether the last char processed was a backslash */ + #escaped: boolean = false; + /** + * Start position of the last unterminated string, -1 if we're not inside + * a string currently. + * + * Example: '{"a": "foo' + * ^ + */ + #lastUnterminatedString: number = -1; + /** + * Start position of the last fully terminated string we've seen. + * + * Example: '{"a": "foo' + * ^ + */ + #lastTerminatedString: number = -1; + /** The bracket stack of expected closing chars. For input '{"a": ["foo', the stack would be ['}', ']']. */ + #stack: string[] = []; + + constructor(text: string = "") { + this.append(text); + } + + get source(): string { + return this.#sourceText; + } + + get json(): JsonObject { + if (this.#cachedJson === undefined) { + this.#cachedJson = this.#parse(); + } + return this.#cachedJson; + } + + /** Whether we're currently inside an unterminated string, e.g. '{"hello' */ + get #inString(): boolean { + return this.#lastUnterminatedString >= 0; + } + + append(delta: string): void { + if (delta) { + // Trim leading whitespace only on the first delta + if (this.#sourceText === "") { + delta = delta.trimStart(); + } + this.#sourceText += delta; + this.#cachedJson = undefined; // Invalidate the cache + } + } + + #autocompleteTail(output: string): string { + // Complete unambiguous partial JSON keywords, + // e.g. '{"a": -' → '{"a": -0' + // '{"a": n' → '{"a": null' + // '{"a": t' → '{"a": true' + // '{"a": f' → '{"a": false' + + if (this.#inString) { + return ""; // Don't complete anything if we're in an unterminated string + } + + const lastChar = output.charAt(output.length - 1); + if (lastChar === "") return ""; + + // Handle incomplete negative numbers + if (lastChar === "-") { + return "0"; // Complete to -0 + } + + // Skip keyword completion for most characters that can't be part of keywords + if (!ALL_KEYWORD_CHARS.includes(lastChar)) return ""; + + // Check the last few characters directly + if (NULL_KEYWORD_CHARS.includes(lastChar)) { + if (output.endsWith("nul")) return "l"; + if (output.endsWith("nu")) return "ll"; + if (output.endsWith("n")) return "ull"; + } + + if (TRUE_KEYWORD_CHARS.includes(lastChar)) { + if (output.endsWith("tru")) return "e"; + if (output.endsWith("tr")) return "ue"; + if (output.endsWith("t")) return "rue"; + } + + if (FALSE_KEYWORD_CHARS.includes(lastChar)) { + if (output.endsWith("fals")) return "e"; + if (output.endsWith("fal")) return "se"; + if (output.endsWith("fa")) return "lse"; + if (output.endsWith("f")) return "alse"; + } + + return ""; + } + + /** + * Updates the internal parsing state by processing any new content + * that has been appended since the last parse. This updates the state with + * facts only. Any interpretation is left to the #parse() method. + */ + #catchup(): void { + const newContent = this.#sourceText.slice(this.#scanIndex); + + // Update internal parsing state by processing only the new content character by character + for (let i = 0; i < newContent.length; i++) { + const ch = newContent[i]; + const absolutePos = this.#scanIndex + i; + + if (this.#inString) { + if (this.#escaped) { + this.#escaped = false; + } else if (ch === "\\") { + this.#escaped = true; + } else if (ch === '"') { + this.#lastTerminatedString = this.#lastUnterminatedString; // Save the terminated string's start + this.#lastUnterminatedString = -1; // Exit string + } + } else { + if (ch === '"') { + this.#lastUnterminatedString = absolutePos; // Enter string + } else if (ch === "{") { + this.#stack.push("}"); + } else if (ch === "[") { + this.#stack.push("]"); + } else if ( + ch === "}" && + this.#stack.length > 0 && + this.#stack[this.#stack.length - 1] === "}" + ) { + this.#stack.pop(); + } else if ( + ch === "]" && + this.#stack.length > 0 && + this.#stack[this.#stack.length - 1] === "]" + ) { + this.#stack.pop(); + } + } + } + + this.#scanIndex = this.#sourceText.length; + } + + #parse(): JsonObject { + this.#catchup(); + + let result = this.#sourceText; // Already trimmed on first append + + if (result.charAt(0) !== "{") { + // Not an object, don't even try to parse it + return EMPTY_OBJECT; + } + + // If it's already valid JSON, return as-is + if (result.endsWith("}")) { + const quickCheck = tryParseJson(result); + if (quickCheck) { + // Due to the '{' check above, we can safely assume it's an object + return quickCheck as JsonObject; + } + } + + // Fix unterminated strings by appending a '"' if needed + // Use our tracked state instead of recalculating + if (this.#inString) { + // If we're in an escaped state (last char was \), remove that incomplete escape + if (this.#escaped) { + result = result.slice(0, -1); // Remove the trailing backslash + } + result += '"'; + } + + // If the last char is a ',' or '.', we can strip it, because it won't + // change the value. Trim whitespace first, then check for comma/period. + result = result.trimEnd(); + result = stripChar(result, ",."); + + // Complete partial keywords at the end (if umambiguous) + result = result + this.#autocompleteTail(result); + + // Use the bracket stack to compute the suffix + const suffix = this.#stack.reduceRight((acc, ch) => acc + ch, ""); + + // Attempt to "just" add the missing ] and }'s. + { + const attempt = tryParseJson(result + suffix); + if (attempt) { + // If it parses, return the result + return attempt as JsonObject; + } + } + + // If there is a parse failure above, it's likely because we're missing + // a "value" for a key in an object. + + if (this.#inString) { + // We're in an unterminated string, just remove it - e.g. '{"abc' + result = result.slice(0, this.#lastUnterminatedString); + } else { + // If the last char is a ":", just remove it - e.g. '{"abc"' or '{"abc":' + result = stripChar(result, ":"); + + // If the last char is a '"', remove that last string + if (result.endsWith('"')) { + result = result.slice(0, this.#lastTerminatedString); + } + } + + // If the last char now is a trailing comma, strip it + result = stripChar(result, ","); + + // Re-add the missing brackets/braces + result += suffix; + + // Run JSON.parse on the result again. it should now work! + return (tryParseJson(result) as JsonObject | undefined) ?? EMPTY_OBJECT; // Still invalid JSON + } +} diff --git a/packages/liveblocks-core/src/lib/__tests__/IncrementalJsonParser.test.ts b/packages/liveblocks-core/src/lib/__tests__/IncrementalJsonParser.test.ts new file mode 100644 index 00000000000..afa744493dc --- /dev/null +++ b/packages/liveblocks-core/src/lib/__tests__/IncrementalJsonParser.test.ts @@ -0,0 +1,403 @@ +import fc from "fast-check"; +import { describe, expect, test } from "vitest"; + +import { IncrementalJsonParser } from "../IncrementalJsonParser"; + +// Helper to express tests more briefly +function parse(input: string) { + return new IncrementalJsonParser(input).json; +} + +describe("IncrementalJsonParser", () => { + test("constructor with initial value", () => { + const parser = new IncrementalJsonParser('{"key":"value"}'); + expect(parser.source).toBe('{"key":"value"}'); + expect(parser.json).toEqual({ key: "value" }); + }); + + test("basic functionality", () => { + const parser = new IncrementalJsonParser(); + + expect(parser.source).toBe(""); + expect(parser.json).toEqual({}); + + parser.append('{"key":"value"}'); + expect(parser.source).toBe('{"key":"value"}'); + expect(parser.json).toEqual({ key: "value" }); + }); + + test("trims leading whitespace only once", () => { + const parser = new IncrementalJsonParser(); + + // First append with leading whitespace should be trimmed + parser.append(" {"); + expect(parser.source).toBe("{"); + + // Subsequent appends should not be trimmed + parser.append(' "key"'); + expect(parser.source).toBe('{ "key"'); + + parser.append(': "value"}'); + expect(parser.source).toBe('{ "key": "value"}'); + expect(parser.json).toEqual({ key: "value" }); + }); + + test("incremental parsing", () => { + const parser = new IncrementalJsonParser(); + + parser.append('{"k'); + expect(parser.json).toEqual({}); + + parser.append('ey":"val'); + expect(parser.json).toEqual({ key: "val" }); + + parser.append('ue"}'); + expect(parser.json).toEqual({ key: "value" }); + }); + + test("bulk append produces same result as character-by-character", () => { + const testString = + '{"complex":{"nested":[1,2,{"deep":"value"}],"more":"data"}}'; + + // Bulk append + const parser1 = new IncrementalJsonParser(testString); + const result1 = parser1.json; + + // Character by character + const parser2 = new IncrementalJsonParser(); + for (const char of testString) { + parser2.append(char); + } + const result2 = parser2.json; + + expect(result1).toEqual(result2); + }); +}); + +describe("caching behavior", () => { + test("resulting json is cached", () => { + const parser = new IncrementalJsonParser('{"test":12'); + const result1 = parser.json; + const result2 = parser.json; + expect(result1).toBe(result2); + }); + + test("cache is invalidated when text is appended", () => { + const parser = new IncrementalJsonParser('{"test":12'); + const result1 = parser.json; + parser.append("3"); + expect(result1).not.toEqual(parser.json); + }); +}); + +describe("parsing basic inputs", () => { + test("basic cases", () => { + expect(parse("")).toEqual({}); + expect(parse(" ")).toEqual({}); + expect(parse("{")).toEqual({}); + expect(parse('{"key"')).toEqual({}); + expect(parse('{"key')).toEqual({}); + expect(parse('{"key":')).toEqual({}); + expect(parse('{"key":""')).toEqual({ key: "" }); + expect(parse('{"key":"')).toEqual({ key: "" }); + expect(parse('{"key":0')).toEqual({ key: 0 }); + expect(parse('{"key":"hi')).toEqual({ key: "hi" }); + expect(parse('{"key":"value"')).toEqual({ key: "value" }); + expect(parse('{"key":"value"}')).toEqual({ key: "value" }); + }); + + test("arrays", () => { + expect(parse('{"arr":[')).toEqual({ arr: [] }); + expect(parse('{"arr":[1')).toEqual({ arr: [1] }); + expect(parse('{"arr":[1,')).toEqual({ arr: [1] }); + expect(parse('{"arr":["hi')).toEqual({ arr: ["hi"] }); + expect(parse('{"arr":["hi"')).toEqual({ arr: ["hi"] }); + expect(parse('{"arr":["hi"]')).toEqual({ arr: ["hi"] }); + expect(parse('{"arr":[1,2,3')).toEqual({ arr: [1, 2, 3] }); + }); + + test("nested structures", () => { + expect(parse('{"arr":[')).toEqual({ arr: [] }); + expect(parse('{"arr":[{')).toEqual({ arr: [{}] }); + expect(parse('{"arr":[{"nested"')).toEqual({ arr: [{}] }); + expect(parse('{"arr":[{"nested":')).toEqual({ arr: [{}] }); + expect(parse('{"arr":[{"nested":42')).toEqual({ arr: [{ nested: 42 }] }); + }); + + test("mixed nesting", () => { + expect(parse('{"a":[1,{"b":')).toEqual({ a: [1, {}] }); + expect(parse('{"a":[1,{"b":[')).toEqual({ a: [1, { b: [] }] }); + expect(parse('{"a":{"b":["c"')).toEqual({ a: { b: ["c"] } }); + }); + + test("strings with special characters", () => { + expect(parse('{"key":"val\\n')).toEqual({ key: "val\n" }); + expect(parse('{"key":"val\\"')).toEqual({ key: 'val"' }); + expect(parse('{"key":"val w/ spaces')).toEqual({ key: "val w/ spaces" }); + expect(parse('{"unicode":"café')).toEqual({ unicode: "café" }); + expect(parse('{"key":"val\\\\n')).toEqual(parse('{"key":"val\\\\n')); + }); + + test("numbers", () => { + expect(parse('{"num":123')).toEqual({ num: 123 }); + expect(parse('{"num":123.')).toEqual({ num: 123 }); + expect(parse('{"num":-42')).toEqual({ num: -42 }); + expect(parse('{"pi":3.')).toEqual({ pi: 3 }); + expect(parse('{"pi":3.14')).toEqual({ pi: 3.14 }); + }); + + test("complex real-world examples", () => { + expect(parse('{"users":[{"id":1,"name":"John')).toEqual({ + users: [{ id: 1, name: "John" }], + }); + expect(parse('{"config":{"debug":true,"timeout":')).toEqual({ + config: { debug: true }, + }); + expect(parse('{"data":[{"items":[{"type":"text","content":')).toEqual({ + data: [{ items: [{ type: "text" }] }], + }); + }); + + test("edge cases for coverage", () => { + // Test properly closed nested objects to trigger stack.pop for '}' + expect(parse('{"a":{"b":{}}')).toEqual({ a: { b: {} } }); + expect(parse('{"nested":{"deep":{"obj":{}}}}')).toEqual({ + nested: { deep: { obj: {} } }, + }); + + // Test escaped characters at end of incomplete strings + expect(parse('{"key":"val\\"')).toEqual({ key: 'val"' }); + expect(parse('{"key":"val\\\\"')).toEqual({ key: "val\\" }); + expect(parse('{"esc":"test\\\\')).toEqual({ esc: "test\\" }); + + // Test combination of escapes and incomplete structure + expect(parse('{"a":"b\\\\","c":')).toEqual({ a: "b\\" }); + + // Test escape handling in string backtracking + expect(parse('{"a":"test\\\\value:')).toEqual({ a: "test\\value:" }); + expect(parse('{"a":"test\\\\value":')).toEqual({}); + + // Test handling of input ending with colon after numeric value + expect(parse('{"a":1:')).toEqual({ a: 1 }); + + // Cover escaped characters inside strings + expect(parse('{"msg":"Say \\"hello\\"')).toEqual({ msg: 'Say "hello"' }); + expect(parse('{"path":"C:\\\\Users')).toEqual({ path: "C:\\Users" }); + + // Escape char at the end of input + expect(parse('{"a":"e\\')).toEqual({ a: "e" }); + + // Cover array closing bracket stack operations + expect(parse('{"arr":[1,2]}')).toEqual({ arr: [1, 2] }); + expect(parse('{"nested":[[[]]]}')).toEqual({ nested: [[[]]] }); + expect(parse('{"mixed":[{"a":1}]}')).toEqual({ mixed: [{ a: 1 }] }); + + // Cover error recovery logic for malformed JSON + + // Test colon removal for incomplete key-value pairs + expect(parse('{"a":1,"incomplete":')).toEqual({ a: 1 }); + + // Test handling of unmatched quotes + expect(parse('{"a":1,"bad":"incomplete')).toEqual({ + a: 1, + bad: "incomplete", + }); + + // Test comma removal for trailing commas + expect(parse('{"a":1,"b":2,')).toEqual({ a: 1, b: 2 }); + expect(parse('{"foo": "bar", ')).toEqual({ foo: "bar" }); + + // Test fallback to {} for completely malformed input + expect(parse("not json at all")).toEqual({}); + expect(parse('{"completely broken syntax",}')).toEqual({}); + expect(parse('{"key":"completely broken syntax",}')).toEqual({}); + expect(parse('{"key","completely broken syntax",}')).toEqual({}); + }); + + test("handles trailing whitespace", () => { + expect(parse('{"key":"value"} \n\t')).toEqual({ key: "value" }); + expect(parse('{"key":"value"} ')).toEqual({ key: "value" }); + expect(parse('{"a":1,"b":2} \n ')).toEqual({ a: 1, b: 2 }); + expect(parse('{"nested":{"obj":{}}} ')).toEqual({ nested: { obj: {} } }); + }); + + test("handles emojis and newlines", () => { + // Complete emojis in strings + expect(parse('{"message":"Hello 👋 world')).toEqual({ + message: "Hello 👋 world", + }); + expect(parse('{"reaction":"🎉","status":"complete')).toEqual({ + reaction: "🎉", + status: "complete", + }); + + // Partial/incomplete emojis (multi-byte sequences) + expect(parse('{"partial":"test 👋')).toEqual({ partial: "test 👋" }); + expect(parse('{"emoji":"🎉🎊')).toEqual({ emoji: "🎉🎊" }); + + // Newlines and whitespace in strings + expect(parse('{"text":"line1\\nline2')).toEqual({ text: "line1\nline2" }); + expect(parse('{"multiline":"first line\\nsecond')).toEqual({ + multiline: "first line\nsecond", + }); + + // Mixed emojis, newlines, and regular content + expect(parse('{"log":"User clicked 👆\\nAction: success ✅')).toEqual({ + log: "User clicked 👆\nAction: success ✅", + }); + + // Emojis with arrays and objects + expect(parse('{"reactions":["👍","👎","❤️')).toEqual({ + reactions: ["👍", "👎", "❤️"], + }); + expect(parse('{"user":{"name":"Alice","status":"🟢 online')).toEqual({ + user: { name: "Alice", status: "🟢 online" }, + }); + }); +}); + +describe("partial keyword recognition", () => { + test("recognizes partial null keyword", () => { + expect(parse('{"key":n')).toEqual({ key: null }); + expect(parse('{"key":nu')).toEqual({ key: null }); + expect(parse('{"key":nul')).toEqual({ key: null }); + }); + + test("recognizes partial true keyword", () => { + expect(parse('{"key":t')).toEqual({ key: true }); + expect(parse('{"key":tr')).toEqual({ key: true }); + expect(parse('{"key":tru')).toEqual({ key: true }); + }); + + test("recognizes partial false keyword", () => { + expect(parse('{"key":f')).toEqual({ key: false }); + expect(parse('{"key":fa')).toEqual({ key: false }); + expect(parse('{"key":fal')).toEqual({ key: false }); + expect(parse('{"key":fals')).toEqual({ key: false }); + }); + + test("recognizes partial keywords in arrays", () => { + expect(parse('{"arr":[t')).toEqual({ arr: [true] }); + expect(parse('{"arr":[n')).toEqual({ arr: [null] }); + expect(parse('{"arr":[f')).toEqual({ arr: [false] }); + }); + + test("recognizes partial keywords after commas", () => { + expect(parse('{"arr":[true,f')).toEqual({ arr: [true, false] }); + expect(parse('{"a":1,"b":n')).toEqual({ a: 1, b: null }); + }); + + test("does not complete keywords in strings", () => { + // This should not complete "n" to "null" because it's inside a string + expect(parse('{"key":"n')).toEqual({ key: "n" }); + }); + + test("fixes the original counterexample", () => { + const input = '{"":false," ":null}'; + const parser = new IncrementalJsonParser(); + let maxLen = 0; + + // Append character by character and track the monotonic property + for (let i = 0; i < input.length; i++) { + parser.append(input[i]); + const result = parser.json; + const len = JSON.stringify(result).length; + + // Output length should never shrink - only stay equal or grow + expect(len).toBeGreaterThanOrEqual(maxLen); + maxLen = Math.max(maxLen, len); + } + + // The final result should be correct + expect(parser.json).toEqual({ "": false, " ": null }); + }); +}); + +describe("property-based testing", () => { + test("regression: escaped chars in keys", () => { + const parser = new IncrementalJsonParser('{"\\"":1'); + expect(parser.json).toEqual({ '"': 1 }); + }); + + test("regression: escaped chars in keys when appended", () => { + const parser = new IncrementalJsonParser('{"\\"'); + expect(parser.json).toEqual({}); + parser.append('":1'); + expect(parser.json).toEqual({ '"': 1 }); + }); + + test("regression: parsing positive decimal numbers", () => { + const parser = new IncrementalJsonParser('{" ":'); + expect(parser.json).toEqual({}); + parser.append("0.007"); + expect(parser.json).toEqual({ " ": 0.007 }); + }); + + test("regression: parsing negative decimal numbers", () => { + const parser = new IncrementalJsonParser('{"x":'); + expect(parser.json).toEqual({}); + parser.append("-0.007"); + expect(parser.json).toEqual({ x: -0.007 }); + }); + + test("regression: incomplete negative number (minus only)", () => { + const parser = new IncrementalJsonParser('{"x":'); + expect(parser.json).toEqual({}); + parser.append("-"); + expect(parser.json).toEqual({ x: -0 }); + parser.append("0"); + expect(parser.json).toEqual({ x: -0 }); + parser.append("."); + expect(parser.json).toEqual({ x: -0 }); + parser.append("00"); + expect(parser.json).toEqual({ x: -0 }); + parser.append("7"); + expect(parser.json).toEqual({ x: -0.007 }); + }); + + test("parsing left-to-right should only ever increment output", () => { + fc.assert( + fc.property( + // Generate a valid JSON object, as a string + fc + .jsonValue() + .filter( + (obj) => + obj !== null && typeof obj === "object" && !Array.isArray(obj) + ) + .map((obj) => JSON.stringify(obj)) + .filter((str) => !str.includes("e")), // Filter out scientific notation + + fc.context(), + + (input, ctx) => { + const parser = new IncrementalJsonParser(); + let maxLen = 0; + + // Append character by character + // No matter what the input is, we should never see the output shrink + for (let i = 0; i < input.length; i++) { + parser.append(input[i]); + const result = parser.json; + const len = JSON.stringify(result).length; + + // Output length should never shrink - only stay equal or grow + if (len < maxLen) { + ctx.log(`Parsed so far: ${input.slice(0, i)}`); + ctx.log(`Next char: ${input[i]}`); + ctx.log( + `❌ FAILURE: Output shrank from ${maxLen} to ${len} at position ${i}` + ); + return false; + } + + maxLen = Math.max(maxLen, len); + } + + return true; + } + ), + { numRuns: 100 } + ); + }); +}); diff --git a/packages/liveblocks-core/src/lib/__tests__/parsePartialJsonObject.test.ts b/packages/liveblocks-core/src/lib/__tests__/parsePartialJsonObject.test.ts deleted file mode 100644 index edfff1a0f99..00000000000 --- a/packages/liveblocks-core/src/lib/__tests__/parsePartialJsonObject.test.ts +++ /dev/null @@ -1,154 +0,0 @@ -import { describe, expect, test } from "vitest"; - -import { parsePartialJsonObject as p } from "../parsePartialJsonObject"; - -describe("parsePartialJsonObject", () => { - test("basic cases", () => { - expect(p("")).toEqual({}); - expect(p(" ")).toEqual({}); - expect(p("{")).toEqual({}); - expect(p('{"key"')).toEqual({}); - expect(p('{"key')).toEqual({}); - expect(p('{"key":')).toEqual({}); - expect(p('{"key":""')).toEqual({ key: "" }); - expect(p('{"key":"')).toEqual({ key: "" }); - expect(p('{"key":0')).toEqual({ key: 0 }); - expect(p('{"key":"hi')).toEqual({ key: "hi" }); - expect(p('{"key":"value"')).toEqual({ key: "value" }); - expect(p('{"key":"value"}')).toEqual({ key: "value" }); - }); - - test("arrays", () => { - expect(p('{"arr":[')).toEqual({ arr: [] }); - expect(p('{"arr":[1')).toEqual({ arr: [1] }); - expect(p('{"arr":[1,')).toEqual({ arr: [1] }); - expect(p('{"arr":["hi')).toEqual({ arr: ["hi"] }); - expect(p('{"arr":["hi"')).toEqual({ arr: ["hi"] }); - expect(p('{"arr":["hi"]')).toEqual({ arr: ["hi"] }); - expect(p('{"arr":[1,2,3')).toEqual({ arr: [1, 2, 3] }); - }); - - test("nested structures", () => { - expect(p('{"arr":[')).toEqual({ arr: [] }); - expect(p('{"arr":[{')).toEqual({ arr: [{}] }); - expect(p('{"arr":[{"nested"')).toEqual({ arr: [{}] }); - expect(p('{"arr":[{"nested":')).toEqual({ arr: [{}] }); - expect(p('{"arr":[{"nested":42')).toEqual({ arr: [{ nested: 42 }] }); - }); - - test("mixed nesting", () => { - expect(p('{"a":[1,{"b":')).toEqual({ a: [1, {}] }); - expect(p('{"a":[1,{"b":[')).toEqual({ a: [1, { b: [] }] }); - expect(p('{"a":{"b":["c"')).toEqual({ a: { b: ["c"] } }); - }); - - test("strings with special characters", () => { - expect(p('{"key":"val\\n')).toEqual({ key: "val\n" }); - expect(p('{"key":"val\\"')).toEqual({ key: 'val"' }); - expect(p('{"key":"val with spaces')).toEqual({ key: "val with spaces" }); - expect(p('{"unicode":"café')).toEqual({ unicode: "café" }); - }); - - test("numbers", () => { - expect(p('{"num":123')).toEqual({ num: 123 }); - expect(p('{"num":123.')).toEqual({ num: 123 }); - expect(p('{"num":-42')).toEqual({ num: -42 }); - expect(p('{"pi":3.')).toEqual({ pi: 3 }); - expect(p('{"pi":3.14')).toEqual({ pi: 3.14 }); - }); - - test("complex real-world examples", () => { - expect(p('{"users":[{"id":1,"name":"John')).toEqual({ - users: [{ id: 1, name: "John" }], - }); - expect(p('{"config":{"debug":true,"timeout":')).toEqual({ - config: { debug: true }, - }); - expect(p('{"data":[{"items":[{"type":"text","content":')).toEqual({ - data: [{ items: [{ type: "text" }] }], - }); - }); - - test("edge cases for coverage", () => { - // Test properly closed nested objects to trigger stack.pop for '}' - expect(p('{"a":{"b":{}}')).toEqual({ a: { b: {} } }); - expect(p('{"nested":{"deep":{"obj":{}}}}')).toEqual({ - nested: { deep: { obj: {} } }, - }); - - // Test escaped characters at end of incomplete strings - expect(p('{"key":"val\\"')).toEqual({ key: 'val"' }); - expect(p('{"key":"val\\\\"')).toEqual({ key: "val\\" }); - expect(p('{"esc":"test\\\\')).toEqual({ esc: "test\\" }); - - // Test combination of escapes and incomplete structure - expect(p('{"a":"b\\\\","c":')).toEqual({ a: "b\\" }); - - // Test escape handling in string backtracking - expect(p('{"a":"test\\\\value:')).toEqual({ a: "test\\value:" }); - expect(p('{"a":"test\\\\value":')).toEqual({}); - - // Test handling of input ending with colon after numeric value - expect(p('{"a":1:')).toEqual({ a: 1 }); - - // Cover escaped characters inside strings - expect(p('{"msg":"Say \\"hello\\"')).toEqual({ msg: 'Say "hello"' }); - expect(p('{"path":"C:\\\\Users')).toEqual({ path: "C:\\Users" }); - - // Cover array closing bracket stack operations - expect(p('{"arr":[1,2]}')).toEqual({ arr: [1, 2] }); - expect(p('{"nested":[[[]]]}')).toEqual({ nested: [[[]]] }); - expect(p('{"mixed":[{"a":1}]}')).toEqual({ mixed: [{ a: 1 }] }); - - // Cover error recovery logic for malformed JSON - - // Test colon removal for incomplete key-value pairs - expect(p('{"a":1,"incomplete":')).toEqual({ a: 1 }); - - // Test handling of unmatched quotes - expect(p('{"a":1,"bad":"incomplete')).toEqual({ a: 1, bad: "incomplete" }); - - // Test comma removal for trailing commas - expect(p('{"a":1,"b":2,')).toEqual({ a: 1, b: 2 }); - - // Test fallback to {} for completely malformed input - expect(p("not json at all")).toEqual({}); - expect(p('{"completely broken syntax",}')).toEqual({}); - expect(p('{"key":"completely broken syntax",}')).toEqual({}); - expect(p('{"key","completely broken syntax",}')).toEqual({}); - }); - - test("handles emojis and newlines", () => { - // Complete emojis in strings - expect(p('{"message":"Hello 👋 world')).toEqual({ - message: "Hello 👋 world", - }); - expect(p('{"reaction":"🎉","status":"complete')).toEqual({ - reaction: "🎉", - status: "complete", - }); - - // Partial/incomplete emojis (multi-byte sequences) - expect(p('{"partial":"test 👋')).toEqual({ partial: "test 👋" }); - expect(p('{"emoji":"🎉🎊')).toEqual({ emoji: "🎉🎊" }); - - // Newlines and whitespace in strings - expect(p('{"text":"line1\\nline2')).toEqual({ text: "line1\nline2" }); - expect(p('{"multiline":"first line\\nsecond')).toEqual({ - multiline: "first line\nsecond", - }); - - // Mixed emojis, newlines, and regular content - expect(p('{"log":"User clicked 👆\\nAction: success ✅')).toEqual({ - log: "User clicked 👆\nAction: success ✅", - }); - - // Emojis with arrays and objects - expect(p('{"reactions":["👍","👎","❤️')).toEqual({ - reactions: ["👍", "👎", "❤️"], - }); - expect(p('{"user":{"name":"Alice","status":"🟢 online')).toEqual({ - user: { name: "Alice", status: "🟢 online" }, - }); - }); -}); diff --git a/packages/liveblocks-core/src/types/ai.ts b/packages/liveblocks-core/src/types/ai.ts index 63f5f82030e..abdb3258293 100644 --- a/packages/liveblocks-core/src/types/ai.ts +++ b/packages/liveblocks-core/src/types/ai.ts @@ -7,8 +7,8 @@ import type { JSONSchema7 } from "json-schema"; import { assertNever } from "../lib/assert"; +import { IncrementalJsonParser } from "../lib/IncrementalJsonParser"; import type { Json, JsonObject } from "../lib/Json"; -import { parsePartialJsonObject } from "../lib/parsePartialJsonObject"; import type { Relax } from "../lib/Relax"; import type { Resolve } from "../lib/Resolve"; import type { Brand } from "../lib/utils"; @@ -337,6 +337,8 @@ export type AiReceivingToolInvocationPart = { /** @internal */ partialArgsText: string; // The raw, partial JSON text value partialArgs: JsonObject; // The interpreted, partial JSON value + /** @internal */ + __appendDelta?: (delta: string) => void; // Internal method for delta updates }; export type AiExecutingToolInvocationPart = { @@ -548,23 +550,10 @@ export function patchContentWithDelta( break; case "tool-stream": { - let _cacheKey = ""; - let _cachedArgs: JsonObject = {}; - - const toolInvocation = { - type: "tool-invocation" as const, - stage: "receiving" as const, - invocationId: delta.invocationId, - name: delta.name, - partialArgsText: "", - get partialArgs(): JsonObject { - if (this.partialArgsText !== _cacheKey) { - _cachedArgs = parsePartialJsonObject(this.partialArgsText); - _cacheKey = this.partialArgsText; - } - return _cachedArgs; - }, - }; + const toolInvocation = createReceivingToolInvocation( + delta.invocationId, + delta.name + ); content.push(toolInvocation); break; } @@ -572,12 +561,12 @@ export function patchContentWithDelta( case "tool-delta": { // Take the last part, expect it to be a tool invocation in receiving // stage. If not, ignore this delta. If it is, append the delta to the - // partialArgsText + // parser if ( lastPart?.type === "tool-invocation" && lastPart.stage === "receiving" ) { - lastPart.partialArgsText += delta.delta; + lastPart.__appendDelta?.(delta.delta); } // Otherwise ignore the delta - it's out of order or unexpected break; @@ -607,3 +596,32 @@ export function patchContentWithDelta( return assertNever(delta, "Unhandled case"); } } + +/** + * Creates a receiving tool invocation part for testing purposes. + * This helper eliminates the need to manually create fake tool invocation objects + * and provides a clean API for tests. + */ +export function createReceivingToolInvocation( + invocationId: string, + name: string, + partialArgsText: string = "" +): AiReceivingToolInvocationPart { + const parser = new IncrementalJsonParser(partialArgsText); + return { + type: "tool-invocation", + stage: "receiving", + invocationId, + name, + get partialArgsText(): string { + return parser.source; + }, + get partialArgs(): JsonObject { + return parser.json; + }, + // Internal method to append deltas + __appendDelta(delta: string) { + parser.append(delta); + }, + } satisfies AiReceivingToolInvocationPart; +}