Skip to content

Commit c23ea4e

Browse files
author
Bertan Ari
committed
fix: preserve large leaked calls and exercise fence guard (#1608)
1 parent 59eb1ba commit c23ea4e

2 files changed

Lines changed: 46 additions & 10 deletions

File tree

‎src/api/providers/__tests__/vscode-lm.spec.ts‎

Lines changed: 42 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -494,15 +494,55 @@ describe("VsCodeLmHandler", () => {
494494

495495
it("does not recover an invoke block quoted inside a fenced code block", async () => {
496496
const block = '<invoke name="calculator"><parameter name="operation">add</parameter></invoke>'
497-
const chunks = await collect(["Do NOT do this:\n```\n" + block + "\n```\n"])
497+
// Open the wrapper outside the fence so only the quoting guard prevents recovery.
498+
const chunks = await collect(["<function_calls>\n```\n" + block + "\n```\n</function_calls>"])
498499

499500
expect(chunks.some((chunk) => chunk.type === "tool_call")).toBe(false)
500501
})
501502

503+
it.each([
504+
{ name: "write_to_file", parameter: "content" },
505+
{ name: "update_todo_list", parameter: "todos" },
506+
])(
507+
"recovers a large $name call split across chunks without leaking markup",
508+
async ({ name, parameter }) => {
509+
const payload = "x".repeat(5000)
510+
streamTextParts([
511+
`<function_calls><invoke name="${name}"><parameter name="${parameter}">`,
512+
payload,
513+
"</parameter></invoke></function_calls>",
514+
])
515+
const chunks = []
516+
for await (const chunk of handler.createMessage("system", [{ role: "user", content: "hi" }], {
517+
taskId: "test-task",
518+
tools: [
519+
{
520+
type: "function",
521+
function: {
522+
name,
523+
parameters: { type: "object", properties: { [parameter]: { type: "string" } } },
524+
},
525+
},
526+
],
527+
})) {
528+
chunks.push(chunk)
529+
}
530+
531+
expect(chunks.filter((chunk) => chunk.type !== "usage")).toEqual([
532+
{
533+
type: "tool_call",
534+
id: expect.stringContaining("vscodelm-salvaged-"),
535+
name,
536+
arguments: JSON.stringify({ [parameter]: payload }),
537+
},
538+
])
539+
},
540+
)
541+
502542
it("flushes an over-long never-closing invoke as plain text before the stream ends", async () => {
503543
// End-of-stream flushing produces identical text, so track production timing to
504544
// prove the buffer releases text before the stream ends.
505-
const filler = "x".repeat(5000)
545+
const filler = "x".repeat(256 * 1024)
506546
const parts = ['<invoke name="calculator">', filler, filler, filler, filler]
507547
let partsProduced = 0
508548

‎src/api/providers/vscode-lm.ts‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -104,12 +104,9 @@ const LEAKED_INVOKE_BLOCK = /<(?:antml:)?invoke\s+name="([^"]+)"\s*>([\s\S]*?)<\
104104
/** Upper bound on an incomplete `<invoke ...` tail held back between chunks. */
105105
const MAX_PARTIAL_INVOKE_CARRY = 64
106106

107-
/**
108-
* Upper bound on buffered text held while waiting for a leaked `<invoke>` block to close. Markup
109-
* that never closes would otherwise withhold the whole response from the user until the stream
110-
* ended, so past this size the buffer is released as ordinary text.
111-
*/
112-
const MAX_SALVAGE_BUFFER_CHARS = 64 * MAX_PARTIAL_INVOKE_CARRY
107+
// Allow substantial file contents and todo lists (256 Ki characters including markup), while
108+
// bounding how much never-closing markup can withhold from the user.
109+
const MAX_SALVAGE_BUFFER_CHARS = 256 * 1024
113110

114111
/** True when `text` already contains a closed `<invoke>` block, so buffering is still productive. */
115112
function hasCompleteInvokeBlock(text: string): boolean {
@@ -1216,8 +1213,7 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
12161213
// stream so the full markup can be parsed and replayed as a structured call.
12171214
if (salvageBuffering) {
12181215
salvageBuffer += chunk.value
1219-
// Markup that never closes must not withhold the response indefinitely; past the
1220-
// cap the buffer is released as plain text and buffering stops for the turn.
1216+
// Release never-closing markup past the cap; a later fresh marker can re-arm recovery.
12211217
if (salvageBuffer.length > MAX_SALVAGE_BUFFER_CHARS && !hasCompleteInvokeBlock(salvageBuffer)) {
12221218
const overflowed = salvageBuffer
12231219
salvageBuffering = false

0 commit comments

Comments
 (0)